From 9c37ce5f9e0e54c44306813e2aa5554cd54835ea Mon Sep 17 00:00:00 2001
From: Eulices Lopez <105620565+eulicesl@users.noreply.github.com>
Date: Wed, 18 Mar 2026 09:08:33 -0400
Subject: [PATCH] fix(ios): harden chat stream acceptance and initial load

---
 .../OpenClawChatUI/ChatViewModel.swift        |  9 ++--
 .../OpenClawKitTests/ChatViewModelTests.swift | 44 +++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift b/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift
index 21ab6d6f0aa..33e0f553459 100644
--- a/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift
+++ b/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift
@@ -42,6 +42,7 @@ public final class OpenClawChatViewModel {
 
     @ObservationIgnored
     private nonisolated(unsafe) var eventTask: Task<Void, Never>?
+    private var hasLoadedInitialState = false
     private var pendingRuns = Set<String>() {
         didSet { self.pendingRunCount = self.pendingRuns.count }
     }
@@ -103,6 +104,8 @@ public final class OpenClawChatViewModel {
     }
 
     public func load() {
+        guard !self.hasLoadedInitialState else { return }
+        self.hasLoadedInitialState = true
         Task { await self.bootstrap() }
     }
 
@@ -905,12 +908,12 @@ public final class OpenClawChatViewModel {
     }
 
     private func shouldAcceptAgentEvent(_ evt: OpenClawAgentEventPayload) -> Bool {
-        if let sessionKey = evt.sessionKey {
-            return Self.matchesCurrentSessionKey(incoming: sessionKey, current: self.sessionKey)
-        }
         if self.pendingRuns.contains(evt.runId) {
             return true
         }
+        if let sessionKey = evt.sessionKey {
+            return Self.matchesCurrentSessionKey(incoming: sessionKey, current: self.sessionKey)
+        }
         if let sessionId {
             return evt.runId == sessionId
         }
diff --git a/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift b/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift
index d0e63151f5a..a123da09abe 100644
--- a/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift
+++ b/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift
@@ -338,6 +338,10 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor
         return ids.last
     }
 
+    func historyRequestCount() async -> Int {
+        await self.state.historyCallCount
+    }
+
     func abortedRunIds() async -> [String] {
         await self.state.abortedRunIds
     }
@@ -514,6 +518,46 @@ extension TestChatTransportState {
         #expect(await MainActor.run { vm.streamingAssistantText } == nil)
     }
 
+    @Test func acceptsPendingRunAgentEventsEvenWhenSessionKeyMismatches() async throws {
+        let history = historyPayload(sessionId: "sess-main")
+        let (transport, vm) = await makeViewModel(historyResponses: [history])
+        try await loadAndWaitBootstrap(vm: vm, sessionId: "sess-main")
+        await sendUserMessage(vm)
+        try await waitUntil("pending run starts") { await transport.lastSentRunId() != nil }
+        let runId = try #require(await transport.lastSentRunId())
+
+        emitAssistantText(
+            transport: transport,
+            runId: runId,
+            text: "own run still streams",
+            sessionKey: "agent:main:other")
+
+        try await waitUntil("assistant stream visible for own run despite key mismatch") {
+            await MainActor.run { vm.streamingAssistantText == "own run still streams" }
+        }
+    }
+
+    @Test func loadIsIdempotentDuringActiveStream() async throws {
+        let history = historyPayload(messages: [chatTextMessage(role: "user", text: "hello", timestamp: 1)])
+        let (transport, vm) = await makeViewModel(historyResponses: [history])
+        try await loadAndWaitBootstrap(vm: vm)
+        await sendUserMessage(vm)
+        try await waitUntil("pending run starts") { await transport.lastSentRunId() != nil }
+        let runId = try #require(await transport.lastSentRunId())
+
+        emitAssistantText(transport: transport, runId: runId, text: "stream survives load")
+        try await waitUntil("assistant stream visible") {
+            await MainActor.run { vm.streamingAssistantText == "stream survives load" }
+        }
+
+        await MainActor.run { vm.load() }
+        try? await Task.sleep(nanoseconds: 150_000_000)
+
+        #expect(await transport.historyRequestCount() == 1)
+        #expect(await MainActor.run { vm.streamingAssistantText } == "stream survives load")
+        #expect(await MainActor.run { vm.pendingRunCount } == 1)
+    }
+
     @Test func acceptsCanonicalSessionKeyEventsForExternalRuns() async throws {
         let now = Date().timeIntervalSince1970 * 1000
         let history1 = historyPayload(messages: [chatTextMessage(role: "user", text: "first", timestamp: now)])