From 11b4ff7789a6ec5a0f805c42da961c6a5170bfe2 Mon Sep 17 00:00:00 2001
From: Jack Amadeo <jackamadeo@squareup.com>
Date: Thu, 21 Aug 2025 22:38:26 -0400
Subject: [PATCH] Block send until extensions are ready

---
 ui/desktop/src/App.tsx                    | 22 ++++--
 ui/desktop/src/components/ChatInput.tsx   | 93 +++++++++++++----------
 ui/desktop/src/contexts/ChatContext.tsx   |  4 +
 ui/desktop/src/utils/appInitialization.ts |  3 +
 4 files changed, 76 insertions(+), 46 deletions(-)

diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx
index e071e3595ad3..887c3b4e7cb2 100644
--- a/ui/desktop/src/App.tsx
+++ b/ui/desktop/src/App.tsx
@@ -406,6 +406,7 @@ export default function App() {
   const [isGoosehintsModalOpen, setIsGoosehintsModalOpen] = useState(false);
   const [isLoadingSharedSession, setIsLoadingSharedSession] = useState(false);
   const [sharedSessionError, setSharedSessionError] = useState<string | null>(null);
+  const [agentWaitingMessage, setAgentWaitingMessage] = useState<string | null>(null);
 
   // Add separate state for pair chat to maintain its own conversation
   const [pairChat, setPairChat] = useState<ChatType>({
@@ -499,16 +500,19 @@ export default function App() {
         getExtensions,
         addExtension,
         setPairChat,
+        setMessage: setAgentWaitingMessage,
         provider: provider as string,
         model: model as string,
       });
     };
 
-    initialize().catch((error) => {
-      console.error('Fatal error during initialization:', error);
-      setFatalError(error instanceof Error ? error.message : 'Unknown error occurred');
-    });
-  }, [getExtensions, addExtension, read, setPairChat]);
+    initialize()
+      .then(() => setAgentWaitingMessage(null))
+      .catch((error) => {
+        console.error('Fatal error during initialization:', error);
+        setFatalError(error instanceof Error ? error.message : 'Unknown error occurred');
+      });
+  }, [getExtensions, addExtension, read, setPairChat, setAgentWaitingMessage]);
 
   useEffect(() => {
     console.log('Sending reactReady signal to Electron');
@@ -926,7 +930,12 @@ export default function App() {
               <Route
                 path="/"
                 element={
-                  <ChatProvider chat={chat} setChat={setChat} contextKey="hub">
+                  <ChatProvider
+                    chat={chat}
+                    setChat={setChat}
+                    contextKey="hub"
+                    agentWaitingMessage={agentWaitingMessage}
+                  >
                     <AppLayout setIsGoosehintsModalOpen={setIsGoosehintsModalOpen} />
                   </ChatProvider>
                 }
@@ -952,6 +961,7 @@ export default function App() {
                         chat={pairChat}
                         setChat={setPairChat}
                         contextKey={`pair-${pairChat.id}`}
+                        agentWaitingMessage={agentWaitingMessage}
                         key={pairChat.id}
                       >
                         <PairRouteWrapper
diff --git a/ui/desktop/src/components/ChatInput.tsx b/ui/desktop/src/components/ChatInput.tsx
index a73898968e9d..17ec5fc78b47 100644
--- a/ui/desktop/src/components/ChatInput.tsx
+++ b/ui/desktop/src/components/ChatInput.tsx
@@ -120,6 +120,7 @@ export default function ChatInput({
   // Draft functionality - get chat context and global draft context
   // We need to handle the case where ChatInput is used without ChatProvider (e.g., in Hub)
   const chatContext = useChatContext(); // This should always be available now
+  const agentIsReady = chatContext === null || chatContext.agentWaitingMessage === null;
   const draftLoadedRef = useRef(false);
 
   // Debug logging for draft context
@@ -895,6 +896,7 @@ export default function ChatInput({
       const canSubmit =
         !isLoading &&
         !isLoadingCompaction &&
+        agentIsReady &&
         (displayValue.trim() ||
           pastedImages.some((img) => img.filePath && !img.error && !img.isLoading) ||
           allDroppedFiles.some((file) => !file.error && !file.isLoading));
@@ -909,6 +911,7 @@ export default function ChatInput({
     const canSubmit =
       !isLoading &&
       !isLoadingCompaction &&
+      agentIsReady &&
       (displayValue.trim() ||
         pastedImages.some((img) => img.filePath && !img.error && !img.isLoading) ||
         allDroppedFiles.some((file) => !file.error && !file.isLoading));
@@ -1089,46 +1092,56 @@ export default function ChatInput({
               <Stop />
             </Button>
           ) : (
-            <Button
-              type="submit"
-              size="sm"
-              shape="round"
-              variant="outline"
-              disabled={
-                !hasSubmittableContent ||
-                isAnyImageLoading ||
-                isAnyDroppedFileLoading ||
-                isRecording ||
-                isTranscribing ||
-                isLoadingCompaction
-              }
-              className={`rounded-full px-10 py-2 flex items-center gap-2 ${
-                !hasSubmittableContent ||
-                isAnyImageLoading ||
-                isAnyDroppedFileLoading ||
-                isRecording ||
-                isTranscribing ||
-                isLoadingCompaction
-                  ? 'bg-slate-600 text-white cursor-not-allowed opacity-50 border-slate-600'
-                  : 'bg-slate-600 text-white hover:bg-slate-700 border-slate-600 hover:cursor-pointer'
-              }`}
-              title={
-                isLoadingCompaction
-                  ? 'Summarizing conversation...'
-                  : isAnyImageLoading
-                    ? 'Waiting for images to save...'
-                    : isAnyDroppedFileLoading
-                      ? 'Processing dropped files...'
-                      : isRecording
-                        ? 'Recording...'
-                        : isTranscribing
-                          ? 'Transcribing...'
-                          : 'Send'
-              }
-            >
-              <Send className="w-4 h-4" />
-              <span className="text-sm">Send</span>
-            </Button>
+            <Tooltip>
+              <TooltipTrigger asChild>
+                <span>
+                  <Button
+                    type="submit"
+                    size="sm"
+                    shape="round"
+                    variant="outline"
+                    disabled={
+                      !hasSubmittableContent ||
+                      isAnyImageLoading ||
+                      isAnyDroppedFileLoading ||
+                      isRecording ||
+                      isTranscribing ||
+                      isLoadingCompaction ||
+                      !agentIsReady
+                    }
+                    className={`rounded-full px-10 py-2 flex items-center gap-2 ${
+                      !hasSubmittableContent ||
+                      isAnyImageLoading ||
+                      isAnyDroppedFileLoading ||
+                      isRecording ||
+                      isTranscribing ||
+                      isLoadingCompaction ||
+                      !agentIsReady
+                        ? 'bg-slate-600 text-white cursor-not-allowed opacity-50 border-slate-600'
+                        : 'bg-slate-600 text-white hover:bg-slate-700 border-slate-600 hover:cursor-pointer'
+                    }`}
+                  >
+                    <Send className="w-4 h-4" />
+                    <span className="text-sm">Send</span>
+                  </Button>
+                </span>
+              </TooltipTrigger>
+              <TooltipContent>
+                <p>
+                  {isLoadingCompaction
+                    ? 'Summarizing conversation...'
+                    : isAnyImageLoading
+                      ? 'Waiting for images to save...'
+                      : isAnyDroppedFileLoading
+                        ? 'Processing dropped files...'
+                        : isRecording
+                          ? 'Recording...'
+                          : isTranscribing
+                            ? 'Transcribing...'
+                            : (chatContext?.agentWaitingMessage ?? 'Send')}
+                </p>
+              </TooltipContent>
+            </Tooltip>
           )}
 
           {/* Recording/transcribing status indicator - positioned above the button row */}
diff --git a/ui/desktop/src/contexts/ChatContext.tsx b/ui/desktop/src/contexts/ChatContext.tsx
index 8ae5ccd47d72..e9f1daabff71 100644
--- a/ui/desktop/src/contexts/ChatContext.tsx
+++ b/ui/desktop/src/contexts/ChatContext.tsx
@@ -21,6 +21,7 @@ interface ChatContextType {
   clearDraft: () => void;
   // Context identification
   contextKey: string; // 'hub' or 'pair-{sessionId}'
+  agentWaitingMessage: string | null;
 }
 
 const ChatContext = createContext<ChatContextType | undefined>(undefined);
@@ -30,12 +31,14 @@ interface ChatProviderProps {
   chat: ChatType;
   setChat: (chat: ChatType) => void;
   contextKey?: string; // Optional context key, defaults to 'hub'
+  agentWaitingMessage: string | null;
 }
 
 export const ChatProvider: React.FC<ChatProviderProps> = ({
   children,
   chat,
   setChat,
+  agentWaitingMessage,
   contextKey = 'hub',
 }) => {
   const draftContext = useDraftContext();
@@ -108,6 +111,7 @@ export const ChatProvider: React.FC<ChatProviderProps> = ({
     setDraft,
     clearDraft,
     contextKey,
+    agentWaitingMessage,
   };
 
   return <ChatContext.Provider value={value}>{children}</ChatContext.Provider>;
diff --git a/ui/desktop/src/utils/appInitialization.ts b/ui/desktop/src/utils/appInitialization.ts
index 8069352acd6e..054142036cac 100644
--- a/ui/desktop/src/utils/appInitialization.ts
+++ b/ui/desktop/src/utils/appInitialization.ts
@@ -14,6 +14,7 @@ interface InitializationDependencies {
   getExtensions?: (b: boolean) => Promise<FixedExtensionEntry[]>;
   addExtension?: (name: string, config: ExtensionConfig, enabled: boolean) => Promise<void>;
   setPairChat: (chat: ChatType | ((prev: ChatType) => ChatType)) => void;
+  setMessage: (message: string | null) => void;
   provider: string;
   model: string;
 }
@@ -22,6 +23,7 @@ export const initializeApp = async ({
   getExtensions,
   addExtension,
   setPairChat,
+  setMessage,
   provider,
   model,
 }: InitializationDependencies) => {
@@ -87,6 +89,7 @@ export const initializeApp = async ({
         initPromises.push(costDbPromise);
       }
 
+      setMessage('starting extensions...');
       await Promise.all(initPromises);
     } catch (error) {
       console.error('Error in system initialization:', error);