diff --git a/Wisp/ViewModels/ChatViewModel.swift b/Wisp/ViewModels/ChatViewModel.swift index 0757650..14014f9 100644 --- a/Wisp/ViewModels/ChatViewModel.swift +++ b/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) } } } diff --git a/WispTests/ChatViewModelTests.swift b/WispTests/ChatViewModelTests.swift index a067e98..6eaf991 100644 --- a/WispTests/ChatViewModelTests.swift +++ b/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() {