Adding context provider for Agent Chat front end to persist the chat while navigating away from the agent chat page

morishwe · morishwe · commit 1d1bae27e6b8 · 2025-11-05T16:13:45.000-07:00
diff --git a/src/ui/src/App.jsx b/src/ui/src/App.jsx
@@ -8,6 +8,7 @@ import '@aws-amplify/ui-react/styles.css';
 
 import { AppContext } from './contexts/app';
 import { AnalyticsProvider } from './contexts/analytics';
+import { AgentChatProvider } from './contexts/agentChat';
 import useAwsConfig from './hooks/use-aws-config';
 import useCurrentSessionCreds from './hooks/use-current-session-creds';
 
@@ -37,14 +38,16 @@ const AppContent = () => {
     setNavigationOpen,
   };
   logger.debug('appContextValue', appContextValue);
-
+  // TODO: Remove the AnalyticsProvider once we migrate full to Agent Chat
   return (
     <div className="App">
       <AppContext.Provider value={appContextValue}>
         <AnalyticsProvider>
-          <HashRouter>
-            <Routes />
-          </HashRouter>
+          <AgentChatProvider>
+            <HashRouter>
+              <Routes />
+            </HashRouter>
+          </AgentChatProvider>
         </AnalyticsProvider>
       </AppContext.Provider>
     </div>
diff --git a/src/ui/src/components/agent-chat/AgentChatLayout.jsx b/src/ui/src/components/agent-chat/AgentChatLayout.jsx
@@ -20,6 +20,7 @@ import remarkGfm from 'remark-gfm';
 import rehypeRaw from 'rehype-raw';
 import useAgentChat from '../../hooks/use-agent-chat';
 import useAppContext from '../../contexts/app';
+import { useAgentChatContext } from '../../contexts/agentChat';
 import PlotDisplay from '../document-agents-layout/PlotDisplay';
 import TableDisplay from '../document-agents-layout/TableDisplay';
 import AgentChatHistoryDropdown from './AgentChatHistoryDropdown';
@@ -33,12 +34,14 @@ const AgentChatLayout = ({
   showHeader = true,
   customStyles = {},
 }) => {
-  const [inputValue, setInputValue] = useState('');
-  const [expandedSections, setExpandedSections] = useState(new Set());
   const [welcomeAnimated, setWelcomeAnimated] = useState(false);
-  const [lastMessageCount, setLastMessageCount] = useState(0);
-  const [enableCodeIntelligence, setEnableCodeIntelligence] = useState(true);
+  const [isLoadingSession, setIsLoadingSession] = useState(false);
   const chatMessagesRef = useRef(null);
+
+  // Get persistent state from context
+  const { agentChatState, updateAgentChatState } = useAgentChatContext();
+  const { inputValue, expandedSections, lastMessageCount, enableCodeIntelligence } = agentChatState;
+
   const { messages, isLoading, waitingForResponse, error, sendMessage, clearError, clearChat, loadChatSession } = useAgentChat(agentConfig);
   const { user } = useAppContext();
 
@@ -59,15 +62,15 @@ const AgentChatLayout = ({
   useEffect(() => {
     const handleSampleQueryInsert = (event) => {
       const { query } = event.detail;
-      setInputValue(query);
+      updateAgentChatState({ inputValue: query });
     };
 
     window.addEventListener('insertSampleQuery', handleSampleQueryInsert);
 
     return () => {
       window.removeEventListener('insertSampleQuery', handleSampleQueryInsert);
     };
-  }, []);
+  }, [updateAgentChatState]);
 
   // Track new messages and scroll to new assistant messages (but not while streaming)
   useEffect(() => {
@@ -88,15 +91,15 @@ const AgentChatLayout = ({
         }, 100);
       }
 
-      setLastMessageCount(messages.length);
+      updateAgentChatState({ lastMessageCount: messages.length });
     }
-  }, [messages, lastMessageCount]);
+  }, [messages, lastMessageCount, updateAgentChatState]);
 
   const handlePromptSubmit = async () => {
     const prompt = inputValue;
     if (!prompt.trim()) return;
 
-    setInputValue('');
+    updateAgentChatState({ inputValue: '' });
     try {
       await sendMessage(prompt, { enableCodeIntelligence });
       // Scroll to the latest user message after sending
@@ -115,7 +118,7 @@ const AgentChatLayout = ({
   };
 
   const handleInputChange = (event) => {
-    setInputValue(event.detail.value);
+    updateAgentChatState({ inputValue: event.detail.value });
   };
 
   const handleKeyDown = (event) => {
@@ -127,26 +130,22 @@ const AgentChatLayout = ({
 
   // Handle expandable section state changes
   const handleExpandedChange = (messageId, expanded) => {
-    setExpandedSections((prev) => {
-      const newSet = new Set(prev);
-      if (expanded) {
-        newSet.add(messageId);
-      } else {
-        newSet.delete(messageId);
-      }
-      return newSet;
-    });
+    const newSet = new Set(expandedSections);
+    if (expanded) {
+      newSet.add(messageId);
+    } else {
+      newSet.delete(messageId);
+    }
+    updateAgentChatState({ expandedSections: newSet });
   };
 
   // Handle session selection from dropdown
   const handleSessionSelect = async (session, sessionMessages) => {
     try {
+      setIsLoadingSession(true);
       console.log('Loading chat session:', session.sessionId);
-      await loadChatSession(session.sessionId, sessionMessages);
 
-      // Reset UI state for loaded session
-      setExpandedSections(new Set());
-      setLastMessageCount(sessionMessages.length);
+      await loadChatSession(session.sessionId, sessionMessages);
 
       // Scroll to bottom after loading
       setTimeout(() => {
@@ -156,6 +155,8 @@ const AgentChatLayout = ({
       }, 100);
     } catch (err) {
       console.error('Failed to load chat session:', err);
+    } finally {
+      setIsLoadingSession(false);
     }
   };
 
@@ -281,7 +282,41 @@ const AgentChatLayout = ({
           </Alert>
         )}
 
-        <div ref={chatMessagesRef} className="chat-messages">
+        <div
+          ref={chatMessagesRef}
+          className="chat-messages"
+          style={{
+            position: 'relative',
+            opacity: isLoadingSession ? 0.5 : 1,
+            pointerEvents: isLoadingSession ? 'none' : 'auto',
+            transition: 'opacity 0.3s ease',
+          }}
+        >
+          {isLoadingSession && (
+            <div
+              style={{
+                position: 'absolute',
+                top: '50%',
+                left: '50%',
+                transform: 'translate(-50%, -50%)',
+                zIndex: 1000,
+                display: 'flex',
+                flexDirection: 'column',
+                alignItems: 'center',
+                gap: '12px',
+                backgroundColor: 'rgba(255, 255, 255, 0.9)',
+                padding: '20px',
+                borderRadius: '8px',
+                boxShadow: '0 4px 12px rgba(0, 0, 0, 0.1)',
+              }}
+            >
+              <Spinner size="large" />
+              <Box fontSize="body-m" color="text-body-secondary">
+                Loading chat history...
+              </Box>
+            </div>
+          )}
+
           {messages.length === 0 ? (
             <div className={`welcome-text ${welcomeAnimated ? 'animate-in' : ''}`}>
               <h2>
@@ -311,7 +346,7 @@ const AgentChatLayout = ({
                   onItemClick={async ({ detail }) => {
                     const selectedPrompt = supportPrompts.find((prompt) => prompt.id === detail.id);
                     if (selectedPrompt) {
-                      setInputValue(selectedPrompt.prompt);
+                      updateAgentChatState({ inputValue: selectedPrompt.prompt });
                     }
                   }}
                 />{' '}
@@ -325,7 +360,7 @@ const AgentChatLayout = ({
                 onChange={handleInputChange}
                 onKeyDown={handleKeyDown}
                 placeholder={placeholder}
-                disabled={isLoading}
+                disabled={isLoading || isLoadingSession}
                 actionButtonIconName="send"
                 onAction={handlePromptSubmit}
                 minRows={3}
@@ -336,7 +371,7 @@ const AgentChatLayout = ({
                 </Box>
                 <Checkbox
                   checked={enableCodeIntelligence}
-                  onChange={({ detail }) => setEnableCodeIntelligence(detail.checked)}
+                  onChange={({ detail }) => updateAgentChatState({ enableCodeIntelligence: detail.checked })}
                   disabled={waitingForResponse}
                 >
                   <Box fontSize="body-s">Enable Code Intelligence Agent</Box>
@@ -349,7 +384,7 @@ const AgentChatLayout = ({
               <AgentChatHistoryDropdown
                 onSessionSelect={handleSessionSelect}
                 onSessionDeleted={handleSessionDeleted}
-                disabled={waitingForResponse}
+                disabled={waitingForResponse || isLoadingSession}
               />
             </Box>
             {messages.length > 0 && (
@@ -358,14 +393,12 @@ const AgentChatLayout = ({
                 iconName="refresh"
                 onClick={() => {
                   clearChat();
-                  setExpandedSections(new Set());
                   setWelcomeAnimated(false);
-                  setLastMessageCount(0);
                   setTimeout(() => {
                     setWelcomeAnimated(true);
                   }, 100);
                 }}
-                disabled={waitingForResponse}
+                disabled={waitingForResponse || isLoadingSession}
               >
                 Clear chat
               </Button>
diff --git a/src/ui/src/contexts/agentChat.jsx b/src/ui/src/contexts/agentChat.jsx
@@ -0,0 +1,103 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import React, { createContext, useContext, useState, useCallback, useMemo } from 'react';
+import PropTypes from 'prop-types';
+import { v4 as uuidv4 } from 'uuid';
+
+const AgentChatContext = createContext(null);
+
+export const AgentChatProvider = ({ children }) => {
+  // State for the agent chat
+  const [agentChatState, setAgentChatState] = useState({
+    messages: [], // Current chat messages
+    sessionId: uuidv4(), // Current session ID
+    isLoading: false,
+    waitingForResponse: false,
+    error: null,
+    expandedSections: new Set(),
+    lastMessageCount: 0,
+    enableCodeIntelligence: true,
+    inputValue: '',
+  });
+
+  // Function to update agent chat state
+  const updateAgentChatState = useCallback((updates) => {
+    setAgentChatState((prevState) => ({
+      ...prevState,
+      ...updates,
+    }));
+  }, []);
+
+  // Function to reset agent chat state (new session)
+  const resetAgentChatState = useCallback(() => {
+    setAgentChatState({
+      messages: [],
+      sessionId: uuidv4(),
+      isLoading: false,
+      waitingForResponse: false,
+      error: null,
+      expandedSections: new Set(),
+      lastMessageCount: 0,
+      enableCodeIntelligence: true,
+      inputValue: '',
+    });
+  }, []);
+
+  // Function to load a specific session
+  const loadAgentChatSession = useCallback((sessionId, messages) => {
+    setAgentChatState((prevState) => ({
+      ...prevState,
+      messages,
+      sessionId,
+      expandedSections: new Set(),
+      lastMessageCount: messages.length,
+      error: null,
+      waitingForResponse: false,
+      isLoading: false,
+    }));
+  }, []);
+
+  // Function to add a message to the current session
+  const addMessageToSession = useCallback((message) => {
+    setAgentChatState((prevState) => ({
+      ...prevState,
+      messages: [...prevState.messages, message],
+    }));
+  }, []);
+
+  // Function to update messages (for streaming updates)
+  const updateMessages = useCallback((updaterFunction) => {
+    setAgentChatState((prevState) => ({
+      ...prevState,
+      messages: updaterFunction(prevState.messages),
+    }));
+  }, []);
+
+  const contextValue = useMemo(
+    () => ({
+      agentChatState,
+      updateAgentChatState,
+      resetAgentChatState,
+      loadAgentChatSession,
+      addMessageToSession,
+      updateMessages,
+    }),
+    [agentChatState, updateAgentChatState, resetAgentChatState, loadAgentChatSession, addMessageToSession, updateMessages],
+  );
+
+  return <AgentChatContext.Provider value={contextValue}>{children}</AgentChatContext.Provider>;
+};
+
+AgentChatProvider.propTypes = {
+  children: PropTypes.node.isRequired,
+};
+
+export const useAgentChatContext = () => {
+  const context = useContext(AgentChatContext);
+  if (!context) {
+    throw new Error('useAgentChatContext must be used within an AgentChatProvider');
+  }
+  return context;
+};
+
+export default AgentChatContext;
diff --git a/src/ui/src/hooks/use-agent-chat.js b/src/ui/src/hooks/use-agent-chat.js
