Copilot context fixes (#174)

* Update copilot context dynamically

* Fix agent context logic

apps/rowboat/app/projects/[projectId]/copilot/app.tsx

@@ -54,12 +54,18 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
     const cancelRef = useRef<any>(null);
     const [statusBar, setStatusBar] = useState<any>(null);
 
+    // Always use effectiveContext for the user's current selection
+    const effectiveContext = discardContext ? null : chatContext;
+
+    // Context locking state
+    const [lockedContext, setLockedContext] = useState<any>(effectiveContext);
+    const [pendingContext, setPendingContext] = useState<any>(effectiveContext);
+    const [isStreaming, setIsStreaming] = useState(false);
+
     // Keep workflow ref up to date
     workflowRef.current = workflow;
 
-    // Get the effective context based on user preference
+    // Copilot streaming state
-    const effectiveContext = discardContext ? null : chatContext;
-
     const {
         streamingResponse,
         loading: loadingResponse,
@@ -102,13 +108,16 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
         setDiscardContext(false);
     }, [chatContext]);
 
-    function handleUserMessage(prompt: string) {
+    // Memoized handleUserMessage for useImperativeHandle and hooks
+    const handleUserMessage = useCallback((prompt: string) => {
+        // Before starting streaming, lock the context to the current pendingContext
+        setLockedContext(pendingContext);
         setMessages(currentMessages => [...currentMessages, {
             role: 'user',
             content: prompt
         }]);
         setIsLastInteracted(true);
-    }
+    }, [setMessages, setIsLastInteracted, pendingContext, setLockedContext]);
 
     // Effect for getting copilot response
     useEffect(() => {
@@ -134,6 +143,34 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
         return () => currentCancel();
     }, [messages, responseError]);
 
+    // --- CONTEXT LOCKING LOGIC ---
+    // Always update pendingContext to the latest effectiveContext
+    useEffect(() => {
+        setPendingContext(effectiveContext);
+    }, [effectiveContext]);
+
+    // Lock/unlock context based on streaming state
+    useEffect(() => {
+        if (loadingResponse) {
+            // Streaming started: lock context to the value at the start
+            setIsStreaming(true);
+            setLockedContext((prev: any) => prev ?? pendingContext); // lock to previous if already set, else to pending
+        } else {
+            // Streaming ended: update lockedContext to the last pendingContext
+            setIsStreaming(false);
+            setLockedContext(pendingContext);
+        }
+    }, [loadingResponse, pendingContext]);
+
+    // After streaming ends, update lockedContext live as effectiveContext changes
+    useEffect(() => {
+        if (!isStreaming) {
+            setLockedContext(effectiveContext);
+        }
+        // If streaming, do not update lockedContext
+    }, [effectiveContext, isStreaming]);
+    // --- END CONTEXT LOCKING LOGIC ---
+
     const handleCopyChat = useCallback(() => {
         if (onCopyJson) {
             onCopyJson({
@@ -145,7 +182,7 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
     useImperativeHandle(ref, () => ({
         handleCopyChat,
         handleUserMessage
-    }), [handleCopyChat]);
+    }), [handleCopyChat, handleUserMessage]);
 
     return (
         <CopilotContext.Provider value={{ workflow: workflowRef.current, dispatch }}>
@@ -159,8 +196,7 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
                         dispatch={dispatch}
                         onStatusBarChange={status => setStatusBar({
                             ...status,
-                            context: effectiveContext,
+                            context: lockedContext,
-                            onCloseContext: () => setDiscardContext(true)
                         })}
                     />
                 </div>
@@ -187,7 +223,6 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
                             </Button>
                         </div>
                     )}
-                    {/* Remove the separate context label here */}
                     <ComposeBoxCopilot
                         handleUserMessage={handleUserMessage}
                         messages={messages}
@@ -196,7 +231,7 @@ const App = forwardRef<{ handleCopyChat: () => void; handleUserMessage: (message
                         shouldAutoFocus={isLastInteracted}
                         onFocus={() => setIsLastInteracted(true)}
                         onCancel={cancel}
-                        statusBar={statusBar || { context: effectiveContext, onCloseContext: () => setDiscardContext(true) }}
+                        statusBar={statusBar || { context: lockedContext }}
                     />
                 </div>
             </div>

apps/rowboat/app/projects/[projectId]/copilot/components/messages.tsx

@@ -1,6 +1,6 @@
 'use client';
 import { Spinner } from "@heroui/react";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useRef, useState, useCallback } from "react";
 import { z } from "zod";
 import { Workflow} from "@/app/lib/types/workflow_types";
 import MarkdownContent from "@/app/lib/components/markdown-content";
@@ -218,8 +218,8 @@ function AssistantMessage({
     const pendingCount = Math.max(0, totalActions - appliedCount);
     const allApplied = pendingCount === 0 && totalActions > 0;
 
-    // Apply a single action
+    // Memoized applyAction for useCallback dependencies
-    const applyAction = (action: any, actionIndex: number) => {
+    const applyAction = useCallback((action: any, actionIndex: number) => {
         // Only apply, do not update appliedActions here
         if (action.action === 'create_new') {
             switch (action.config_type) {
@@ -276,10 +276,10 @@ function AssistantMessage({
                     break;
             }
         }
-    };
+    }, [dispatch]);
 
-    // Apply All: batch apply all unapplied actions and update state once
+    // Memoized handleApplyAll for useEffect dependencies
-    const handleApplyAll = () => {
+    const handleApplyAll = useCallback(() => {
         // Find all unapplied action indices
         const unapplied = cardBlocks
             .filter(({ part, actionIndex }) => part.type === 'action' && !appliedActions.has(actionIndex))
@@ -296,7 +296,7 @@ function AssistantMessage({
             unapplied.forEach(({ actionIndex }) => next.add(actionIndex));
             return next;
         });
-    };
+    }, [cardBlocks, appliedActions, setAppliedActions, applyAction]);
 
     // Manual single apply (from card)
     const handleSingleApply = (action: any, actionIndex: number) => {
@@ -426,7 +426,7 @@ function AssistantMessage({
                 handleApplyAll,
             });
         }
-    }, [allCardsLoaded, allApplied, appliedCount, pendingCount, streamingLine, completedSummary, hasPanelWarning]);
+    }, [allCardsLoaded, allApplied, appliedCount, pendingCount, streamingLine, completedSummary, hasPanelWarning, handleApplyAll, onStatusBarChange]);
 
     // Render all cards inline, not in a panel
     return (