Fix chat history reliability: use wisplog as authoritative source

Wisp/ViewModels/ChatViewModel.swift

@@ -711,10 +711,24 @@ final class ChatViewModel {
         let (parsed, parsedSessionId) = Self.parseWispLog(output)
         guard !parsed.isEmpty else { return }
 
+        // Don't overwrite messages if a new streaming session started while the
+        // wisplog was being fetched (user sent another message in the window).
+        guard !isStreaming else { return }
+
         messages = parsed
         if let parsedSessionId { sessionId = parsedSessionId }
         rebuildToolUseIndex()
-        persistMessages(modelContext: modelContext)
+
+        // Mirror restoreFromSessionFile: surface a trailing user message as a draft,
+        // or mark the session complete and clear the exec ID so reconnectIfNeeded
+        // doesn't attempt a pointless reattach on the next open.
+        if let last = messages.last, last.role == .user {
+            restoreUndeliveredDraft(modelContext: modelContext)
+        } else {
+            execSessionId = nil
+            saveSession(modelContext: modelContext, isComplete: true)
+            persistMessages(modelContext: modelContext)
+        }
     }
 
     func persistMessages(modelContext: ModelContext) {
@@ -918,23 +932,31 @@ final class ChatViewModel {
     func reconnectIfNeeded(apiClient: SpritesAPIClient, modelContext: ModelContext) {
         guard !isStreaming else { return }
 
-        if messages.isEmpty {
-            // No local messages — the wisp log file (or legacy JSONL) on the sprite
-            // contains the full conversation. Load it so the chat isn't blank.
-            // This handles the case where SwiftData was cleared or the app was reinstalled.
+        if messages.isEmpty || messages.first?.role == .assistant {
+            // No local messages, or messages appear truncated (first message should always
+            // be from the user — if it's an assistant message the persisted state is corrupt,
+            // likely because a reconnect after backgrounding wrote a partial history starting
+            // mid-conversation). Reload the full conversation from the wisp log file.
             if !isLoadingHistory {
                 Task { await loadFromWispLog(apiClient: apiClient, modelContext: modelContext) }
             }
             return
         }
 
-        // If the last session completed cleanly, content is already loaded from
-        // persistence — no need to hit the network at all.
+        // If the last session completed cleanly, no exec reconnect is needed.
         if let chat = fetchChat(modelContext: modelContext), chat.lastSessionComplete {
-            // Edge case: app killed between persistMessages and saveSession(isComplete:false).
-            // The session was previously complete but a new user message was appended and
-            // persisted before the exec session was created. Restore it as a draft.
-            restoreUndeliveredDraft(modelContext: modelContext)
+            if messages.last?.role == .user {
+                // Edge case: app killed between persistMessages and saveSession(isComplete:false).
+                // A new user message was appended and persisted before the exec session was
+                // created. Restore it as a draft — don't also reload the wisplog here since
+                // the wisplog would include this undelivered prompt and put it back in messages.
+                restoreUndeliveredDraft(modelContext: modelContext)
+            } else if !isLoadingHistory {
+                // Session ended cleanly with an assistant response. Verify against the
+                // wisplog in the background to catch any history truncation from earlier
+                // reconnects or backgrounding that our streaming reload may not have fixed.
+                Task { await loadFromWispLog(apiClient: apiClient, modelContext: modelContext) }
+            }
             return
         }
 
@@ -1129,7 +1151,40 @@ final class ChatViewModel {
             return
         }
 
-        // If timed out with no data, clear Claude lock files and retry once
+        // If timed out with no data, first check whether Claude actually finished and
+        // the output was written to the wisplog before our WebSocket could connect.
+        // This happens when a short response completes in under ~1s (faster than the
+        // WebSocket handshake), causing processExecStream to see no stdout at all.
+        // Loading from the wisplog avoids a redundant retry and a duplicate user prompt.
+        if case .timedOut = streamResult, !retriedAfterTimeout, sessionId != nil {
+            let logPath = Self.wispLogPath(for: chatId)
+            let (wispOutput, wispSuccess) = await apiClient.runExec(
+                spriteName: spriteName,
+                command: "cat \(shellEscape(logPath)) 2>/dev/null",
+                timeout: 15
+            )
+            if wispSuccess, !wispOutput.isEmpty {
+                let (parsed, parsedSessionId) = Self.parseWispLog(wispOutput)
+                // The wisplog has a completed response if it has at least as many messages
+                // as the current array and ends with a non-empty assistant message.
+                if parsed.count >= messages.count,
+                   parsed.last?.role == .assistant,
+                   !(parsed.last?.content.isEmpty ?? true) {
+                    logger.info("[Chat] Exec timed out but wisplog has completed response — loading it")
+                    messages = parsed
+                    if let parsedSessionId { sessionId = parsedSessionId }
+                    rebuildToolUseIndex()
+                    execSessionId = nil
+                    saveSession(modelContext: modelContext, isComplete: true)
+                    persistMessages(modelContext: modelContext)
+                    if !Task.isCancelled { status = .idle }
+                    return
+                }
+            }
+        }
+
+        // If timed out with no data and wisplog had nothing useful, clear Claude lock
+        // files and retry once.
         if case .timedOut = streamResult, !retriedAfterTimeout {
             logger.info("Timeout — clearing Claude lock files and retrying")
             retriedAfterTimeout = true
@@ -1174,6 +1229,13 @@ final class ChatViewModel {
             messages.append(userMessage)
             persistMessages(modelContext: modelContext)
             await executeClaudeCommand(prompt: prompt, apiClient: apiClient, modelContext: modelContext)
+        } else {
+            // Reload from the wisplog to correct any history truncation from reconnects
+            // or backgrounding. The guard in loadFromWispLog aborts if the user sends
+            // another message before the cat completes.
+            let capturedApiClient = apiClient
+            let capturedModelContext = modelContext
+            Task { await loadFromWispLog(apiClient: capturedApiClient, modelContext: capturedModelContext) }
         }
     }
 
@@ -1409,6 +1471,11 @@ final class ChatViewModel {
             messages.append(userMessage)
             persistMessages(modelContext: modelContext)
             await executeClaudeCommand(prompt: prompt, apiClient: apiClient, modelContext: modelContext)
+        } else if case .completed = streamResult, !Task.isCancelled {
+            // Same post-completion wisplog reload as executeClaudeCommand.
+            let capturedApiClient = apiClient
+            let capturedModelContext = modelContext
+            Task { await loadFromWispLog(apiClient: capturedApiClient, modelContext: capturedModelContext) }
         }
     }
 

WispTests/ChatViewModelTests.swift

@@ -821,6 +821,32 @@ struct ChatViewModelTests {
         #expect(vm.streamTask != nil)
     }
 
+    @Test func reconnectIfNeeded_corruptedMessages_firstIsAssistant_triggersWispLogLoad() throws {
+        // Regression: after a reconnect following phone-sleep during streaming,
+        // SwiftData can end up with a partial history that starts with an assistant
+        // message (e.g. the first user+assistant exchange was lost). The wisplog
+        // on the sprite still has the full history. Verify that reconnectIfNeeded
+        // detects this corruption and does NOT try to reattach to a (nonexistent)
+        // exec session — it should take the wisplog-reload path instead.
+        let ctx = try makeModelContext()
+        let (vm, _) = makeChatViewModel(modelContext: ctx)
+
+        // Simulate truncated SwiftData: first message is an assistant (wrong — should be user)
+        vm.messages = [
+            ChatMessage(role: .assistant, content: [.text("Here's the deep dive...")]),
+            ChatMessage(role: .user, content: [.text("Follow-up question")]),
+        ]
+        vm.setExecSessionId("exec-abc") // would normally trigger reattach
+
+        vm.reconnectIfNeeded(apiClient: SpritesAPIClient(), modelContext: ctx)
+
+        // Should NOT have started a reattach task (the exec-session path).
+        // An async wisplog load is scheduled instead — but it's async so we verify
+        // the exec-reattach task was NOT created and messages were not synchronously changed.
+        #expect(vm.streamTask == nil)
+        #expect(vm.messages.count == 2)
+    }
+
     // MARK: - UUID persistence
 
     @Test func persistMessages_savesUUIDsToChat() throws {
@@ -1473,6 +1499,26 @@ struct ChatViewModelTests {
         #expect(messages[1].textContent == "I was saying...")
     }
 
+    @Test func parseWispLog_thinkingOnlyBlockFollowedByTextBlock() {
+        // Regression: extended thinking emits a thinking-only assistant event before the
+        // text event. Both events belong to the same turn and should produce a single
+        // assistant message containing only the text (thinking blocks are ignored).
+        let ndjson = """
+        {"type":"wisp_user_prompt","text":"Do we have similar tests?","timestamp":""}
+        {"type":"system","session_id":"sess-1"}
+        {"type":"assistant","message":{"role":"assistant","content":[{"type":"thinking","thinking":"Let me check..."}]}}
+        {"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Looking back — no, no similar tests exist."}]}}
+        {"type":"result","session_id":"sess-1","is_error":false}
+        """
+
+        let (messages, _) = ChatViewModel.parseWispLog(ndjson)
+
+        #expect(messages.count == 2)
+        #expect(messages[0].role == .user)
+        #expect(messages[1].role == .assistant)
+        #expect(messages[1].textContent == "Looking back — no, no similar tests exist.")
+    }
+
     // MARK: - convertJSONLToWisp
 
     @Test func convertJSONLToWisp_userPromptString() {