From 7478e6e485a78cbee2f8ea0e7701f87b55c311fd Mon Sep 17 00:00:00 2001 From: joshavant <830519+joshavant@users.noreply.github.com> Date: Fri, 5 Jun 2026 20:56:51 -0500 Subject: [PATCH] Fix chat session sync ownership --- apps/ios/Sources/Design/ChatProTab.swift | 28 +-- .../OpenClawChatUI/ChatViewModel.swift | 60 ++++++- .../OpenClawKitTests/ChatViewModelTests.swift | 168 +++++++++++++++++- 3 files changed, 220 insertions(+), 36 deletions(-) diff --git a/apps/ios/Sources/Design/ChatProTab.swift b/apps/ios/Sources/Design/ChatProTab.swift index c606ee63e5df..8104f213f39d 100644 --- a/apps/ios/Sources/Design/ChatProTab.swift +++ b/apps/ios/Sources/Design/ChatProTab.swift @@ -6,7 +6,6 @@ struct ChatProTab: View { @Environment(NodeAppModel.self) private var appModel @Environment(\.colorScheme) private var colorScheme @State private var viewModel: OpenClawChatViewModel? - @State private var programmaticSessionSwitchCounts: [String: Int] = [:] var body: some View { NavigationStack { @@ -107,14 +106,6 @@ struct ChatProTab: View { sessionKey: sessionKey, transport: IOSGatewayChatTransport(gateway: self.appModel.operatorSession), onSessionChanged: { sessionKey in - if self.consumeProgrammaticSessionSwitch(sessionKey) { - // Programmatic switches complete asynchronously; stale completions - // must repair back to the current app-model session, not become focus. - if sessionKey != self.appModel.chatSessionKey { - self.syncChatViewModel() - } - return - } self.appModel.focusChatSession(sessionKey) }, diagnosticsLog: { message in @@ -123,24 +114,7 @@ struct ChatProTab: View { return } guard viewModel.sessionKey != sessionKey else { return } - self.recordProgrammaticSessionSwitch(sessionKey) - viewModel.switchSession(to: sessionKey) - } - - private func recordProgrammaticSessionSwitch(_ sessionKey: String) { - self.programmaticSessionSwitchCounts[sessionKey, default: 0] += 1 - } - - private func consumeProgrammaticSessionSwitch(_ sessionKey: String) -> Bool { - guard let count = self.programmaticSessionSwitchCounts[sessionKey] else { - return false - } - if count <= 1 { - self.programmaticSessionSwitchCounts[sessionKey] = nil - } else { - self.programmaticSessionSwitchCounts[sessionKey] = count - 1 - } - return true + viewModel.syncSession(to: sessionKey) } private var talkControl: OpenClawChatTalkControl { diff --git a/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift b/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift index 9fcd6722f2bd..365d8c60b28c 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift @@ -72,6 +72,11 @@ public final class OpenClawChatViewModel { private var lastCompactAt: Date? private let compactCooldown: TimeInterval = 60 + private enum SessionSwitchIntent { + case userInitiated + case externalSync + } + private var pendingToolCallsById: [String: OpenClawChatPendingToolCall] = [:] { didSet { self.pendingToolCalls = self.pendingToolCallsById.values @@ -151,7 +156,11 @@ public final class OpenClawChatViewModel { } public func switchSession(to sessionKey: String) { - Task { await self.performSwitchSession(to: sessionKey) } + self.applySessionSwitch(to: sessionKey, intent: .userInitiated) + } + + public func syncSession(to sessionKey: String) { + self.applySessionSwitch(to: sessionKey, intent: .externalSync) } public func selectThinkingLevel(_ level: String) { @@ -261,7 +270,9 @@ public final class OpenClawChatViewModel { self.diagnosticsLog?(message) } - private func bootstrap() async { + private func bootstrap(sessionKey requestedSessionKey: String? = nil) async { + let sessionKey = requestedSessionKey ?? self.sessionKey + guard sessionKey == self.sessionKey else { return } self.isLoading = true self.errorText = nil self.healthOK = false @@ -269,15 +280,24 @@ public final class OpenClawChatViewModel { self.pendingToolCallsById = [:] self.streamingAssistantText = nil self.sessionId = nil - defer { self.isLoading = false } + defer { + if self.sessionKey == sessionKey { + self.isLoading = false + } + } do { do { - try await self.transport.setActiveSessionKey(self.sessionKey) + try await self.transport.setActiveSessionKey(sessionKey) } catch { // Best-effort only; history/send/health still work without push events. } + guard self.sessionKey == sessionKey else { + await self.restoreActiveSessionAfterStaleBootstrap(staleSessionKey: sessionKey) + return + } - let payload = try await self.transport.requestHistory(sessionKey: self.sessionKey) + let payload = try await self.transport.requestHistory(sessionKey: sessionKey) + guard self.sessionKey == sessionKey else { return } self.messages = Self.reconcileMessageIDs( previous: self.messages, incoming: Self.decodeMessages(payload.messages ?? [])) @@ -290,15 +310,37 @@ public final class OpenClawChatViewModel { } self.syncThinkingLevelOptions() await self.pollHealthIfNeeded(force: true) + guard self.sessionKey == sessionKey else { return } await self.fetchSessions(limit: 50) + guard self.sessionKey == sessionKey else { return } await self.fetchModels() + guard self.sessionKey == sessionKey else { return } self.errorText = nil } catch { + guard self.sessionKey == sessionKey else { return } self.errorText = error.localizedDescription chatUILogger.error("bootstrap failed \(error.localizedDescription, privacy: .public)") } } + private func restoreActiveSessionAfterStaleBootstrap(staleSessionKey: String) async { + var lastSubscribedSessionKey = staleSessionKey + while true { + let currentSessionKey = self.sessionKey + guard currentSessionKey != lastSubscribedSessionKey else { return } + do { + // A stale bootstrap may complete its subscribe side effect after the winning switch. + // Reassert and recheck so push events stay aligned with the visible session. + try await self.transport.setActiveSessionKey(currentSessionKey) + } catch { + // Best-effort only; the current bootstrap still owns history/send/health. + return + } + guard self.sessionKey != currentSessionKey else { return } + lastSubscribedSessionKey = currentSessionKey + } + } + private func refreshPendingRunAfterForeground() async { guard self.pendingRunCount > 0 else { return } self.logDiagnostic( @@ -757,14 +799,16 @@ public final class OpenClawChatViewModel { } } - private func performSwitchSession(to sessionKey: String) async { + private func applySessionSwitch(to sessionKey: String, intent: SessionSwitchIntent) { let next = sessionKey.trimmingCharacters(in: .whitespacesAndNewlines) guard !next.isEmpty else { return } guard next != self.sessionKey else { return } self.sessionKey = next - self.onSessionChanged?(next) + if intent == .userInitiated { + self.onSessionChanged?(next) + } self.modelSelectionID = Self.defaultModelSelectionID - await self.bootstrap() + Task { await self.bootstrap(sessionKey: next) } } private func performStartNewSession() async { diff --git a/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift b/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift index 762144123fd8..072adf389a85 100644 --- a/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift +++ b/apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift @@ -97,6 +97,7 @@ private func makeViewModel( historyResponses: [OpenClawChatHistoryPayload], sessionsResponses: [OpenClawChatSessionsListResponse] = [], modelResponses: [[OpenClawChatModelChoice]] = [], + setActiveSessionHook: (@Sendable (String) async throws -> Void)? = nil, createSessionHook: (@Sendable (String, String?) async throws -> Void)? = nil, resetSessionHook: (@Sendable (String) async throws -> Void)? = nil, compactSessionHook: (@Sendable (String) async throws -> Void)? = nil, @@ -105,6 +106,7 @@ private func makeViewModel( waitForRunCompletionHook: (@Sendable (String, Int) async -> Bool)? = nil, healthResponses: [Bool] = [true], initialThinkingLevel: String? = nil, + onSessionChanged: (@MainActor (String) -> Void)? = nil, onThinkingLevelChanged: (@MainActor @Sendable (String) -> Void)? = nil) async -> (TestChatTransport, OpenClawChatViewModel) { @@ -112,6 +114,7 @@ private func makeViewModel( historyResponses: historyResponses, sessionsResponses: sessionsResponses, modelResponses: modelResponses, + setActiveSessionHook: setActiveSessionHook, createSessionHook: createSessionHook, resetSessionHook: resetSessionHook, compactSessionHook: compactSessionHook, @@ -124,6 +127,7 @@ private func makeViewModel( sessionKey: sessionKey, transport: transport, initialThinkingLevel: initialThinkingLevel, + onSessionChanged: onSessionChanged, onThinkingLevelChanged: onThinkingLevelChanged) } return (transport, vm) @@ -276,11 +280,30 @@ private actor AsyncCounter { } } +private actor SessionSubscribeGate { + private var waiters: [CheckedContinuation] = [] + + func wait() async { + await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + func release() { + let waiters = self.waiters + self.waiters = [] + for waiter in waiters { + waiter.resume() + } + } +} + private actor TestChatTransportState { var historyCallCount: Int = 0 var sessionsCallCount: Int = 0 var modelsCallCount: Int = 0 var healthCallCount: Int = 0 + var activeSessionKeys: [String] = [] var createdSessionKeys: [String] = [] var createdParentSessionKeys: [String?] = [] var resetSessionKeys: [String] = [] @@ -299,6 +322,7 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor private let historyResponses: [OpenClawChatHistoryPayload] private let sessionsResponses: [OpenClawChatSessionsListResponse] private let modelResponses: [[OpenClawChatModelChoice]] + private let setActiveSessionHook: (@Sendable (String) async throws -> Void)? private let createSessionHook: (@Sendable (String, String?) async throws -> Void)? private let resetSessionHook: (@Sendable (String) async throws -> Void)? private let compactSessionHook: (@Sendable (String) async throws -> Void)? @@ -314,6 +338,7 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor historyResponses: [OpenClawChatHistoryPayload], sessionsResponses: [OpenClawChatSessionsListResponse] = [], modelResponses: [[OpenClawChatModelChoice]] = [], + setActiveSessionHook: (@Sendable (String) async throws -> Void)? = nil, createSessionHook: (@Sendable (String, String?) async throws -> Void)? = nil, resetSessionHook: (@Sendable (String) async throws -> Void)? = nil, compactSessionHook: (@Sendable (String) async throws -> Void)? = nil, @@ -325,6 +350,7 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor self.historyResponses = historyResponses self.sessionsResponses = sessionsResponses self.modelResponses = modelResponses + self.setActiveSessionHook = setActiveSessionHook self.createSessionHook = createSessionHook self.resetSessionHook = resetSessionHook self.compactSessionHook = compactSessionHook @@ -343,7 +369,12 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor self.stream } - func setActiveSessionKey(_: String) async throws {} + func setActiveSessionKey(_ sessionKey: String) async throws { + await self.state.activeSessionKeysAppend(sessionKey) + if let setActiveSessionHook { + try await setActiveSessionHook(sessionKey) + } + } func createSession( key: String, @@ -483,6 +514,10 @@ private final class TestChatTransport: @unchecked Sendable, OpenClawChatTranspor await self.state.patchedModels } + func activeSessionKeys() async -> [String] { + await self.state.activeSessionKeys + } + func patchedThinkingLevels() async -> [String] { await self.state.patchedThinkingLevels } @@ -525,6 +560,10 @@ extension TestChatTransportState { self.healthCallCount = v } + fileprivate func activeSessionKeysAppend(_ v: String) { + self.activeSessionKeys.append(v) + } + fileprivate func sentRunIdsAppend(_ v: String) { self.sentRunIds.append(v) } @@ -2162,6 +2201,133 @@ extension TestChatTransportState { #expect(await transport.patchedModels() == ["openai/gpt-5.4", "openai/gpt-5.4-pro"]) } + @Test @MainActor func switchSessionNotifiesSessionChangedCallback() async throws { + var changedSessionKeys: [String] = [] + let (_, vm) = await makeViewModel( + historyResponses: [ + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "other", sessionId: "sess-other"), + ], + onSessionChanged: { changedSessionKeys.append($0) }) + + try await loadAndWaitBootstrap(vm: vm, sessionId: "sess-main") + + vm.switchSession(to: "other") + + try await waitUntil("user switch bootstrapped target session") { + await MainActor.run { vm.sessionKey == "other" && vm.sessionId == "sess-other" } + } + #expect(changedSessionKeys == ["other"]) + } + + @Test @MainActor func syncSessionDoesNotNotifySessionChangedCallback() async throws { + var changedSessionKeys: [String] = [] + let (_, vm) = await makeViewModel( + historyResponses: [ + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "other", sessionId: "sess-other"), + ], + onSessionChanged: { changedSessionKeys.append($0) }) + + try await loadAndWaitBootstrap(vm: vm, sessionId: "sess-main") + + vm.syncSession(to: "other") + + try await waitUntil("external sync bootstrapped target session") { + await MainActor.run { vm.sessionKey == "other" && vm.sessionId == "sess-other" } + } + #expect(changedSessionKeys.isEmpty) + } + + @Test @MainActor func staleSyncBootstrapRestoresCurrentActiveSessionSubscription() async throws { + let staleSubscribeGate = SessionSubscribeGate() + let (transport, vm) = await makeViewModel( + historyResponses: [ + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "other", sessionId: "sess-other"), + ], + setActiveSessionHook: { sessionKey in + if sessionKey == "other" { + await staleSubscribeGate.wait() + } + }) + + try await loadAndWaitBootstrap(vm: vm, sessionId: "sess-main") + + vm.syncSession(to: "other") + try await waitUntil("stale subscribe is in flight") { + await transport.activeSessionKeys().last == "other" + } + + vm.syncSession(to: "main") + try await waitUntil("current session subscribed") { + let sessionKey = await MainActor.run { vm.sessionKey } + let activeSessionKeys = await transport.activeSessionKeys() + return sessionKey == "main" && + Array(activeSessionKeys.suffix(2)) == ["other", "main"] + } + + await staleSubscribeGate.release() + + try await waitUntil("current session resubscribed after stale subscribe") { + Array(await transport.activeSessionKeys().suffix(3)) == ["other", "main", "main"] + } + } + + @Test @MainActor func staleSyncRepairReassertsLatestActiveSessionSubscription() async throws { + let staleSubscribeGate = SessionSubscribeGate() + let staleRepairGate = SessionSubscribeGate() + let mainSubscribeCount = AsyncCounter() + let (transport, vm) = await makeViewModel( + historyResponses: [ + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "main", sessionId: "sess-main"), + historyPayload(sessionKey: "final", sessionId: "sess-final"), + ], + setActiveSessionHook: { sessionKey in + if sessionKey == "other" { + await staleSubscribeGate.wait() + } + if sessionKey == "main" { + let count = await mainSubscribeCount.increment() + if count == 3 { + await staleRepairGate.wait() + } + } + }) + + try await loadAndWaitBootstrap(vm: vm, sessionId: "sess-main") + + vm.syncSession(to: "other") + try await waitUntil("stale subscribe is in flight") { + await transport.activeSessionKeys().last == "other" + } + + vm.syncSession(to: "main") + try await waitUntil("main session subscribed") { + Array(await transport.activeSessionKeys().suffix(2)) == ["other", "main"] + } + + await staleSubscribeGate.release() + try await waitUntil("stale repair is in flight") { + Array(await transport.activeSessionKeys().suffix(3)) == ["other", "main", "main"] + } + + vm.syncSession(to: "final") + try await waitUntil("newest session subscribed") { + let sessionKey = await MainActor.run { vm.sessionKey } + let activeSessionKeys = await transport.activeSessionKeys() + return sessionKey == "final" && activeSessionKeys.last == "final" + } + + await staleRepairGate.release() + + try await waitUntil("newest session resubscribed after stale repair") { + Array(await transport.activeSessionKeys().suffix(3)) == ["main", "final", "final"] + } + } + @Test func switchingSessionsIgnoresLateModelPatchCompletionFromPreviousSession() async throws { let now = Date().timeIntervalSince1970 * 1000 let sessions = OpenClawChatSessionsListResponse(