Skip to content

Commit 71be307

Browse files
joshavantjoshavant
committed
fix(ios): scope final message reconciliation
1 parent 1ee2212 commit 71be307

2 files changed

Lines changed: 160 additions & 36 deletions

File tree

‎apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatViewModel.swift‎

Lines changed: 152 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,11 @@ public final class OpenClawChatViewModel {
4747
}
4848

4949
private var pendingLocalUserEchoMessageIDsByRunID: [String: UUID] = [:]
50-
private var provisionalFinalMessageIDs = Set<UUID>()
50+
// Final chat events and durable session-message rows arrive independently.
51+
// Keep each provisional final scoped to the run's user turn so a later identical
52+
// answer in the same session does not adopt or suppress the wrong row.
53+
private var runMessageScopesByRunID: [String: RunMessageScope] = [:]
54+
private var provisionalFinalMessagesByID: [UUID: ProvisionalFinalMessage] = [:]
5155
private var sessionGeneration: UInt64 = 0
5256
private var bootstrapGeneration: UInt64 = 0
5357
// A newer same-session history request only invalidates older responses after it applies.
@@ -112,6 +116,17 @@ public final class OpenClawChatViewModel {
112116
var timestamp: Double?
113117
}
114118

119+
private struct RunMessageScope {
120+
var session: SessionSnapshot
121+
var latestUserTurn: LatestUserTurn?
122+
}
123+
124+
private struct ProvisionalFinalMessage {
125+
var reconciliationKey: String
126+
var runId: String?
127+
var scope: RunMessageScope
128+
}
129+
115130
private var pendingToolCallsById: [String: OpenClawChatPendingToolCall] = [:] {
116131
didSet {
117132
self.pendingToolCalls = self.pendingToolCallsById.values
@@ -360,7 +375,8 @@ public final class OpenClawChatViewModel {
360375
Self.reconcileMessageIDs(previous: self.messages, incoming: incoming)
361376
}
362377
self.prunePendingLocalUserEchoMessageIDs()
363-
self.pruneProvisionalFinalMessageIDs()
378+
self.pruneProvisionalFinalMessages()
379+
self.pruneRunMessageScopes()
364380
self.sessionId = payload.sessionId
365381
// Incomplete refreshes can arrive before durable assistant history.
366382
// The latest visible user turn must survive answered before it can reject older replies.
@@ -580,14 +596,84 @@ public final class OpenClawChatViewModel {
580596
return [role, toolCallId, toolName, contentFingerprint].joined(separator: "|")
581597
}
582598

583-
private func hasCanonicalFinalMessageMatching(_ message: OpenClawChatMessage) -> Bool {
599+
private static func normalizedRunID(_ runId: String?) -> String? {
600+
let trimmed = runId?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
601+
return trimmed.isEmpty ? nil : trimmed
602+
}
603+
604+
private func currentRunMessageScope() -> RunMessageScope {
605+
RunMessageScope(
606+
session: self.currentSessionSnapshot(),
607+
latestUserTurn: Self.latestUserTurn(in: self.messages))
608+
}
609+
610+
private func runMessageScope(for runId: String?) -> RunMessageScope {
611+
guard let runId = Self.normalizedRunID(runId),
612+
let scope = self.runMessageScopesByRunID[runId],
613+
self.isCurrentSession(scope.session)
614+
else {
615+
return self.currentRunMessageScope()
616+
}
617+
return scope
618+
}
619+
620+
private static func isSameUserTurnBoundary(_ lhs: LatestUserTurn?, _ rhs: LatestUserTurn?) -> Bool {
621+
switch (lhs, rhs) {
622+
case (nil, nil):
623+
return true
624+
case let (lhs?, rhs?):
625+
if let lhsKey = lhs.refreshKey, let rhsKey = rhs.refreshKey {
626+
return lhsKey == rhsKey && lhs.occurrence == rhs.occurrence
627+
}
628+
return lhs.refreshKey == nil &&
629+
rhs.refreshKey == nil &&
630+
lhs.occurrence == rhs.occurrence &&
631+
lhs.timestamp == rhs.timestamp
632+
default:
633+
return false
634+
}
635+
}
636+
637+
private static func indexAfterLatestUserTurn(
638+
_ latestUserTurn: LatestUserTurn?,
639+
in messages: [OpenClawChatMessage])
640+
-> [OpenClawChatMessage].Index
641+
{
642+
guard let latestUserTurn else { return messages.startIndex }
643+
if let refreshKey = latestUserTurn.refreshKey {
644+
var occurrence = 0
645+
for index in messages.indices {
646+
guard self.userRefreshIdentityKey(for: messages[index]) == refreshKey else { continue }
647+
occurrence += 1
648+
if occurrence == latestUserTurn.occurrence {
649+
return messages.index(after: index)
650+
}
651+
}
652+
} else if let timestamp = latestUserTurn.timestamp,
653+
let index = messages.lastIndex(where: { message in
654+
message.role.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "user" &&
655+
message.timestamp == timestamp
656+
})
657+
{
658+
return messages.index(after: index)
659+
}
660+
661+
return messages.lastIndex(where: { message in
662+
message.role.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "user"
663+
}).map { messages.index(after: $0) } ?? messages.startIndex
664+
}
665+
666+
private func hasCanonicalFinalMessageMatching(
667+
_ message: OpenClawChatMessage,
668+
scope: RunMessageScope) -> Bool
669+
{
584670
guard let key = Self.finalMessageReconciliationKey(for: message) else { return false }
585-
let searchStart = self.messages.lastIndex(where: { existing in
586-
existing.role.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "user"
587-
}).map { self.messages.index(after: $0) } ?? self.messages.startIndex
671+
guard self.isCurrentSession(scope.session) else { return false }
672+
let searchStart = Self.indexAfterLatestUserTurn(scope.latestUserTurn, in: self.messages)
673+
guard searchStart < self.messages.endIndex else { return false }
588674

589675
return self.messages[searchStart...].contains { existing in
590-
!self.provisionalFinalMessageIDs.contains(existing.id) &&
676+
self.provisionalFinalMessagesByID[existing.id] == nil &&
591677
Self.finalMessageReconciliationKey(for: existing) == key
592678
}
593679
}
@@ -600,10 +686,24 @@ public final class OpenClawChatViewModel {
600686
}
601687
}
602688

603-
private func pruneProvisionalFinalMessageIDs() {
604-
guard !self.provisionalFinalMessageIDs.isEmpty else { return }
689+
private func pruneProvisionalFinalMessages() {
690+
guard !self.provisionalFinalMessagesByID.isEmpty else { return }
605691
let visibleMessageIDs = Set(messages.map(\.id))
606-
self.provisionalFinalMessageIDs = self.provisionalFinalMessageIDs.intersection(visibleMessageIDs)
692+
self.provisionalFinalMessagesByID = self.provisionalFinalMessagesByID.filter { entry in
693+
visibleMessageIDs.contains(entry.key) && self.isCurrentSession(entry.value.scope.session)
694+
}
695+
}
696+
697+
private func pruneRunMessageScopes() {
698+
self.runMessageScopesByRunID = self.runMessageScopesByRunID.filter { entry in
699+
self.isCurrentSession(entry.value.session)
700+
}
701+
guard self.runMessageScopesByRunID.count > 64 else { return }
702+
let referencedRunIDs = Set(self.pendingRuns)
703+
.union(self.provisionalFinalMessagesByID.values.compactMap(\.runId))
704+
self.runMessageScopesByRunID = self.runMessageScopesByRunID.filter { entry in
705+
referencedRunIDs.contains(entry.key)
706+
}
607707
}
608708

609709
private func adoptPendingLocalUserEcho(incoming: OpenClawChatMessage) -> Bool {
@@ -637,18 +737,20 @@ public final class OpenClawChatViewModel {
637737

638738
private func adoptProvisionalFinalMessage(incoming: OpenClawChatMessage) -> Bool {
639739
guard let incomingKey = Self.finalMessageReconciliationKey(for: incoming) else { return false }
640-
let searchStart = self.messages.lastIndex(where: { existing in
641-
existing.role.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "user"
642-
}).map { self.messages.index(after: $0) } ?? self.messages.startIndex
740+
let canonicalScope = self.currentRunMessageScope()
741+
let searchStart = Self.indexAfterLatestUserTurn(canonicalScope.latestUserTurn, in: self.messages)
742+
guard searchStart < self.messages.endIndex else { return false }
643743

644744
guard let matchIndex = messages[searchStart...].lastIndex(where: { existing in
645-
self.provisionalFinalMessageIDs.contains(existing.id) &&
646-
Self.finalMessageReconciliationKey(for: existing) == incomingKey
745+
guard let provisional = self.provisionalFinalMessagesByID[existing.id] else { return false }
746+
return provisional.reconciliationKey == incomingKey &&
747+
Self.isSameUserTurnBoundary(provisional.scope.latestUserTurn, canonicalScope.latestUserTurn)
647748
}) else {
648749
return false
649750
}
650751

651752
let existing = self.messages[matchIndex]
753+
let provisional = self.provisionalFinalMessagesByID[existing.id]
652754
var updated = self.messages
653755
updated[matchIndex] = OpenClawChatMessage(
654756
id: existing.id,
@@ -660,9 +762,13 @@ public final class OpenClawChatViewModel {
660762
usage: incoming.usage,
661763
stopReason: incoming.stopReason,
662764
errorMessage: incoming.errorMessage)
663-
self.provisionalFinalMessageIDs.remove(existing.id)
765+
self.provisionalFinalMessagesByID.removeValue(forKey: existing.id)
766+
if let runId = provisional?.runId {
767+
self.runMessageScopesByRunID.removeValue(forKey: runId)
768+
}
664769
self.messages = Self.dedupeMessages(updated)
665-
self.pruneProvisionalFinalMessageIDs()
770+
self.pruneProvisionalFinalMessages()
771+
self.pruneRunMessageScopes()
666772
return true
667773
}
668774

@@ -912,6 +1018,7 @@ public final class OpenClawChatViewModel {
9121018
content: userContent,
9131019
timestamp: userMessageTimestamp))
9141020
self.pendingLocalUserEchoMessageIDsByRunID[runId] = userMessageID
1021+
self.runMessageScopesByRunID[runId] = self.currentRunMessageScope()
9151022

9161023
// Clear input immediately for responsive UX (before network await)
9171024
self.input = ""
@@ -935,9 +1042,11 @@ public final class OpenClawChatViewModel {
9351042
+ "localRunId=\(runId) remoteRunId=\(response.runId)")
9361043
if response.runId != runId {
9371044
let pendingUserMessageID = self.pendingLocalUserEchoMessageIDsByRunID.removeValue(forKey: runId)
1045+
let runScope = self.runMessageScopesByRunID.removeValue(forKey: runId)
9381046
self.clearPendingRun(runId)
9391047
self.pendingRuns.insert(response.runId)
9401048
self.pendingLocalUserEchoMessageIDsByRunID[response.runId] = pendingUserMessageID
1049+
self.runMessageScopesByRunID[response.runId] = runScope
9411050
self.armPendingRunTimeout(runId: response.runId)
9421051
}
9431052
if response.status == "ok" {
@@ -966,6 +1075,7 @@ public final class OpenClawChatViewModel {
9661075
} catch {
9671076
guard self.isCurrentSession(sessionSnapshot) else { return }
9681077
self.removePendingLocalUserEcho(for: runId)
1078+
self.runMessageScopesByRunID.removeValue(forKey: runId)
9691079
self.clearPendingRun(runId)
9701080
self.errorText = error.localizedDescription
9711081
self.logDiagnostic(
@@ -1027,7 +1137,8 @@ public final class OpenClawChatViewModel {
10271137
self.modelSelectionID = Self.defaultModelSelectionID
10281138
self.messages = []
10291139
self.pendingLocalUserEchoMessageIDsByRunID.removeAll()
1030-
self.provisionalFinalMessageIDs.removeAll()
1140+
self.runMessageScopesByRunID.removeAll()
1141+
self.provisionalFinalMessagesByID.removeAll()
10311142
self.sessionId = nil
10321143
self.pendingToolCallsById = [:]
10331144
self.streamingAssistantText = nil
@@ -1062,7 +1173,8 @@ public final class OpenClawChatViewModel {
10621173
self.modelSelectionID = Self.defaultModelSelectionID
10631174
self.messages = []
10641175
self.pendingLocalUserEchoMessageIDsByRunID.removeAll()
1065-
self.provisionalFinalMessageIDs.removeAll()
1176+
self.runMessageScopesByRunID.removeAll()
1177+
self.provisionalFinalMessagesByID.removeAll()
10661178
self.sessionId = nil
10671179
self.pendingToolCallsById = [:]
10681180
self.streamingAssistantText = nil
@@ -1090,6 +1202,8 @@ public final class OpenClawChatViewModel {
10901202
return
10911203
}
10921204

1205+
self.runMessageScopesByRunID.removeAll()
1206+
self.provisionalFinalMessagesByID.removeAll()
10931207
self.startBootstrap()
10941208
}
10951209

@@ -1599,7 +1713,8 @@ public final class OpenClawChatViewModel {
15991713

16001714
let reconciled = Self.reconcileMessageIDs(previous: self.messages, incoming: self.messages + [sanitized])
16011715
self.messages = Self.dedupeMessages(reconciled)
1602-
self.pruneProvisionalFinalMessageIDs()
1716+
self.pruneProvisionalFinalMessages()
1717+
self.pruneRunMessageScopes()
16031718
}
16041719

16051720
private func handleChatEvent(_ chat: OpenClawChatEventPayload) {
@@ -1686,16 +1801,28 @@ public final class OpenClawChatViewModel {
16861801
stopReason: "stop")
16871802
}
16881803

1689-
if self.hasCanonicalFinalMessageMatching(message) {
1804+
let runId = Self.normalizedRunID(chat.runId)
1805+
let scope = self.runMessageScope(for: runId)
1806+
guard self.isCurrentSession(scope.session) else { return }
1807+
guard let reconciliationKey = Self.finalMessageReconciliationKey(for: message) else { return }
1808+
1809+
if self.hasCanonicalFinalMessageMatching(message, scope: scope) {
1810+
if let runId {
1811+
self.runMessageScopesByRunID.removeValue(forKey: runId)
1812+
}
16901813
return
16911814
}
16921815

16931816
let reconciled = Self.reconcileMessageIDs(previous: self.messages, incoming: self.messages + [message])
16941817
self.messages = Self.dedupeMessages(reconciled)
16951818
if self.messages.contains(where: { $0.id == message.id }) {
1696-
self.provisionalFinalMessageIDs.insert(message.id)
1819+
self.provisionalFinalMessagesByID[message.id] = ProvisionalFinalMessage(
1820+
reconciliationKey: reconciliationKey,
1821+
runId: runId,
1822+
scope: scope)
16971823
}
1698-
self.pruneProvisionalFinalMessageIDs()
1824+
self.pruneProvisionalFinalMessages()
1825+
self.pruneRunMessageScopes()
16991826
}
17001827

17011828
private static func isAssistantMessage(_ message: OpenClawChatMessage) -> Bool {
@@ -1847,7 +1974,7 @@ public final class OpenClawChatViewModel {
18471974
}
18481975

18491976
private func removePendingLocalUserEcho(for runId: String) {
1850-
guard let messageID = self.pendingLocalUserEchoMessageIDsByRunID[runId] else { return }
1977+
guard let messageID = pendingLocalUserEchoMessageIDsByRunID[runId] else { return }
18511978
self.messages.removeAll { $0.id == messageID }
18521979
self.pendingLocalUserEchoMessageIDsByRunID[runId] = nil
18531980
}
@@ -1938,7 +2065,7 @@ public final class OpenClawChatViewModel {
19382065
guard let lastUserIndex = messages.lastIndex(where: { $0.role.lowercased() == "user" }) else {
19392066
return nil
19402067
}
1941-
guard let refreshKey = self.userRefreshIdentityKey(for: messages[lastUserIndex]) else {
2068+
guard let refreshKey = userRefreshIdentityKey(for: messages[lastUserIndex]) else {
19422069
return LatestUserTurn(
19432070
refreshKey: nil,
19442071
occurrence: 0,

‎apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatViewModelTests.swift‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ private func chatTextMessage(role: String, text: String, timestamp: Double, cont
88
if let contentId {
99
content["id"] = contentId
1010
}
11-
AnyCodable([
11+
return AnyCodable([
1212
"role": role,
1313
"content": [content],
1414
"timestamp": timestamp,
@@ -937,16 +937,11 @@ struct ChatViewModelTests {
937937
@Test func `later identical session reply does not adopt prior turn provisional final`() async throws {
938938
let sessionId = "sess-main"
939939
let now = Date().timeIntervalSince1970 * 1000
940-
let finalRefreshGate = SessionSubscribeGate()
941-
let historyCount = AsyncCounter()
942940
let history = historyPayload(sessionId: sessionId)
943941
let (transport, vm) = await makeViewModel(
944942
historyResponses: [history, history],
945-
requestHistoryHook: { _ in
946-
let count = await historyCount.increment()
947-
if count == 2 {
948-
await finalRefreshGate.wait()
949-
}
943+
sendMessageHook: { runId in
944+
OpenClawChatSendResponse(runId: runId, status: "pending")
950945
})
951946
try await loadAndWaitBootstrap(vm: vm, sessionId: sessionId)
952947

@@ -989,8 +984,6 @@ struct ChatViewModelTests {
989984
return okReplies.count == 2 && vm.messages.last?.timestamp == now + 4
990985
}
991986
}
992-
993-
await finalRefreshGate.release()
994987
}
995988

996989
@Test func `completion wait refreshes history and clears pending run`() async throws {
@@ -1649,7 +1642,11 @@ struct ChatViewModelTests {
16491642
}
16501643

16511644
@Test func `dedupes gateway echo of local user message`() async throws {
1652-
let (transport, vm) = await makeViewModel(historyResponses: [historyPayload()])
1645+
let (transport, vm) = await makeViewModel(
1646+
historyResponses: [historyPayload()],
1647+
sendMessageHook: { runId in
1648+
OpenClawChatSendResponse(runId: runId, status: "pending")
1649+
})
16531650

16541651
await MainActor.run { vm.load() }
16551652
try await waitUntil("bootstrap history loaded") { await MainActor.run { vm.messages.isEmpty } }

0 commit comments

Comments
 (0)