Merge pull request #60 from mcintyre94/investigate-reconnect-on-disconnect

Fix reconnect loop dropping final messages when service stops mid-reconnect

Author: mcintyre94

Wisp/Services/SpritesAPIClient.swift

@@ -280,6 +280,26 @@ final class SpritesAPIClient {
         return try await request(method: "GET", path: "/sprites/\(spriteName)/services/\(serviceName)")
     }
 
+    // ServiceLogsProvider conformance — bridges the default-argument version to the protocol signature.
+    func streamServiceLogs(spriteName: String, serviceName: String) -> AsyncThrowingStream<ServiceLogEvent, Error> {
+        streamServiceLogs(spriteName: spriteName, serviceName: serviceName, duration: "3600s")
+    }
+}
+
+// MARK: - ServiceLogsProvider
+
+/// Minimal protocol covering the two API calls used by the reconnect loop,
+/// allowing the loop to be tested without a live network connection.
+@MainActor
+protocol ServiceLogsProvider {
+    func streamServiceLogs(spriteName: String, serviceName: String) -> AsyncThrowingStream<ServiceLogEvent, Error>
+    func getServiceStatus(spriteName: String, serviceName: String) async throws -> ServiceInfo
+}
+
+extension SpritesAPIClient: ServiceLogsProvider {}
+
+extension SpritesAPIClient {
+
     /// Delete a service (5s timeout to avoid blocking callers if sprite is unresponsive).
     func deleteService(spriteName: String, serviceName: String) async throws {
         let _: EmptyResponse = try await request(method: "DELETE", path: "/sprites/\(spriteName)/services/\(serviceName)", timeout: 5)

Wisp/ViewModels/ChatViewModel.swift

@@ -956,8 +956,11 @@ final class ChatViewModel {
     /// so existing content stays on screen with no flash. Only genuinely new events
     /// are appended. If the service is still running after a replay, polls and
     /// re-replays until the service stops or a result event arrives.
-    private func reconnectToServiceLogs(
-        apiClient: SpritesAPIClient,
+    /// Core reconnect loop — fetches full log history on repeat until a result event
+    /// arrives or the service is confirmed stopped. Separated from
+    /// `reconnectToServiceLogs` so it can be tested against a mock API client.
+    func runReconnectLoop(
+        apiClient: some ServiceLogsProvider,
         modelContext: ModelContext
     ) async {
         status = .reconnecting
@@ -986,6 +989,7 @@ final class ChatViewModel {
         // Replay loop — each iteration fetches full log history.
         // processServiceStream skips events whose UUID is already in
         // processedEventUUIDs, so content is never cleared mid-stream.
+        var retriedAfterServiceStopped = false
         while !Task.isCancelled {
             receivedSystemEvent = false
             receivedResultEvent = false
@@ -1017,15 +1021,27 @@ final class ChatViewModel {
             // If we got a result event, Claude is done
             if receivedResultEvent { break }
 
-            // Check if service is still running before retrying
-            if let serviceInfo = try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName),
-               serviceInfo.state.status == "running" {
+            // Check if service is still running
+            let isRunning = (try? await apiClient.getServiceStatus(spriteName: spriteName, serviceName: serviceName))?.state.status == "running"
+
+            if isRunning {
                 logger.info("[Chat] Service still running, will re-poll after delay")
                 try? await Task.sleep(for: .seconds(2))
                 continue
             }
 
-            // Service not running or status check failed — we're done
+            // Service has stopped (or status check failed / service gone). The GET stream
+            // may have been killed by iOS just as Claude finished writing its final events —
+            // a race between the connection dying and the result arriving. Allow one extra
+            // retry so we catch any events that landed in the log after the stream closed.
+            if !retriedAfterServiceStopped {
+                retriedAfterServiceStopped = true
+                logger.info("[Chat] Service stopped without result event — retrying once for final events")
+                try? await Task.sleep(for: .seconds(1))
+                continue
+            }
+
+            // Already retried after stop — give up
             break
         }
 
@@ -1039,6 +1055,13 @@ final class ChatViewModel {
             status = .idle
         }
         persistMessages(modelContext: modelContext)
+    }
+
+    private func reconnectToServiceLogs(
+        apiClient: SpritesAPIClient,
+        modelContext: ModelContext
+    ) async {
+        await runReconnectLoop(apiClient: apiClient, modelContext: modelContext)
 
         if let queued = queuedPrompt, !Task.isCancelled {
             let prompt = buildPrompt(text: queued, attachments: queuedAttachments)

WispTests/ChatViewModelTests.swift

@@ -26,6 +26,33 @@ struct ChatViewModelTests {
         return (vm, chat)
     }
 
+    // MARK: - Mock API client
+
+    private final class MockServiceLogsProvider: ServiceLogsProvider {
+        var streams: [AsyncThrowingStream<ServiceLogEvent, Error>]
+        var statuses: [String]
+        private(set) var streamCallCount = 0
+        private(set) var statusCallCount = 0
+
+        init(streams: [AsyncThrowingStream<ServiceLogEvent, Error>], statuses: [String]) {
+            self.streams = streams
+            self.statuses = statuses
+        }
+
+        func streamServiceLogs(spriteName: String, serviceName: String) -> AsyncThrowingStream<ServiceLogEvent, Error> {
+            let idx = streamCallCount
+            streamCallCount += 1
+            return idx < streams.count ? streams[idx] : AsyncThrowingStream { $0.finish() }
+        }
+
+        func getServiceStatus(spriteName: String, serviceName: String) async throws -> ServiceInfo {
+            let idx = statusCallCount
+            statusCallCount += 1
+            let status = idx < statuses.count ? statuses[idx] : "stopped"
+            return ServiceInfo(name: serviceName, state: ServiceInfo.ServiceState(status: status))
+        }
+    }
+
     // MARK: - handleEvent: system
 
     @Test func handleEvent_systemSetsModelName() throws {
@@ -723,6 +750,65 @@ struct ChatViewModelTests {
         #expect(vm.inputText == "")
     }
 
+    // MARK: - reconnectToServiceLogs: retriedAfterServiceStopped
+
+    @Test func reconnectToServiceLogs_retriesOnceWhenServiceStoppedWithNoResult_thenDeliversResult() async throws {
+        let ctx = try makeModelContext()
+        let (vm, _) = makeChatViewModel(modelContext: ctx)
+
+        // First stream: delivers a system event but no result — simulates the
+        // stream dying just before Claude finishes.
+        let systemLine = #"{"type":"system","session_id":"s1","model":"claude-sonnet-4-20250514"}"# + "\n"
+        let stream1 = AsyncThrowingStream<ServiceLogEvent, Error> { continuation in
+            continuation.yield(ServiceLogEvent(type: .stdout, data: systemLine, exitCode: nil, timestamp: nil, logFiles: nil))
+            continuation.finish()
+        }
+
+        // Second stream: delivers the result event that landed after the first stream closed.
+        let resultLine = #"{"type":"result","session_id":"s1","subtype":"success"}"# + "\n"
+        let stream2 = AsyncThrowingStream<ServiceLogEvent, Error> { continuation in
+            continuation.yield(ServiceLogEvent(type: .stdout, data: systemLine, exitCode: nil, timestamp: nil, logFiles: nil))
+            continuation.yield(ServiceLogEvent(type: .stdout, data: resultLine, exitCode: nil, timestamp: nil, logFiles: nil))
+            continuation.finish()
+        }
+
+        let mock = MockServiceLogsProvider(streams: [stream1, stream2], statuses: ["stopped"])
+
+        await vm.runReconnectLoop(apiClient: mock, modelContext: ctx)
+
+        #expect(mock.streamCallCount == 2, "Should replay logs twice: once on initial reconnect, once on retry")
+        #expect(mock.statusCallCount == 1, "Should only check status once (before the retry)")
+        guard case .idle = vm.status else {
+            Issue.record("Expected idle status after reconnect completes, got \(vm.status)")
+            return
+        }
+    }
+
+    @Test func reconnectToServiceLogs_givesUpAfterOneRetryWhenServiceStillStopped() async throws {
+        let ctx = try makeModelContext()
+        let (vm, _) = makeChatViewModel(modelContext: ctx)
+
+        // Both streams return no result event, and the service stays stopped.
+        let systemLine = #"{"type":"system","session_id":"s1","model":"claude-sonnet-4-20250514"}"# + "\n"
+        let makeStream = {
+            AsyncThrowingStream<ServiceLogEvent, Error> { continuation in
+                continuation.yield(ServiceLogEvent(type: .stdout, data: systemLine, exitCode: nil, timestamp: nil, logFiles: nil))
+                continuation.finish()
+            }
+        }
+
+        let mock = MockServiceLogsProvider(streams: [makeStream(), makeStream()], statuses: ["stopped", "stopped"])
+
+        await vm.runReconnectLoop(apiClient: mock, modelContext: ctx)
+
+        #expect(mock.streamCallCount == 2, "Should attempt exactly two replays: initial + one retry")
+        #expect(mock.statusCallCount == 2, "Should check status once per iteration that yields no result event")
+        guard case .idle = vm.status else {
+            Issue.record("Expected idle status after giving up, got \(vm.status)")
+            return
+        }
+    }
+
     @Test func stashDraft_leavesInputReadyForNextMessage() throws {
         let ctx = try makeModelContext()
         let (vm, _) = makeChatViewModel(modelContext: ctx)