refactor: update get session to manual call (#26)

Signed-off-by: Lin Wang <wonglam@amazon.com>

Author: wanglam

public/contexts/core_context.tsx

@@ -8,10 +8,12 @@ import {
   useOpenSearchDashboards,
 } from '../../../../src/plugins/opensearch_dashboards_react/public';
 import { AppPluginStartDependencies, SetupDependencies } from '../types';
+import { SessionLoadService } from '../services/session_load_service';
 
 export interface AssistantServices extends Required<OpenSearchDashboardsServices> {
   setupDeps: SetupDependencies;
   startDeps: AppPluginStartDependencies;
+  sessionLoad: SessionLoadService;
 }
 
 export const useCore: () => OpenSearchDashboardsReactContextValue<

public/hooks/use_chat_actions.tsx

@@ -53,16 +53,24 @@ export const useChatActions = (): AssistantActions => {
     }
   };
 
-  const loadChat = (sessionId?: string, title?: string) => {
+  const loadChat = async (sessionId?: string, title?: string) => {
     abortControllerRef?.abort();
+    core.services.sessionLoad.abortController?.abort();
     chatContext.setSessionId(sessionId);
     chatContext.setTitle(title);
     // Chat page will always visible in fullscreen mode, we don't need to change the tab anymore
     if (!chatContext.flyoutFullScreen) {
       chatContext.setSelectedTabId('chat');
     }
     chatContext.setFlyoutComponent(null);
-    if (!sessionId) chatStateDispatch({ type: 'reset' });
+    if (!sessionId) {
+      chatStateDispatch({ type: 'reset' });
+      return;
+    }
+    const session = await core.services.sessionLoad.load(sessionId);
+    if (session) {
+      chatStateDispatch({ type: 'receive', payload: session.messages });
+    }
   };
 
   const openChatUI = () => {

public/hooks/use_sessions.ts

@@ -6,46 +6,10 @@
 import { useCallback, useEffect, useReducer, useState } from 'react';
 import { HttpFetchQuery, SavedObjectsFindOptions } from '../../../../src/core/public';
 import { ASSISTANT_API } from '../../common/constants/llm';
-import { ISession, ISessionFindResponse } from '../../common/types/chat_saved_object_attributes';
-import { useChatContext } from '../contexts/chat_context';
+import { ISessionFindResponse } from '../../common/types/chat_saved_object_attributes';
 import { useCore } from '../contexts/core_context';
 import { GenericReducer, genericReducer, genericReducerWithAbortController } from './fetch_reducer';
 
-export const useGetSession = () => {
-  const chatContext = useChatContext();
-  const core = useCore();
-  const reducer: GenericReducer<ISession> = genericReducer;
-  const [state, dispatch] = useReducer(reducer, { loading: false });
-  const [refreshToggle, setRefreshToggle] = useState(false);
-
-  const refresh = useCallback(() => {
-    setRefreshToggle((flag) => !flag);
-  }, []);
-
-  useEffect(() => {
-    const abortController = new AbortController();
-    dispatch({ type: 'request' });
-    if (!chatContext.sessionId) {
-      dispatch({ type: 'success', payload: undefined });
-      return;
-    }
-
-    core.services.http
-      .get<ISession>(`${ASSISTANT_API.SESSION}/${chatContext.sessionId}`, {
-        signal: abortController.signal,
-      })
-      .then((payload) => dispatch({ type: 'success', payload }))
-      .catch((error) => dispatch({ type: 'failure', error }));
-
-    return () => {
-      abortController.abort();
-    };
-    // refreshToggle is used to force refresh session to get latest data
-  }, [chatContext.sessionId, refreshToggle]);
-
-  return { ...state, refresh };
-};
-
 export const useGetSessions = (options: Partial<SavedObjectsFindOptions> = {}) => {
   const core = useCore();
   const reducer: GenericReducer<ISessionFindResponse> = genericReducer;

public/plugin.tsx

@@ -21,6 +21,7 @@ import {
   ContentRenderer,
   SetupDependencies,
 } from './types';
+import { SessionLoadService } from './services/session_load_service';
 
 export const [getCoreStart, setCoreStart] = createGetterSetter<CoreStart>('CoreStart');
 
@@ -56,6 +57,7 @@ export class AssistantPlugin
         ...coreStart,
         setupDeps,
         startDeps,
+        sessionLoad: new SessionLoadService(coreStart.http),
       });
       const account = await getAccount();
       const username = account.data.user_name;

public/services/session_load_service.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BehaviorSubject, from } from 'rxjs';
+import { HttpStart } from '../../../../src/core/public';
+import { ISession } from '../../common/types/chat_saved_object_attributes';
+import { ASSISTANT_API } from '../../common/constants/llm';
+
+export class SessionLoadService {
+  status$: BehaviorSubject<
+    'idle' | 'loading' | { status: 'error'; error: Error }
+  > = new BehaviorSubject<'idle' | 'loading' | { status: 'error'; error: Error }>('idle');
+  abortController?: AbortController;
+
+  constructor(private _http: HttpStart) {}
+
+  load = async (sessionId: string) => {
+    this.abortController?.abort();
+    this.status$.next('loading');
+    this.abortController = new AbortController();
+    try {
+      return await this._http.get<ISession>(`${ASSISTANT_API.SESSION}/${sessionId}`, {
+        signal: this.abortController.signal,
+      });
+    } catch (error) {
+      this.status$.next({ status: 'error', error });
+    } finally {
+      this.status$.next('idle');
+    }
+  };
+}

public/tabs/chat/chat_page.tsx

@@ -4,34 +4,36 @@
  */
 
 import { EuiFlyoutBody, EuiFlyoutFooter, EuiPage, EuiPageBody, EuiSpacer } from '@elastic/eui';
-import React, { useEffect, useState } from 'react';
+import React, { useCallback, useState } from 'react';
 import cs from 'classnames';
 import { useChatContext } from '../../contexts/chat_context';
 import { useChatState } from '../../hooks/use_chat_state';
-import { useGetSession } from '../../hooks/use_sessions';
 import { ChatPageContent } from './chat_page_content';
 import { ChatInputControls } from './controls/chat_input_controls';
+import { useObservable } from 'react-use';
+import { useCore } from '../../contexts/core_context';
 
 interface ChatPageProps {
   className?: string;
 }
 
 export const ChatPage: React.FC<ChatPageProps> = (props) => {
+  const core = useCore();
   const chatContext = useChatContext();
   const { chatState, chatStateDispatch } = useChatState();
   const [showGreetings, setShowGreetings] = useState(false);
-  const {
-    data: session,
-    loading: messagesLoading,
-    error: messagesLoadingError,
-    refresh,
-  } = useGetSession();
+  const sessionLoadStatus = useObservable(core.services.sessionLoad.status$);
+  const messagesLoading = sessionLoadStatus === 'loading';
 
-  useEffect(() => {
+  const refresh = useCallback(async () => {
+    if (!chatContext.sessionId) {
+      return;
+    }
+    const session = await core.services.sessionLoad.load(chatContext.sessionId);
     if (session) {
       chatStateDispatch({ type: 'receive', payload: session.messages });
     }
-  }, [session]);
+  }, [chatContext.sessionId, chatStateDispatch]);
 
   return (
     <>
@@ -42,7 +44,9 @@ export const ChatPage: React.FC<ChatPageProps> = (props) => {
               showGreetings={showGreetings}
               setShowGreetings={setShowGreetings}
               messagesLoading={messagesLoading}
-              messagesLoadingError={messagesLoadingError}
+              messagesLoadingError={
+                typeof sessionLoadStatus !== 'string' ? sessionLoadStatus?.error : undefined
+              }
               onRefresh={refresh}
             />
           </EuiPageBody>