ctx-mgmt: ctx session management (dev mode only) (#2415)

2025-05-02 16:12:56 -04:00
parent 2366a3ad01
commit 8ba40bdccc
19 changed files with 710 additions and 118 deletions
@@ -15,12 +15,15 @@ import { SearchView } from './conversation/SearchView';
 import { createRecipe } from '../recipe';
 import { AgentHeader } from './AgentHeader';
 import LayingEggLoader from './LayingEggLoader';
-import { fetchSessionDetails } from '../sessions';
+import { fetchSessionDetails, generateSessionId } from '../sessions';
 import 'react-toastify/dist/ReactToastify.css';
 import { useMessageStream } from '../hooks/useMessageStream';
 import { SessionSummaryModal } from './context_management/SessionSummaryModal';
 import { Recipe } from '../recipe';
-import { ContextManagerProvider, useChatContextManager } from './context_management/ContextManager';
+import {
+  ChatContextManagerProvider,
+  useChatContextManager,
+} from './context_management/ContextManager';
 import { ContextLengthExceededHandler } from './context_management/ContextLengthExceededHandler';
 import {
  Message,
@@ -64,14 +67,14 @@ export default function ChatView({
  setIsGoosehintsModalOpen: (isOpen: boolean) => void;
 }) {
  return (
-    <ContextManagerProvider>
+    <ChatContextManagerProvider>
      <ChatContent
        chat={chat}
        setChat={setChat}
        setView={setView}
        setIsGoosehintsModalOpen={setIsGoosehintsModalOpen}
      />
-    </ContextManagerProvider>
+    </ChatContextManagerProvider>
  );
 }

@@ -93,7 +96,9 @@ function ChatContent({
  const [showGame, setShowGame] = useState(false);
  const [isGeneratingRecipe, setIsGeneratingRecipe] = useState(false);
  const [sessionTokenCount, setSessionTokenCount] = useState<number>(0);
+  const [ancestorMessages, setAncestorMessages] = useState<Message[]>([]);
  const [droppedFiles, setDroppedFiles] = useState<string[]>([]);
+
  const scrollRef = useRef<ScrollAreaHandle>(null);

  const {
@@ -108,7 +113,6 @@ function ChatContent({

  useEffect(() => {
    // Log all messages when the component first mounts
-    console.log('Initial messages when resuming session:', chat.messages);
    window.electron.logInfo(
      'Initial messages when resuming session: ' + JSON.stringify(chat.messages, null, 2)
    );
@@ -129,6 +133,7 @@ function ChatContent({
    setInput: _setInput,
    handleInputChange: _handleInputChange,
    handleSubmit: _submitMessage,
+    updateMessageStreamBody,
  } = useMessageStream({
    api: getApiUrl('/reply'),
    initialMessages: chat.messages,
@@ -159,6 +164,36 @@ function ChatContent({
    },
  });

+  // for CLE events -- create a new session id for the next set of messages
+  useEffect(() => {
+    // If we're in a continuation session, update the chat ID
+    if (summarizedThread.length > 0) {
+      const newSessionId = generateSessionId();
+
+      // Update the session ID in the chat object
+      setChat({
+        ...chat,
+        id: newSessionId!,
+        title: `Continued from ${chat.id}`,
+        messageHistoryIndex: summarizedThread.length,
+      });
+
+      // Update the body used by useMessageStream to send future messages to the new session
+      if (summarizedThread.length > 0 && updateMessageStreamBody) {
+        updateMessageStreamBody({
+          session_id: newSessionId,
+          session_working_dir: window.appConfig.get('GOOSE_WORKING_DIR'),
+        });
+      }
+    }
+
+    // only update if summarizedThread length changes from 0 -> 1+
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    summarizedThread.length > 0,
+  ]);
+
  // Listen for make-agent-from-chat event
  useEffect(() => {
    const handleMakeAgent = async () => {
@@ -201,7 +236,8 @@ function ChatContent({
        window.electron.logInfo('Opening recipe editor window');
      } catch (error) {
        window.electron.logInfo('Failed to create recipe:');
-        window.electron.logInfo(error.message);
+        const errorMessage = error instanceof Error ? error.message : String(error);
+        window.electron.logInfo(errorMessage);
      } finally {
        setIsGeneratingRecipe(false);
      }
@@ -233,25 +269,33 @@ function ChatContent({
  // Handle submit
  const handleSubmit = (e: React.FormEvent) => {
    window.electron.startPowerSaveBlocker();
-    const customEvent = e as CustomEvent;
+    const customEvent = e as unknown as CustomEvent;
    const content = customEvent.detail?.value || '';

    if (content.trim()) {
      setLastInteractionTime(Date.now());

-      if (process.env.ALPHA && summarizedThread.length > 0) {
-        // First reset the messages with the summary
-        resetMessagesWithSummary(messages, setMessages);
+      if (summarizedThread.length > 0) {
+        // move current `messages` to `ancestorMessages` and `messages` to `summarizedThread`
+        resetMessagesWithSummary(
+          messages,
+          setMessages,
+          ancestorMessages,
+          setAncestorMessages,
+          summaryContent
+        );

-        // Then append the new user message
+        // update the chat with new sessionId
+
+        // now call the llm
        setTimeout(() => {
          append(createUserMessage(content));
          if (scrollRef.current?.scrollToBottom) {
            scrollRef.current.scrollToBottom();
          }
-        }, 150); // Small delay to ensure state updates properly
+        }, 150);
      } else {
-        // Normal flow - just append the message
+        // Normal flow (existing code)
        append(createUserMessage(content));
        if (scrollRef.current?.scrollToBottom) {
          scrollRef.current.scrollToBottom();
@@ -324,6 +368,8 @@ function ChatContent({
        // Create tool responses for all interrupted tool requests

        let responseMessage: Message = {
+          display: true,
+          sendToLLM: true,
          role: 'user',
          created: Date.now(),
          content: [],
@@ -352,7 +398,7 @@ function ChatContent({

  // Filter out standalone tool response messages for rendering
  // They will be shown as part of the tool invocation in the assistant message
-  const filteredMessages = messages.filter((message) => {
+  const filteredMessages = [...ancestorMessages, ...messages].filter((message) => {
    // Only filter out when display is explicitly false
    if (message.display === false) return false;

@@ -466,10 +512,13 @@ function ChatContent({
                  ) : (
                    <>
                      {/* Only render GooseMessage if it's not a CLE message (and we are not in alpha mode) */}
-                      {process.env.ALPHA && hasContextLengthExceededContent(message) ? (
+                      {process.env.NODE_ENV === 'development' &&
+                      hasContextLengthExceededContent(message) ? (
                        <ContextLengthExceededHandler
                          messages={messages}
                          messageId={message.id ?? message.created.toString()}
+                          chatId={chat.id}
+                          workingDir={window.appConfig.get('GOOSE_WORKING_DIR') as string}
                        />
                      ) : (
                        <GooseMessage
@@ -529,7 +578,7 @@ function ChatContent({
      </Card>

      {showGame && <FlappyGoose onClose={() => setShowGame(false)} />}
-      {process.env.ALPHA && (
+      {process.env.NODE_ENV === 'development' && (
        <SessionSummaryModal
          isOpen={isSummaryModalOpen}
          onClose={closeSummaryModal}
@@ -5,21 +5,30 @@ import { useChatContextManager } from './ContextManager';
 interface ContextLengthExceededHandlerProps {
  messages: Message[];
  messageId: string;
+  chatId: string;
+  workingDir: string;
 }

 export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandlerProps> = ({
  messages,
  messageId,
+  chatId,
+  workingDir,
 }) => {
-  const { fetchSummary, summaryContent, isLoadingSummary, errorLoadingSummary, openSummaryModal } =
-    useChatContextManager();
+  const {
+    summaryContent,
+    isLoadingSummary,
+    errorLoadingSummary,
+    openSummaryModal,
+    handleContextLengthExceeded,
+  } = useChatContextManager();

  const [hasFetchStarted, setHasFetchStarted] = useState(false);

  // Find the relevant message to check if it's the latest
  const isCurrentMessageLatest =
-    messageId === messages[messages.length - 1].id ||
-    messageId === messages[messages.length - 1].created.toString();
+    messageId === messages[messages.length - 1]?.id ||
+    messageId === String(messages[messages.length - 1]?.created);

  // Only allow interaction for the most recent context length exceeded event
  const shouldAllowSummaryInteraction = isCurrentMessageLatest;
@@ -27,31 +36,60 @@ export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandler
  // Use a ref to track if we've started the fetch
  const fetchStartedRef = useRef(false);

+  // Function to trigger the async operation properly
+  const triggerContextLengthExceeded = () => {
+    setHasFetchStarted(true);
+    fetchStartedRef.current = true;
+
+    // Call the async function without awaiting it in useEffect
+    handleContextLengthExceeded(messages, chatId, workingDir).catch((err) => {
+      console.error('Error handling context length exceeded:', err);
+    });
+  };
+
  useEffect(() => {
-    // Automatically fetch summary if conditions are met
    if (
      !summaryContent &&
      !hasFetchStarted &&
      shouldAllowSummaryInteraction &&
      !fetchStartedRef.current
    ) {
-      setHasFetchStarted(true);
-      fetchStartedRef.current = true;
-      fetchSummary(messages);
+      // Use the wrapper function instead of calling the async function directly
+      triggerContextLengthExceeded();
    }
-  }, [fetchSummary, hasFetchStarted, messages, shouldAllowSummaryInteraction, summaryContent]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [
+    hasFetchStarted,
+    messages,
+    shouldAllowSummaryInteraction,
+    summaryContent,
+    chatId,
+    workingDir,
+  ]);

-  // Handle retry
+  // Handle retry - Call the async function properly
  const handleRetry = () => {
    if (!shouldAllowSummaryInteraction) return;
-    fetchSummary(messages);
+
+    // Reset states for retry
+    setHasFetchStarted(false);
+    fetchStartedRef.current = false;
+
+    // Trigger the process again
+    triggerContextLengthExceeded();
  };

  // Render the notification UI
  return (
    <div className="flex flex-col items-start mt-1 pl-4">
+      {/* Horizontal line with text in the middle - shown regardless of loading state */}
+      <div className="relative flex items-center py-2 w-full">
+        <div className="flex-grow border-t border-gray-300"></div>
+        <div className="flex-grow border-t border-gray-300"></div>
+      </div>
+
      {isLoadingSummary && shouldAllowSummaryInteraction ? (
-        // Only show loading indicator during loading state
+        // Show loading indicator during loading state
        <div className="flex items-center text-xs text-gray-400">
          <span className="mr-2">Preparing summary...</span>
          <span className="animate-spin h-3 w-3 border-2 border-gray-400 rounded-full border-t-transparent"></span>
@@ -59,15 +97,17 @@ export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandler
      ) : (
        // Show different UI based on whether it's already handled
        <>
-          <span className="text-xs text-gray-400 italic">{'Session summarized'}</span>
-
+          <span className="text-xs text-gray-400">{`Your conversation has exceeded the model's context capacity`}</span>
+          <span className="text-xs text-gray-400">{`Messages above this line remain viewable but are not included in the active context`}</span>
          {/* Only show the button if its last message */}
          {shouldAllowSummaryInteraction && (
            <button
              onClick={() => (errorLoadingSummary ? handleRetry() : openSummaryModal())}
              className="text-xs text-textStandard hover:text-textSubtle transition-colors mt-1 flex items-center"
            >
-              {errorLoadingSummary ? 'Retry loading summary' : 'View or edit summary'}
+              {errorLoadingSummary
+                ? 'Retry loading summary'
+                : 'View or edit summary (you may continue your conversation based on the summary)'}
            </button>
          )}
        </>
@@ -3,7 +3,7 @@ import { Message } from '../../types/message';
 import { manageContextFromBackend, convertApiMessageToFrontendMessage } from './index';

 // Define the context management interface
-interface ContextManagerState {
+interface ChatContextManagerState {
  summaryContent: string;
  summarizedThread: Message[];
  isSummaryModalOpen: boolean;
@@ -11,31 +11,73 @@ interface ContextManagerState {
  errorLoadingSummary: boolean;
 }

-interface ContextManagerActions {
+interface ChatContextManagerActions {
  fetchSummary: (messages: Message[]) => Promise<void>;
  updateSummary: (newSummaryContent: string) => void;
  resetMessagesWithSummary: (
    messages: Message[],
-    setMessages: (messages: Message[]) => void
+    setMessages: (messages: Message[]) => void,
+    ancestorMessages: Message[],
+    setAncestorMessages: (messages: Message[]) => void,
+    summaryContent: string
  ) => void;
  openSummaryModal: () => void;
  closeSummaryModal: () => void;
  hasContextLengthExceededContent: (message: Message) => boolean;
+  handleContextLengthExceeded: (
+    messages: Message[],
+    chatId: string,
+    workingDir: string
+  ) => Promise<void>;
 }

 // Create the context
-const ContextManagerContext = createContext<
-  (ContextManagerState & ContextManagerActions) | undefined
+const ChatContextManagerContext = createContext<
+  (ChatContextManagerState & ChatContextManagerActions) | undefined
 >(undefined);

 // Create the provider component
-export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
+export const ChatContextManagerProvider: React.FC<{ children: React.ReactNode }> = ({
+  children,
+}) => {
  const [summaryContent, setSummaryContent] = useState<string>('');
  const [summarizedThread, setSummarizedThread] = useState<Message[]>([]);
  const [isSummaryModalOpen, setIsSummaryModalOpen] = useState<boolean>(false);
  const [isLoadingSummary, setIsLoadingSummary] = useState<boolean>(false);
  const [errorLoadingSummary, setErrorLoadingSummary] = useState<boolean>(false);

+  const handleContextLengthExceeded = async (messages: Message[]): Promise<void> => {
+    setIsLoadingSummary(true);
+    setErrorLoadingSummary(false);
+
+    try {
+      // 2. Now get the summary from the backend
+      const summaryResponse = await manageContextFromBackend({
+        messages: messages,
+        manageAction: 'summarize',
+      });
+
+      // Convert API messages to frontend messages
+      const convertedMessages = summaryResponse.messages.map(
+        (apiMessage) => convertApiMessageToFrontendMessage(apiMessage, false, true) // do not show to user but send to llm
+      );
+
+      // Extract summary from the first message
+      const summaryMessage = convertedMessages[0].content[0];
+      if (summaryMessage.type === 'text') {
+        const summary = summaryMessage.text;
+        setSummaryContent(summary);
+        setSummarizedThread(convertedMessages);
+      }
+
+      setIsLoadingSummary(false);
+    } catch (err) {
+      console.error('Error handling context length exceeded:', err);
+      setErrorLoadingSummary(true);
+      setIsLoadingSummary(false);
+    }
+  };
+
  const fetchSummary = async (messages: Message[]) => {
    setIsLoadingSummary(true);
    setErrorLoadingSummary(false);
@@ -47,8 +89,8 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
      });

      // Convert API messages to frontend messages
-      const convertedMessages = response.messages.map((apiMessage) =>
-        convertApiMessageToFrontendMessage(apiMessage)
+      const convertedMessages = response.messages.map(
+        (apiMessage) => convertApiMessageToFrontendMessage(apiMessage, false, true) // do not show to user but send to llm
      );

      // Extract the summary text from the first message
@@ -101,27 +143,69 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (

  const resetMessagesWithSummary = (
    messages: Message[],
-    setMessages: (messages: Message[]) => void
+    setMessages: (messages: Message[]) => void,
+    ancestorMessages: Message[],
+    setAncestorMessages: (messages: Message[]) => void,
+    summaryContent: string
  ) => {
-    // Update summarizedThread with metadata
-    const updatedSummarizedThread = summarizedThread.map((msg) => ({
-      ...msg,
-      display: false,
-      sendToLLM: true,
-    }));
+    // Create a copy of the summarized thread
+    const updatedSummarizedThread = [...summarizedThread];

-    // Update list of messages with other metadata
-    const updatedMessages = messages.map((msg) => ({
-      ...msg,
-      display: true,
-      sendToLLM: false,
-    }));
+    // Make sure there's at least one message in the summarized thread
+    if (updatedSummarizedThread.length > 0) {
+      // Get the first message
+      const firstMessage = { ...updatedSummarizedThread[0] };

-    // Make a copy that combines both
-    const newMessages = [...updatedMessages, ...updatedSummarizedThread];
+      // Make a copy of the content array
+      const contentCopy = [...firstMessage.content];
+
+      // Assuming the first content item is of type TextContent
+      if (contentCopy.length > 0 && 'text' in contentCopy[0]) {
+        // Update the text with the new summary content
+        contentCopy[0] = {
+          ...contentCopy[0],
+          text: summaryContent,
+        };
+
+        // Update the first message with the new content
+        firstMessage.content = contentCopy;
+
+        // Update the first message in the thread
+        updatedSummarizedThread[0] = firstMessage;
+      }
+    }
+
+    // Update metadata for the summarized thread
+    const finalUpdatedThread = updatedSummarizedThread.map((msg, index) => ({
+      ...msg,
+      display: index === 0, // First message has display: true, others false
+      sendToLLM: true, // All messages have sendToLLM: true
+    }));

    // Update the messages state
-    setMessages(newMessages);
+    setMessages(finalUpdatedThread);
+
+    // If ancestorMessages already has items, extend it instead of replacing it
+    if (ancestorMessages.length > 0) {
+      // Convert current messages to ancestor format
+      const newAncestorMessages = messages.map((msg) => ({
+        ...msg,
+        display: true,
+        sendToLLM: false,
+      }));
+
+      // Append new ancestor messages to existing ones
+      setAncestorMessages([...ancestorMessages, ...newAncestorMessages]);
+    } else {
+      // Initial set of ancestor messages
+      const newAncestorMessages = messages.map((msg) => ({
+        ...msg,
+        display: true,
+        sendToLLM: false,
+      }));
+
+      setAncestorMessages(newAncestorMessages);
+    }

    // Clear the summarized thread and content
    setSummarizedThread([]);
@@ -155,14 +239,19 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
    openSummaryModal,
    closeSummaryModal,
    hasContextLengthExceededContent,
+    handleContextLengthExceeded,
  };

-  return <ContextManagerContext.Provider value={value}>{children}</ContextManagerContext.Provider>;
+  return (
+    <ChatContextManagerContext.Provider value={value}>
+      {children}
+    </ChatContextManagerContext.Provider>
+  );
 };

 // Create a hook to use the context
 export const useChatContextManager = () => {
-  const context = useContext(ContextManagerContext);
+  const context = useContext(ChatContextManagerContext);
  if (context === undefined) {
    throw new Error('useContextManager must be used within a ContextManagerProvider');
  }
@@ -93,6 +93,12 @@ export function SessionSummaryModal({
        ref={textareaRef}
        defaultValue={summaryContent}
        className="bg-gray-50 dark:bg-gray-800 p-4 rounded-lg text-gray-700 dark:text-gray-300 border border-gray-200 dark:border-gray-700 text-sm w-full min-h-[200px] focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent"
+        style={{
+          textRendering: 'optimizeLegibility',
+          WebkitFontSmoothing: 'antialiased',
+          MozOsxFontSmoothing: 'grayscale',
+          transform: 'translateZ(0)', // Force hardware acceleration
+        }}
      />
    </div>
  );
@@ -49,10 +49,14 @@ export async function manageContextFromBackend({
 }

 // Function to convert API Message to frontend Message
-export function convertApiMessageToFrontendMessage(apiMessage: ApiMessage): FrontendMessage {
+export function convertApiMessageToFrontendMessage(
+  apiMessage: ApiMessage,
+  display?: boolean,
+  sendToLLM?: boolean
+): FrontendMessage {
  return {
-    display: false,
-    sendToLLM: false,
+    display: display ?? true,
+    sendToLLM: sendToLLM ?? true,
    id: generateId(),
    role: apiMessage.role,
    created: apiMessage.created,
@@ -35,8 +35,8 @@ const SessionListView: React.FC<SessionListViewProps> = ({ setView, onSelectSess
    setIsLoading(true);
    setError(null);
    try {
-      const response = await fetchSessions();
-      setSessions(response.sessions);
+      const sessions = await fetchSessions();
+      setSessions(sessions);
    } catch (err) {
      console.error('Failed to load sessions:', err);
      setError('Failed to load sessions. Please try again later.');
@@ -47,6 +47,8 @@ const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {

  const handleResumeSession = () => {
    if (selectedSession) {
+      console.log('Selected session object:', JSON.stringify(selectedSession, null, 2));
+
      // Get the working directory from the session metadata
      const workingDir = selectedSession.metadata.working_dir;