From ed6e6cd1ed7cfbb210181de5cb484a9a66e30070 Mon Sep 17 00:00:00 2001
From: chihlasm <michael@chihlas.com>
Date: Fri, 3 Apr 2026 04:00:47 +0000
Subject: [PATCH] fix: stabilize cockpit and assistant session handoff

---
 frontend/src/hooks/useAssistantSession.ts | 51 ++++++++++++++++-------
 frontend/src/pages/CockpitPage.tsx        |  9 +++-
 2 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/frontend/src/hooks/useAssistantSession.ts b/frontend/src/hooks/useAssistantSession.ts
index d7cdeba0..3e6bef39 100644
--- a/frontend/src/hooks/useAssistantSession.ts
+++ b/frontend/src/hooks/useAssistantSession.ts
@@ -310,23 +310,34 @@ export function useAssistantSession() {
   // Callback for pages to handle triage updates from chat responses.
   const onTriageUpdateRef = useRef<((update: TriageUpdate) => void) | null>(null)
 
-  const handleSend = async () => {
-    if (!input.trim() || !activeChatId || loadingRef.current) return
+  const sendMessage = useCallback(async (
+    rawMessage: string,
+    options?: {
+      uploadIds?: string[]
+      clearComposer?: boolean
+    }
+  ) => {
+    const message = rawMessage.trim()
+    if (!message || !activeChatId || loadingRef.current) return
     loadingRef.current = true
 
     const sendChatId = activeChatId
-    const userMessage = input.trim()
-    const completedUploadIds = pendingUploads
+    const uploadIds = options?.uploadIds
+    const completedUploadIds = uploadIds ?? pendingUploads
       .filter((u) => u.status === 'done' && u.result?.id)
       .map((u) => u.result!.id)
-    setInput('')
-    setPendingUploads([])
-    setMessages(prev => [...prev, { role: 'user', content: userMessage }])
+    if (options?.clearComposer !== false) {
+      setInput('')
+      if (!uploadIds) {
+        setPendingUploads([])
+      }
+    }
+    setMessages(prev => [...prev, { role: 'user', content: message }])
     setLoading(true)
 
     try {
       const response = await aiSessionsApi.sendChatMessage(sendChatId, {
-        message: userMessage,
+        message,
         upload_ids: completedUploadIds.length > 0 ? completedUploadIds : undefined,
       })
       if (currentChatRef.current !== sendChatId) return
@@ -335,7 +346,7 @@ export function useAssistantSession() {
       setChats(prev =>
         prev.map(c =>
           c.id === sendChatId
-            ? { ...c, message_count: c.message_count + 2, title: c.message_count === 0 ? userMessage.slice(0, 100) : c.title, updated_at: new Date().toISOString() }
+            ? { ...c, message_count: c.message_count + 2, title: c.message_count === 0 ? message.slice(0, 100) : c.title, updated_at: new Date().toISOString() }
             : c
         )
       )
@@ -353,12 +364,17 @@ export function useAssistantSession() {
         requestAnimationFrame(() => inputRef.current?.focus())
       }
     }
+  }, [activeChatId, pendingUploads, processResponse])
+
+  const handleSend = async () => {
+    await sendMessage(input)
   }
 
   // Handle prefill from command palette / dashboard handoff
   const handlePrefill = useCallback((_prefillRoute: string) => {
-    const state = location.state as { prefill?: string; uploadIds?: string[] } | null
+    const state = location.state as { prefill?: string; logs?: string; uploadIds?: string[] } | null
     const prefill = state?.prefill
+    const logs = state?.logs?.trim()
     const uploadIds = state?.uploadIds
     if (!prefill || prefillHandledRef.current) return
     prefillHandledRef.current = true
@@ -372,11 +388,18 @@ export function useAssistantSession() {
       setActiveQuestions([])
       setActiveActions([])
       setLoading(true)
+      if (logs) {
+        setShowLogs(true)
+        setLogContent(logs)
+      }
 
       try {
+        const initialMessage = logs
+          ? `${prefill}\n\nAttached logs/output:\n\`\`\`\n${logs}\n\`\`\``
+          : prefill
         const session = await aiSessionsApi.createChatSession({
           intake_type: 'free_text',
-          intake_content: { text: prefill },
+          intake_content: { text: initialMessage },
         })
         const prefillChatId = session.session_id
         currentChatRef.current = prefillChatId
@@ -390,10 +413,10 @@ export function useAssistantSession() {
         }
         setChats(prev => [chatItem, ...prev])
         setActiveChatId(prefillChatId)
-        setMessages([{ role: 'user', content: prefill }])
+        setMessages([{ role: 'user', content: initialMessage }])
 
         const response = await aiSessionsApi.sendChatMessage(prefillChatId, {
-          message: prefill,
+          message: initialMessage,
           upload_ids: uploadIds?.length ? uploadIds : undefined,
         })
         if (currentChatRef.current !== prefillChatId) return
@@ -561,7 +584,7 @@ export function useAssistantSession() {
     setShowTaskLane, setActiveQuestions, setActiveActions,
     // Handlers
     loadChats, selectChat, handleNewChat, handleDeleteChat, handleTogglePin,
-    handleSend, handleConclude, handleResumeNew,
+    handleSend, sendMessage, handleConclude, handleResumeNew,
     handleKeyDown, handlePaste,
     handleDragOver, handleDragEnter, handleDragLeave, handleDrop,
     handleFileSelect, handleRemoveUpload, retryUpload,
diff --git a/frontend/src/pages/CockpitPage.tsx b/frontend/src/pages/CockpitPage.tsx
index a5370f13..17ce2a01 100644
--- a/frontend/src/pages/CockpitPage.tsx
+++ b/frontend/src/pages/CockpitPage.tsx
@@ -85,6 +85,12 @@ export default function CockpitPage() {
     prevMessageCountRef.current = session.messages.length
   }, [session.messages.length, showOnboarding])
 
+  // Reset local step UI when switching cases or when a new action set arrives.
+  useEffect(() => {
+    setActiveStepIndex(0)
+    setCompletedSteps(new Set())
+  }, [session.activeChatId, session.activeActions])
+
   // ── Triage handlers ──
 
   const handleTriageFieldSave = useCallback(async (field: keyof TriageMeta, value: string) => {
@@ -306,8 +312,7 @@ export default function CockpitPage() {
                   <FlowPilotAsks
                     questions={session.activeQuestions}
                     onAnswer={(answer) => {
-                      session.setInput(answer)
-                      setTimeout(() => session.handleSend(), 10)
+                      void session.sendMessage(answer, { clearComposer: false })
                     }}
                     loading={session.loading}
                   />