mito-ai: fix agent target notebook

Author: aarondr77

mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx

@@ -389,10 +389,6 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 0: reset the state for a new message
         resetForNewMessage()
 
-        // Ensure a notebook exists before proceeding
-        const agentTargetNotebookPanel = await ensureNotebookExists(notebookTracker, documentManager);
-        agentTargetNotebookPanelRef.current = agentTargetNotebookPanel;
-
         // Step 1: Add the user's message to the chat history
         const newChatHistoryManager = getDuplicateChatHistoryManager()
 
@@ -403,15 +399,15 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
 
         const agentExecutionMetadata = newChatHistoryManager.addAgentExecutionMessage(
             activeThreadIdRef.current, 
-            agentTargetNotebookPanel,
+            agentTargetNotebookPanelRef.current,
             input,
             additionalContext
         )
         if (messageIndex !== undefined) {
             agentExecutionMetadata.index = messageIndex
         }
 
-        agentExecutionMetadata.base64EncodedActiveCellOutput = await getBase64EncodedCellOutputInNotebook(agentTargetNotebookPanel, sendCellIDOutput)
+        agentExecutionMetadata.base64EncodedActiveCellOutput = await getBase64EncodedCellOutputInNotebook(agentTargetNotebookPanelRef.current, sendCellIDOutput)
 
         setChatHistoryManager(newChatHistoryManager)
         setLoadingStatus('thinking');
@@ -708,6 +704,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         app,
         streamlitPreviewManager,
         websocketClient,
+        documentManager,
         chatHistoryManagerRef,
         activeThreadIdRef,
         activeRequestControllerRef,

mito-ai/src/Extensions/AiChat/hooks/useAgentExecution.ts

@@ -6,6 +6,7 @@
 import { useRef, useState } from 'react';
 import { JupyterFrontEnd } from '@jupyterlab/application';
 import { INotebookTracker } from '@jupyterlab/notebook';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { UUID } from '@lumino/coreutils';
 import { IStreamlitPreviewManager } from '../../AppPreview/StreamlitPreviewPlugin';
 import { CompletionWebsocketClient } from '../../../websockets/completions/CompletionsWebsocketClient';
@@ -18,6 +19,7 @@ import { getCodeBlockFromMessage } from '../../../utils/strings';
 import { getAIOptimizedCellsInNotebookPanel, setActiveCellByIDInNotebookPanel } from '../../../utils/notebook';
 import { AgentReviewStatus } from '../ChatTaskpane';
 import { LoadingStatus } from './useChatState';
+import { ensureNotebookExists } from '../utils';
 
 export type AgentExecutionStatus = 'working' | 'stopping' | 'idle';
 
@@ -28,6 +30,7 @@ interface UseAgentExecutionProps {
     app: JupyterFrontEnd;
     streamlitPreviewManager: IStreamlitPreviewManager;
     websocketClient: CompletionWebsocketClient;
+    documentManager: IDocumentManager;
     chatHistoryManagerRef: React.MutableRefObject<ChatHistoryManager>;
     activeThreadIdRef: React.MutableRefObject<string>;
     activeRequestControllerRef: React.MutableRefObject<AbortController | null>;
@@ -63,6 +66,7 @@ export const useAgentExecution = ({
     app,
     streamlitPreviewManager,
     websocketClient,
+    documentManager,
     chatHistoryManagerRef,
     activeThreadIdRef,
     activeRequestControllerRef,
@@ -133,7 +137,10 @@ export const useAgentExecution = ({
         messageIndex?: number,
         additionalContext?: Array<{ type: string, value: string }>
     ): Promise<void> => {
-        agentTargetNotebookPanelRef.current = notebookTracker.currentWidget;
+        
+        // Ensure a notebook exists before proceeding with agent execution
+        const agentTargetNotebookPanel = await ensureNotebookExists(notebookTracker, documentManager);
+        agentTargetNotebookPanelRef.current = agentTargetNotebookPanel;
 
         agentReview.acceptAllAICode();
         agentReview.setNotebookSnapshotPreAgentExecution(getAIOptimizedCellsInNotebookPanel(agentTargetNotebookPanelRef.current));