Polish streaming thinking traces

2026-08-30 15:26:31 -07:00
parent 922172fc60
commit 26d917e96a
10 changed files with 1128 additions and 122 deletions
@@ -40,9 +40,72 @@ struct SidebarItem: Identifiable, Hashable {
    var isRunning: Bool
 }

+enum StreamingAssistantContentUpdate: Equatable, Sendable {
+    case unchanged
+    case replace(String)
+}
+
+struct StreamingAssistantState: Equatable, Sendable {
+    private(set) var segmentText = ""
+    private(set) var initiatedToolCallIDs: Set<String> = []
+    private(set) var ephemeralTrace: String?
+    private(set) var traceRevision = 0
+
+    mutating func receiveDelta(_ text: String) -> StreamingAssistantContentUpdate {
+        guard !text.isEmpty else { return .unchanged }
+        segmentText += text
+        ephemeralTrace = nil
+        return .replace(segmentText)
+    }
+
+    mutating func receiveToolCall(id: String, status: String) -> StreamingAssistantContentUpdate {
+        guard status.lowercased() == "initiated", initiatedToolCallIDs.insert(id).inserted else {
+            return .unchanged
+        }
+
+        let trace = segmentText.trimmingCharacters(in: .whitespacesAndNewlines)
+        segmentText = ""
+        if !trace.isEmpty {
+            ephemeralTrace = trace
+            traceRevision += 1
+        }
+        return .replace("")
+    }
+
+    mutating func receiveDone(_ text: String) -> StreamingAssistantContentUpdate {
+        segmentText = ""
+        ephemeralTrace = nil
+        return .replace(text)
+    }
+
+    mutating func receiveError() -> StreamingAssistantContentUpdate {
+        segmentText = ""
+        ephemeralTrace = nil
+        return .unchanged
+    }
+}
+
+struct StreamingAssistantPresentation: Equatable, Sendable {
+    var messageID: String
+    var ephemeralTrace: String?
+    var traceRevision: Int
+}
+
 private struct PendingChatState {
    var chatID: String?
    var messages: [Message]
+    var streamingAssistantState = StreamingAssistantState()
+
+    var streamingAssistantPresentation: StreamingAssistantPresentation? {
+        guard let messageID = messages.last(where: { $0.id.hasPrefix("temp-assistant-") })?.id else {
+            return nil
+        }
+        return StreamingAssistantPresentation(
+            messageID: messageID,
+            ephemeralTrace: streamingAssistantState.ephemeralTrace,
+            traceRevision: streamingAssistantState.traceRevision
+        )
+    }
 }

 private enum ActiveSendContext: Hashable {
@@ -125,6 +188,7 @@ final class SybilViewModel {
    var quickQuestionSubmittedModel: String?
    var isQuickQuestionSending = false
    var isConvertingQuickQuestion = false
+    private var quickQuestionStreamingAssistantState = StreamingAssistantState()

    @ObservationIgnored
    private var hasBootstrapped = false
@@ -360,6 +424,30 @@ final class SybilViewModel {
        return canonical
    }

+    var displayedStreamingAssistantPresentation: StreamingAssistantPresentation? {
+        if case let .chat(chatID) = selectedItem,
+           let pending = pendingChatStates[chatID] {
+            return pending.streamingAssistantPresentation
+        }
+
+        if draftKind == .chat, let pending = pendingDraftChatState {
+            return pending.streamingAssistantPresentation
+        }
+
+        return nil
+    }
+
+    var quickQuestionStreamingAssistantPresentation: StreamingAssistantPresentation? {
+        guard let messageID = quickQuestionMessages.last(where: { $0.id.hasPrefix("temp-assistant-quick-") })?.id else {
+            return nil
+        }
+        return StreamingAssistantPresentation(
+            messageID: messageID,
+            ephemeralTrace: quickQuestionStreamingAssistantState.ephemeralTrace,
+            traceRevision: quickQuestionStreamingAssistantState.traceRevision
+        )
+    }
+
    var displayedSearch: SearchDetail? {
        if case let .search(searchID) = selectedItem,
           let activeSearch = activeSearchDetails[searchID] {
@@ -620,6 +708,7 @@ final class SybilViewModel {
        quickQuestionTask = nil
        quickQuestionRunID = nil
        isQuickQuestionSending = false
+        quickQuestionStreamingAssistantState = StreamingAssistantState()
    }

    @discardableResult
@@ -1140,6 +1229,7 @@ final class SybilViewModel {
        quickQuestionSubmittedPrompt = prompt
        quickQuestionSubmittedProvider = provider
        quickQuestionSubmittedModel = model
+        quickQuestionStreamingAssistantState = StreamingAssistantState()
        quickQuestionMessages = [
            Message(
                id: "temp-assistant-quick-\(UUID().uuidString)",
@@ -1184,6 +1274,8 @@ final class SybilViewModel {
                return
            }

+            let update = quickQuestionStreamingAssistantState.receiveError()
+            applyQuickQuestionAssistantContentUpdate(update)
            quickQuestionError = normalizeAPIError(error)
            SybilLog.error(SybilLog.ui, "Quick question failed", error: error)
        }
@@ -1198,17 +1290,16 @@ final class SybilViewModel {
            upsertQuickQuestionToolCallMessage(payload)

        case let .delta(payload):
-            guard !payload.text.isEmpty else { return }
-            mutateQuickQuestionAssistantMessage { existing in
-                existing + payload.text
-            }
+            let update = quickQuestionStreamingAssistantState.receiveDelta(payload.text)
+            applyQuickQuestionAssistantContentUpdate(update)

        case let .done(payload):
-            mutateQuickQuestionAssistantMessage { _ in
-                payload.text
-            }
+            let update = quickQuestionStreamingAssistantState.receiveDone(payload.text)
+            applyQuickQuestionAssistantContentUpdate(update)

        case let .error(payload):
+            let update = quickQuestionStreamingAssistantState.receiveError()
+            applyQuickQuestionAssistantContentUpdate(update)
            await streamStatus.setError(payload.message)

        case .ignored:
@@ -1607,21 +1698,28 @@ final class SybilViewModel {
            }

            serverActiveChatIDs.remove(chatID)
-            pendingChatStates[chatID] = nil
            await refreshCollections(preferredSelection: selectedItem, refreshSelection: false, attachVisibleActiveRun: false)

            if selectedItem == selection, draftKind == nil {
-                selectedChat = try await client.getChat(chatID: chatID)
-                selectedSearch = nil
+                let refreshedChat = try await client.getChat(chatID: chatID)
+                if selectedItem == selection, draftKind == nil {
+                    selectedChat = refreshedChat
+                    selectedSearch = nil
+                }
            }
+            pendingChatStates[chatID] = nil
        } catch {
            serverActiveChatIDs.remove(chatID)
-            pendingChatStates[chatID] = nil
+            defer {
+                pendingChatStates[chatID] = nil
+            }

            if isCancellation(error) {
                return
            }

+            applyPendingAssistantError(chatID: chatID)
+
            if isActiveStreamNotFound(error) {
                SybilLog.info(SybilLog.app, "Active chat stream \(chatID) no longer exists")
            } else if shouldSuppressInactiveTransportError(error) {
@@ -1635,8 +1733,11 @@ final class SybilViewModel {

            if selectedItem == selection, draftKind == nil {
                do {
-                    selectedChat = try await client().getChat(chatID: chatID)
-                    selectedSearch = nil
+                    let refreshedChat = try await client().getChat(chatID: chatID)
+                    if selectedItem == selection, draftKind == nil {
+                        selectedChat = refreshedChat
+                        selectedSearch = nil
+                    }
                } catch {
                    SybilLog.warning(SybilLog.app, "Chat refresh after attach failure failed: \(SybilLog.describe(error))")
                }
@@ -1856,6 +1957,13 @@ final class SybilViewModel {
        } else {
            pendingDraftChatState = PendingChatState(chatID: nil, messages: optimisticMessages)
        }
+        defer {
+            if let resolvedChatID = chatID {
+                pendingChatStates[resolvedChatID] = nil
+            } else {
+                clearPendingChatState(for: sendContext)
+            }
+        }
        requestChatBottomPin()

        if chatID == nil {
@@ -2045,17 +2153,13 @@ final class SybilViewModel {
            upsertPendingToolCallMessage(payload, chatID: chatID)

        case let .delta(payload):
-            guard !payload.text.isEmpty else { return }
-            mutatePendingAssistantMessage(chatID: chatID) { existing in
-                existing + payload.text
-            }
+            applyPendingAssistantDelta(payload.text, chatID: chatID)

        case let .done(payload):
-            mutatePendingAssistantMessage(chatID: chatID) { _ in
-                payload.text
-            }
+            finalizePendingAssistant(payload.text, chatID: chatID)

        case let .error(payload):
+            applyPendingAssistantError(chatID: chatID)
            await streamStatus.setError(payload.message)

        case .ignored:
@@ -2233,29 +2337,68 @@ final class SybilViewModel {
        }
    }

-    private func mutatePendingAssistantMessage(chatID: String, _ transform: (String) -> String) {
+    private func applyPendingAssistantDelta(_ text: String, chatID: String) {
        guard var pending = pendingChatStates[chatID], !pending.messages.isEmpty else {
            return
        }

-        let index = pending.messages.indices.last { pending.messages[$0].id.hasPrefix("temp-assistant-") }
-        guard let index else {
-            return
-        }
-
-        var message = pending.messages[index]
-        message.content = transform(message.content)
-        pending.messages[index] = message
+        let update = pending.streamingAssistantState.receiveDelta(text)
+        applyStreamingAssistantContentUpdate(
+            update,
+            to: &pending.messages,
+            assistantMessagePrefix: "temp-assistant-"
+        )
        pendingChatStates[chatID] = pending
    }

-    private func mutateQuickQuestionAssistantMessage(_ transform: (String) -> String) {
-        let index = quickQuestionMessages.indices.last { quickQuestionMessages[$0].id.hasPrefix("temp-assistant-quick-") }
-        guard let index else {
+    private func finalizePendingAssistant(_ text: String, chatID: String) {
+        guard var pending = pendingChatStates[chatID], !pending.messages.isEmpty else {
            return
        }

-        quickQuestionMessages[index].content = transform(quickQuestionMessages[index].content)
+        let update = pending.streamingAssistantState.receiveDone(text)
+        applyStreamingAssistantContentUpdate(
+            update,
+            to: &pending.messages,
+            assistantMessagePrefix: "temp-assistant-"
+        )
+        pendingChatStates[chatID] = pending
+    }
+
+    private func applyPendingAssistantError(chatID: String) {
+        guard var pending = pendingChatStates[chatID] else {
+            return
+        }
+
+        let update = pending.streamingAssistantState.receiveError()
+        applyStreamingAssistantContentUpdate(
+            update,
+            to: &pending.messages,
+            assistantMessagePrefix: "temp-assistant-"
+        )
+        pendingChatStates[chatID] = pending
+    }
+
+    private func applyQuickQuestionAssistantContentUpdate(_ update: StreamingAssistantContentUpdate) {
+        applyStreamingAssistantContentUpdate(
+            update,
+            to: &quickQuestionMessages,
+            assistantMessagePrefix: "temp-assistant-quick-"
+        )
+    }
+
+    private func applyStreamingAssistantContentUpdate(
+        _ update: StreamingAssistantContentUpdate,
+        to messages: inout [Message],
+        assistantMessagePrefix: String
+    ) {
+        guard case let .replace(content) = update,
+              let index = messages.indices.last(where: { messages[$0].id.hasPrefix(assistantMessagePrefix) })
+        else {
+            return
+        }
+
+        messages[index].content = content
    }

    private func upsertPendingToolCallMessage(_ payload: CompletionStreamToolCall, chatID: String) {
@@ -2263,6 +2406,16 @@ final class SybilViewModel {
            return
        }

+        let update = pending.streamingAssistantState.receiveToolCall(
+            id: payload.toolCallId,
+            status: payload.status
+        )
+        applyStreamingAssistantContentUpdate(
+            update,
+            to: &pending.messages,
+            assistantMessagePrefix: "temp-assistant-"
+        )
+
        if let existingIndex = pending.messages.firstIndex(where: { $0.toolCallMetadata?.toolCallId == payload.toolCallId || $0.id == "temp-tool-\(payload.toolCallId)" }) {
            pending.messages[existingIndex] = toolCallMessage(for: payload, id: pending.messages[existingIndex].id)
            pendingChatStates[chatID] = pending
@@ -2281,6 +2434,12 @@ final class SybilViewModel {
    }

    private func upsertQuickQuestionToolCallMessage(_ payload: CompletionStreamToolCall) {
+        let update = quickQuestionStreamingAssistantState.receiveToolCall(
+            id: payload.toolCallId,
+            status: payload.status
+        )
+        applyQuickQuestionAssistantContentUpdate(update)
+
        if let existingIndex = quickQuestionMessages.firstIndex(where: { $0.toolCallMetadata?.toolCallId == payload.toolCallId || $0.id == "temp-tool-\(payload.toolCallId)" }) {
            quickQuestionMessages[existingIndex] = toolCallMessage(for: payload, id: quickQuestionMessages[existingIndex].id)
            return
@@ -2430,7 +2589,8 @@ final class SybilViewModel {

    private func clearPendingChatState(for context: ActiveSendContext) {
        switch context {
-        case .draftChat:
+        case let .draftChat(identity):
+            guard draftKind == .chat, draftIdentity == identity else { return }
            pendingDraftChatState = nil
        case let .chat(chatID):
            pendingChatStates[chatID] = nil