Merge pull request #65 from mcintyre94/chat-reconnect-delay-troubleshoot-cd59686f

Batch reconnect replay into a single observable mutation

Author: mcintyre94

Wisp/Models/Local/SpriteChat.swift

@@ -21,6 +21,7 @@ final class SpriteChat {
     var forkContext: String?
     var worktreePath: String?
     var worktreeBranch: String?
+    var lastSessionComplete: Bool = false
 
     var displayName: String {
         customName ?? "Chat \(chatNumber)"

Wisp/ViewModels/ChatViewModel.swift

@@ -6,7 +6,7 @@ import UIKit
 
 private let logger = Logger(subsystem: "com.wisp.app", category: "Chat")
 
-enum ChatStatus: Sendable {
+enum ChatStatus: Sendable, Equatable {
     case idle
     case connecting
     case streaming
@@ -81,6 +81,17 @@ final class ChatViewModel {
     var processedEventUUIDs: Set<String> = []
     private var hasPlayedFirstTextHaptic = false
 
+    /// When true, handleEvent buffers content changes instead of mutating the
+    /// @Observable ChatMessage directly. Flushed in one shot after replay ends,
+    /// reducing the N per-event SwiftUI re-renders to a single update.
+    private var isReplaying = false
+    private var replayContentBuffer: [ChatContent] = []
+    private var replayToolUseBuffer: [String: (messageIndex: Int, toolName: String)] = [:]
+    /// When true (reconnecting to a still-running service), isReplaying is cleared
+    /// on the first new event so live events stream in incrementally rather than
+    /// being batched until the connection drops.
+    private var isReplayingLiveService = false
+
     init(spriteName: String, chatId: UUID, currentServiceName: String?, workingDirectory: String) {
         self.spriteName = spriteName
         self.chatId = chatId
@@ -553,7 +564,7 @@ final class ChatViewModel {
         // Cancel the stale stream and reconnect via service logs
         streamTask?.cancel()
         streamTask = Task {
-            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext)
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: true)
         }
     }
 
@@ -600,12 +611,19 @@ final class ChatViewModel {
     func reconnectIfNeeded(apiClient: SpritesAPIClient, modelContext: ModelContext) {
         guard !isStreaming, !messages.isEmpty else { return }
 
+        // 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 {
+            return
+        }
+
         streamTask = Task {
             // Only reconnect if the service exists (running or stopped with logs)
-            guard let _ = try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName)
+            guard let serviceInfo = try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName)
             else { return }
 
-            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext)
+            let isRunning = serviceInfo.state.status == "running"
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: isRunning)
         }
     }
 
@@ -632,8 +650,8 @@ final class ChatViewModel {
             }
         }
 
-        // Persist the new service name immediately for reconnect
-        saveSession(modelContext: modelContext)
+        // Persist the new service name immediately for reconnect; clear any prior completion flag
+        saveSession(modelContext: modelContext, isComplete: false)
 
         guard let claudeToken = apiClient.claudeToken else {
             status = .error("No Claude token configured")
@@ -749,7 +767,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)
+            await reconnectToServiceLogs(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: true)
             return
         }
 
@@ -863,6 +881,15 @@ final class ChatViewModel {
             if let uuid = parsedEvent.uuid {
                 processedEventUUIDs.insert(uuid)
             }
+            // First new event after skipping historical ones on a live-service reconnect:
+            // flush the replay buffer (empty, since skipped events don't buffer anything)
+            // and switch to incremental rendering so live events stream in as they arrive.
+            if isReplayingLiveService {
+                applyReplayBuffer()
+                isReplaying = false
+                isReplayingLiveService = false
+                if case .reconnecting = status { status = .streaming }
+            }
             handleEvent(parsedEvent, modelContext: modelContext)
         }
 
@@ -981,8 +1008,16 @@ final class ChatViewModel {
     /// `reconnectToServiceLogs` so it can be tested against a mock API client.
     func runReconnectLoop(
         apiClient: some ServiceLogsProvider,
-        modelContext: ModelContext
+        modelContext: ModelContext,
+        serviceIsRunning: Bool = false
     ) async {
+        isReplaying = true
+        isReplayingLiveService = serviceIsRunning
+        defer {
+            applyReplayBuffer()
+            isReplaying = false
+            isReplayingLiveService = false
+        }
         status = .reconnecting
         let priorUUIDs = processedEventUUIDs.count
         logger.info("[Chat] Reconnecting to service logs (\(priorUUIDs) prior UUIDs)")
@@ -1033,6 +1068,9 @@ final class ChatViewModel {
                 modelContext: modelContext
             )
 
+            // Flush buffered content to the observable message in one shot
+            applyReplayBuffer()
+
             let currentUUIDs = processedEventUUIDs.count
             logger.info("[Chat] Reconnect stream ended: result=\(streamResult), content=\(assistantMessage.content.count), uuids=\(currentUUIDs)")
 
@@ -1079,9 +1117,10 @@ final class ChatViewModel {
 
     private func reconnectToServiceLogs(
         apiClient: SpritesAPIClient,
-        modelContext: ModelContext
+        modelContext: ModelContext,
+        serviceIsRunning: Bool = false
     ) async {
-        await runReconnectLoop(apiClient: apiClient, modelContext: modelContext)
+        await runReconnectLoop(apiClient: apiClient, modelContext: modelContext, serviceIsRunning: serviceIsRunning)
 
         if let queued = queuedPrompt, !Task.isCancelled {
             let prompt = buildPrompt(text: queued, attachments: queuedAttachments)
@@ -1110,28 +1149,44 @@ final class ChatViewModel {
             for block in assistantEvent.message.content {
                 switch block {
                 case .text(let text):
-                    if !hasPlayedFirstTextHaptic {
+                    if !isReplaying && !hasPlayedFirstTextHaptic {
                         hasPlayedFirstTextHaptic = true
                         fireHaptic(.medium)
                     }
                     // Merge consecutive text blocks
-                    if case .text(let existing) = message.content.last {
-                        message.content[message.content.count - 1] = .text(existing + text)
+                    if isReplaying {
+                        if case .text(let existing) = replayContentBuffer.last {
+                            replayContentBuffer[replayContentBuffer.count - 1] = .text(existing + text)
+                        } else {
+                            replayContentBuffer.append(.text(text))
+                        }
                     } else {
-                        message.content.append(.text(text))
+                        if case .text(let existing) = message.content.last {
+                            message.content[message.content.count - 1] = .text(existing + text)
+                        } else {
+                            message.content.append(.text(text))
+                        }
                     }
                 case .toolUse(let toolUse):
                     let card = ToolUseCard(
                         toolUseId: toolUse.id,
                         toolName: toolUse.name,
                         input: toolUse.input
                     )
-                    message.content.append(.toolUse(card))
                     logger.info("Tool use: \(toolUse.name, privacy: .public) id=\(toolUse.id, privacy: .public)")
-                    toolUseIndex[toolUse.id] = (
-                        messageIndex: messages.count - 1,
-                        toolName: toolUse.name
-                    )
+                    if isReplaying {
+                        replayContentBuffer.append(.toolUse(card))
+                        replayToolUseBuffer[toolUse.id] = (
+                            messageIndex: messages.count - 1,
+                            toolName: toolUse.name
+                        )
+                    } else {
+                        message.content.append(.toolUse(card))
+                        toolUseIndex[toolUse.id] = (
+                            messageIndex: messages.count - 1,
+                            toolName: toolUse.name
+                        )
+                    }
                     if ["Write", "Edit"].contains(toolUse.name) {
                         turnHasMutations = true
                     }
@@ -1144,22 +1199,34 @@ final class ChatViewModel {
             guard let message = currentAssistantMessage else { return }
 
             for result in toolResultEvent.message.content {
-                let toolName = toolUseIndex[result.toolUseId]?.toolName ?? "Unknown"
+                let toolName = replayToolUseBuffer[result.toolUseId]?.toolName
+                    ?? toolUseIndex[result.toolUseId]?.toolName
+                    ?? "Unknown"
                 let resultCard = ToolResultCard(
                     toolUseId: result.toolUseId,
                     toolName: toolName,
                     content: result.content ?? .null
                 )
-                message.content.append(.toolResult(resultCard))
-
-                // Link result back to matching tool use card
-                for item in message.content {
-                    if case .toolUse(let toolCard) = item, toolCard.toolUseId == result.toolUseId {
-                        toolCard.result = resultCard
-                        break
+                if isReplaying {
+                    replayContentBuffer.append(.toolResult(resultCard))
+                    // Link result to matching tool use card in the buffer
+                    for item in replayContentBuffer {
+                        if case .toolUse(let toolCard) = item, toolCard.toolUseId == result.toolUseId {
+                            toolCard.result = resultCard
+                            break
+                        }
+                    }
+                } else {
+                    message.content.append(.toolResult(resultCard))
+                    // Link result back to matching tool use card
+                    for item in message.content {
+                        if case .toolUse(let toolCard) = item, toolCard.toolUseId == result.toolUseId {
+                            toolCard.result = resultCard
+                            break
+                        }
                     }
+                    fireHaptic(.light)
                 }
-                fireHaptic(.light)
             }
 
         case .result(let resultEvent):
@@ -1168,10 +1235,10 @@ final class ChatViewModel {
             }
             receivedResultEvent = true
             sessionId = resultEvent.sessionId
-            saveSession(modelContext: modelContext)
+            saveSession(modelContext: modelContext, isComplete: true)
 
             let autoCheckpointEnabled = UserDefaults.standard.bool(forKey: "autoCheckpoint")
-            if turnHasMutations, autoCheckpointEnabled, let apiClient {
+            if !isReplaying, turnHasMutations, autoCheckpointEnabled, let apiClient {
                 let assistantMsg = currentAssistantMessage
                 let sprite = spriteName
                 Task { [weak assistantMsg] in
@@ -1429,11 +1496,12 @@ final class ChatViewModel {
         return try? modelContext.fetch(descriptor).first
     }
 
-    private func saveSession(modelContext: ModelContext) {
+    private func saveSession(modelContext: ModelContext, isComplete: Bool? = nil) {
         guard let chat = fetchChat(modelContext: modelContext) else { return }
         chat.claudeSessionId = sessionId
         chat.currentServiceName = serviceName
         chat.lastUsed = Date()
+        if let isComplete { chat.lastSessionComplete = isComplete }
         try? modelContext.save()
     }
 
@@ -1481,6 +1549,27 @@ final class ChatViewModel {
         return !timedOut
     }
 
+    /// Flush the replay content buffer into the current assistant message in one shot,
+    /// replacing N per-event observable mutations with a single array assignment.
+    private func applyReplayBuffer() {
+        guard !replayContentBuffer.isEmpty, let message = currentAssistantMessage else {
+            replayContentBuffer = []
+            replayToolUseBuffer = [:]
+            return
+        }
+        for item in replayContentBuffer {
+            if case .text(let newText) = item,
+               case .text(let existing) = message.content.last {
+                message.content[message.content.count - 1] = .text(existing + newText)
+            } else {
+                message.content.append(item)
+            }
+        }
+        toolUseIndex.merge(replayToolUseBuffer) { _, new in new }
+        replayContentBuffer = []
+        replayToolUseBuffer = [:]
+    }
+
     private func fireHaptic(_ style: UIImpactFeedbackGenerator.FeedbackStyle) {
         UIImpactFeedbackGenerator(style: style).impactOccurred()
     }