Keep chat streams alive across navigation

Wisp.xcodeproj/project.pbxproj

@@ -14,6 +14,7 @@
 		3855BDFB5819B6D35F57ED9E /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 099181E9B8BFE6F3475C5D36 /* Assets.xcassets */; };
 		3A36571818CB2EBA1ED78500 /* ToolInputDetailView.swift in Sources */ = {isa = PBXBuildFile; fileRef = AB04F3FB6C647DE1F6DC0FD0 /* ToolInputDetailView.swift */; };
 		4FE7C138A35A53B4C38BEF98 /* ClaudeStreamParser.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8F1378232DC7FA71A8B5444 /* ClaudeStreamParser.swift */; };
+		AA33BB44CC55DD66EE77FF02 /* ChatSessionManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA33BB44CC55DD66EE77FF01 /* ChatSessionManager.swift */; };
 		5A1F582DDCBB994013E4BA2E /* DashboardView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0B99AA33329D50DEDB7343B7 /* DashboardView.swift */; };
 		5ABE241326BF34FF2E3C3E51 /* CreateSpriteSheet.swift in Sources */ = {isa = PBXBuildFile; fileRef = EA1634863D8543868BA11B02 /* CreateSpriteSheet.swift */; };
 		61CBE460AC96CBE6BF8AF0B0 /* SpriteSession.swift in Sources */ = {isa = PBXBuildFile; fileRef = D6BE6E0D68BBB35DE56600E1 /* SpriteSession.swift */; };
@@ -45,6 +46,7 @@
 		AA3400BB00CC00DD00EE0001 /* QuickChatView.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3400BB00CC00DD00EE0002 /* QuickChatView.swift */; };
 		AA3500BB00CC00DD00EE0001 /* BashQuickView.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3500BB00CC00DD00EE0002 /* BashQuickView.swift */; };
 		AA3700BB00CC00DD00EE0001 /* BashQuickViewModelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */; };
+		AA3800BB00CC00DD00EE0001 /* ChatSessionManagerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */; };
 		FE9876543210ABCDEF987601 /* QuickChatViewModelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */; };
 		AA00BB11CC22DD3300000001 /* WispAskCard.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA00BB11CC22DD3300000002 /* WispAskCard.swift */; };
 		AA11BB22CC33DD44EE55FF02 /* GitHubDeviceFlowClient.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA11BB22CC33DD44EE55FF01 /* GitHubDeviceFlowClient.swift */; };
@@ -210,6 +212,7 @@
 		C4D5E6F7A8B90A41637E8FA1 /* CheckpointMarkerView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CheckpointMarkerView.swift; sourceTree = "<group>"; };
 		C8C4C3A1F2B2D57313431E50 /* Extensions.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Extensions.swift; sourceTree = "<group>"; };
 		C8F1378232DC7FA71A8B5444 /* ClaudeStreamParser.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ClaudeStreamParser.swift; sourceTree = "<group>"; };
+		AA33BB44CC55DD66EE77FF01 /* ChatSessionManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatSessionManager.swift; sourceTree = "<group>"; };
 		CC11DD22EE33FF4400110002 /* ChatAttachmentButton.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatAttachmentButton.swift; sourceTree = "<group>"; };
 		CC11DD22EE33FF4400110004 /* SpriteFileBrowserView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SpriteFileBrowserView.swift; sourceTree = "<group>"; };
 		CC11DD22EE33FF4400110005 /* FileEntryTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FileEntryTests.swift; sourceTree = "<group>"; };
@@ -249,6 +252,7 @@
 		AA3400BB00CC00DD00EE0002 /* QuickChatView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QuickChatView.swift; sourceTree = "<group>"; };
 		AA3500BB00CC00DD00EE0002 /* BashQuickView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BashQuickView.swift; sourceTree = "<group>"; };
 		AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BashQuickViewModelTests.swift; sourceTree = "<group>"; };
+		AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatSessionManagerTests.swift; sourceTree = "<group>"; };
 		FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QuickChatViewModelTests.swift; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
@@ -307,6 +311,7 @@
 			isa = PBXGroup;
 			children = (
 				AA22BB33CC44DD55EE66FF02 /* ClaudeQuestionTool.swift */,
+				AA33BB44CC55DD66EE77FF01 /* ChatSessionManager.swift */,
 				C8F1378232DC7FA71A8B5444 /* ClaudeStreamParser.swift */,
 				2554BB4DE8EC315CA48767BF /* ExecSession.swift */,
 				BB11CC22DD33EE44FF660A01 /* GitHubAPIClient.swift */,
@@ -471,6 +476,7 @@
 				EE22FF334455AA66BB770001 /* ExecSessionURLTests.swift */,
 				CC11DD22EE33FF4400110005 /* FileEntryTests.swift */,
 				AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */,
+				AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */,
 				FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */,
 			);
 			path = WispTests;
@@ -728,6 +734,7 @@
 				D2B1EF5C6A7F8B9304C5E6F7 /* ClaudeModel.swift in Sources */,
 				2A18611EA392DA2BFFE74250 /* ClaudeStreamEvent.swift in Sources */,
 				4FE7C138A35A53B4C38BEF98 /* ClaudeStreamParser.swift in Sources */,
+				AA33BB44CC55DD66EE77FF02 /* ChatSessionManager.swift in Sources */,
 				99BA5812AF56D1C7501D1F0C /* CreateCheckpointSheet.swift in Sources */,
 				5ABE241326BF34FF2E3C3E51 /* CreateSpriteSheet.swift in Sources */,
 				5A1F582DDCBB994013E4BA2E /* DashboardView.swift in Sources */,
@@ -813,6 +820,7 @@
 				0DFF41BB14AC3922E95DF781 /* WorktreeTests.swift in Sources */,
 				CC11DD22EE33FF4400110006 /* FileEntryTests.swift in Sources */,
 				AA3700BB00CC00DD00EE0001 /* BashQuickViewModelTests.swift in Sources */,
+				AA3800BB00CC00DD00EE0001 /* ChatSessionManagerTests.swift in Sources */,
 				FE9876543210ABCDEF987601 /* QuickChatViewModelTests.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;

Wisp/App/WispApp.swift

@@ -5,6 +5,7 @@ import SwiftUI
 struct WispApp: App {
     @State private var apiClient = SpritesAPIClient()
     @State private var browserCoordinator = InAppBrowserCoordinator()
+    @State private var chatSessionManager = ChatSessionManager()
     @AppStorage("theme") private var theme: String = "system"
 
     init() {
@@ -27,6 +28,7 @@ struct WispApp: App {
             RootView()
                 .environment(apiClient)
                 .environment(browserCoordinator)
+                .environment(chatSessionManager)
                 .preferredColorScheme(preferredColorScheme)
                 .onChange(of: apiClient.isAuthenticated, initial: true) {
                     browserCoordinator.authToken = apiClient.spritesToken

Wisp/Services/ChatSessionManager.swift

@@ -0,0 +1,58 @@
+import Foundation
+import SwiftData
+
+/// App-wide cache of ChatViewModels, keyed by chat UUID.
+/// Keeps streams alive across chat switches and sprite navigation so connections
+/// are only torn down on explicit actions (interrupt, chat deletion, app termination).
+@Observable
+@MainActor
+final class ChatSessionManager {
+    private var cache: [UUID: ChatViewModel] = [:]
+
+    /// Returns the cached VM for a chat, or creates and loads a new one.
+    func viewModel(
+        for chat: SpriteChat,
+        spriteName: String,
+        apiClient: SpritesAPIClient,
+        modelContext: ModelContext
+    ) -> ChatViewModel {
+        if let existing = cache[chat.id] {
+            return existing
+        }
+        let vm = ChatViewModel(
+            spriteName: spriteName,
+            chatId: chat.id,
+            currentServiceName: chat.currentServiceName,
+            workingDirectory: chat.workingDirectory
+        )
+        vm.loadSession(apiClient: apiClient, modelContext: modelContext)
+        cache[chat.id] = vm
+        return vm
+    }
+
+    func isStreaming(chatId: UUID) -> Bool {
+        cache[chatId]?.isStreaming ?? false
+    }
+
+    /// Called when the app returns to the foreground — reconnects any VM that had a live stream.
+    func resumeAllAfterBackground(apiClient: SpritesAPIClient, modelContext: ModelContext) {
+        for vm in cache.values {
+            vm.resumeAfterBackground(apiClient: apiClient, modelContext: modelContext)
+            vm.reconnectIfNeeded(apiClient: apiClient, modelContext: modelContext)
+        }
+    }
+
+    /// Removes and detaches a VM when its chat is deleted.
+    func remove(chatId: UUID, modelContext: ModelContext) {
+        cache[chatId]?.detach(modelContext: modelContext)
+        cache.removeValue(forKey: chatId)
+    }
+
+    /// Detaches and removes all cached VMs. Call when clearing all chats.
+    func detachAll(modelContext: ModelContext) {
+        for vm in cache.values {
+            vm.detach(modelContext: modelContext)
+        }
+        cache.removeAll()
+    }
+}

Wisp/ViewModels/ChatViewModel.swift

@@ -594,7 +594,7 @@ final class ChatViewModel {
         // Cancel the stale stream and reconnect via service logs
         streamTask?.cancel()
         streamTask = Task {
-            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: true)
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped: false)
         }
     }
 
@@ -642,19 +642,41 @@ final class ChatViewModel {
         // If the last session completed cleanly, content is already loaded from
         // persistence — no need to hit the network at all.
         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 service was created. Restore it as a draft.
+            restoreUndeliveredDraft(modelContext: modelContext)
             return
         }
 
         streamTask = Task {
             // Only reconnect if the service exists (running or stopped with logs)
             guard let serviceInfo = try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName)
-            else { return }
+            else {
+                // Service was never created (e.g. app killed before the service call).
+                // Restore any trailing user message as a draft rather than leaving a
+                // stale bubble with no response.
+                restoreUndeliveredDraft(modelContext: modelContext)
+                return
+            }
 
-            let isRunning = serviceInfo.state.status == "running"
-            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: isRunning)
+            let alreadyStopped = serviceInfo.state.status != "running"
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped: alreadyStopped)
         }
     }
 
+    /// If the last message is a user message with no response, remove it from history
+    /// and restore its text to the input box as a draft.
+    func restoreUndeliveredDraft(modelContext: ModelContext) {
+        guard let last = messages.last, last.role == .user else { return }
+        let text = last.textContent
+        messages.removeLast()
+        persistMessages(modelContext: modelContext)
+        guard !text.isEmpty, inputText.isEmpty else { return }
+        inputText = text
+        saveDraft(modelContext: modelContext)
+    }
+
     // MARK: - Private
 
     private func executeClaudeCommand(
@@ -796,7 +818,7 @@ final class ChatViewModel {
         // Attempt reconnection on disconnect
         if case .disconnected = streamResult {
             logger.info("[Chat] Disconnected mid-stream, attempting reconnect via service logs")
-            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: true)
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped: false)
             return
         }
 
@@ -933,6 +955,7 @@ final class ChatViewModel {
                     receivedData = true
                     timeoutTask.cancel()
                     if case .connecting = status { status = .streaming }
+                    else if case .reconnecting = status { status = .streaming }
 
                     // Two-level NDJSON: ServiceLogEvent.data contains Claude NDJSON.
                     // The logs endpoint prefixes each line with a timestamp
@@ -961,6 +984,7 @@ final class ChatViewModel {
                     receivedData = true
                     timeoutTask.cancel()
                     if case .connecting = status { status = .streaming }
+                    else if case .reconnecting = status { status = .streaming }
                     if let text = event.data {
                         logger.warning("Service stderr: \(text.prefix(500), privacy: .public)")
                     }
@@ -994,6 +1018,7 @@ final class ChatViewModel {
 
                 case .started:
                     if case .connecting = status { status = .streaming }
+                    else if case .reconnecting = status { status = .streaming }
 
                 case .stopping, .stopped:
                     break
@@ -1038,10 +1063,10 @@ final class ChatViewModel {
     func runReconnectLoop(
         apiClient: some ServiceLogsProvider,
         modelContext: ModelContext,
-        serviceIsRunning: Bool = false
+        serviceAlreadyStopped: Bool = false
     ) async {
         isReplaying = true
-        isReplayingLiveService = serviceIsRunning
+        isReplayingLiveService = !serviceAlreadyStopped
         defer {
             applyReplayBuffer()
             isReplaying = false
@@ -1060,7 +1085,11 @@ final class ChatViewModel {
         if let existing = currentAssistantMessage {
             assistantMessage = existing
             if !hasPriorEvents { assistantMessage.content = [] }
-        } else if let last = messages.last(where: { $0.role == .assistant }) {
+        } else if let last = messages.last, last.role == .assistant {
+            // Only reuse an assistant message if it's already at the tail — meaning
+            // the last exchange ended with a partial response that needs replaying.
+            // If the last message is a user message, fall through to create a new
+            // assistant message after it (each service is scoped to one user message).
             assistantMessage = last
             if !hasPriorEvents { assistantMessage.content = [] }
             currentAssistantMessage = last
@@ -1108,10 +1137,12 @@ final class ChatViewModel {
             // If we got a result event, Claude is done
             if receivedResultEvent { break }
 
-            // Check if service is still running
-            let isRunning = (try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName))?.state.status == "running"
+            // If we already knew the service was stopped before entering, no need to re-check
+            if serviceAlreadyStopped { break }
 
-            if isRunning {
+            // Check if service is still running before retrying
+            if let serviceInfo = try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName),
+               serviceInfo.state.status == "running" {
                 logger.info("[Chat] Service still running, will re-poll after delay")
                 try? await Task.sleep(for: .seconds(2))
                 continue
@@ -1147,9 +1178,9 @@ final class ChatViewModel {
     private func reconnectToServiceLogs(
         apiClient: SpritesAPIClient,
         modelContext: ModelContext,
-        serviceIsRunning: Bool = false
+        serviceAlreadyStopped: Bool = false
     ) async {
-        await runReconnectLoop(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: serviceIsRunning)
+        await runReconnectLoop(apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped: serviceAlreadyStopped)
 
         if let queued = queuedPrompt, !Task.isCancelled {
             let prompt = buildPrompt(text: queued, attachments: queuedAttachments)

Wisp/Views/SpriteDetail/Chat/ChatSwitcherSheet.swift

@@ -4,6 +4,7 @@ import SwiftData
 struct ChatSwitcherSheet: View {
     @Bindable var viewModel: SpriteChatListViewModel
     @Environment(SpritesAPIClient.self) private var apiClient
+    @Environment(ChatSessionManager.self) private var chatSessionManager
     @Environment(\.modelContext) private var modelContext
     @Environment(\.dismiss) private var dismiss
     @State private var chatToDelete: SpriteChat?
@@ -39,6 +40,7 @@ struct ChatSwitcherSheet: View {
                         }
                         if !chat.isClosed {
                             Button {
+                                chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
                                 viewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
                             } label: {
                                 Label("Close", systemImage: "xmark.circle")
@@ -71,6 +73,7 @@ struct ChatSwitcherSheet: View {
 
                         if !chat.isClosed {
                             Button {
+                                chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
                                 viewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
                             } label: {
                                 Label("Close", systemImage: "xmark.circle")
@@ -87,6 +90,7 @@ struct ChatSwitcherSheet: View {
                         titleVisibility: .visible
                     ) {
                         Button("Delete", role: .destructive) {
+                            chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
                             viewModel.deleteChat(chat, apiClient: apiClient, modelContext: modelContext)
                             chatToDelete = nil
                         }