feat(keyboard): ship AI mode surface with streaming search answers

Add the AI keyboard tab, Agent settings, and user-owned LLM key path for 1.7.0, including streaming answers and web-search transports without the built-in DeepSeek fallback.
2026-08-11 01:06:27 +08:00
parent f6212dcd2c
commit 3de665d254
81 changed files with 4467 additions and 398 deletions
@@ -54,9 +54,17 @@ final class KeyboardFlowCoordinator {
    private var editHostConfirmed = false
    private var cancelledEditUtteranceIDs: Set<UUID> = []
    private var cancelledDictationUtteranceIDs: Set<UUID> = []
+    private var cancelledAIUtteranceIDs: Set<UUID> = []
    var onEditHostRecordingConfirmed: () -> Void = {}
    var onEditResult: (FlowResult) -> Void = { _ in }
    var onEditFailure: (String) -> Void = { _ in }
+    var onAIUtterancePrepared: (UUID) -> Void = { _ in }
+    var onAIRecordingStarted: (UUID) -> Void = { _ in }
+    var onAIRecognitionStarted: (UUID) -> Void = { _ in }
+    var onAITranscript: (String, UUID, FlowResult.Status) -> Void = { _, _, _ in }
+    var onAIStreamingAnswer: (String, UUID) -> Void = { _, _ in }
+    var onAIResult: (FlowResult) -> Void = { _ in }
+    var onAIFailure: (String, UUID?) -> Void = { _, _ in }
    /// Utterance whose final result we already inserted (or failed). Prevents
    /// `adoptHostBusyStateIfNeeded` from re-entering `.processing` after a
    /// stale App Group snapshot still says `reason=processing`.
@@ -178,6 +186,7 @@ final class KeyboardFlowCoordinator {
        adoptPendingResultIfNeeded()
        consumePendingFlowDeliveryIfNeeded()
        consumeEditStartFailureIfNeeded()
+        consumeAIStartFailureIfNeeded()

        recoverFromDeadHostIfNeeded()

@@ -243,6 +252,15 @@ final class KeyboardFlowCoordinator {
        completeEditResult(result, outcome: .rejected)
    }

+    private func consumeAIStartFailureIfNeeded() {
+        guard currentUtteranceRequest?.isAIQuestion == true,
+              let result = matchingResult(),
+              isTerminalFailure(result) else {
+            return
+        }
+        completeAIFailure(result)
+    }
+
    private func recomputeMicVoiceAvailability() {
        FlowSessionBridge.reloadFromDisk()
        let readySnapshot = FlowSessionBridge.readySnapshot()
@@ -617,13 +635,76 @@ final class KeyboardFlowCoordinator {
        startUtterance(.editLastInput(reference))
    }

+    func beginAIRecording(
+        conversationID: UUID
+    ) -> FlowUtteranceStartDisposition {
+        startUtterance(.aiQuestion(conversationID: conversationID))
+    }
+
+    func stopAIRecording() {
+        guard currentUtteranceRequest?.isAIQuestion == true else { return }
+        pressEnded()
+    }
+
+    func cancelAIRecording() {
+        guard currentUtteranceRequest?.isAIQuestion == true else { return }
+        recordWhenHostReady = false
+        recordAfterHandoff = false
+        isPendingFlowStart = false
+        flowStartDeadline = 0
+        coldStartDebouncer.reset()
+        stopHostReadyWait()
+        stopFlowWatchdog()
+        stopUtteranceCountdown()
+        ExtensionScreenWakeLock.release()
+        state.level = 0
+        state.phase = .idle
+        state.lastTranscript = ""
+
+        guard let utteranceID = currentUtteranceId,
+              isFlowRecording || isAwaitingFlowResult else {
+            clearUnissuedUtterance()
+            isFlowRecording = false
+            isAwaitingFlowResult = false
+            recomputeMicVoiceAvailability()
+            return
+        }
+
+        cancelledAIUtteranceIDs.insert(utteranceID)
+        writeCommand(.abort)
+        isFlowRecording = false
+        isAwaitingFlowResult = true
+        startFlowResultWatchdog()
+        recomputeMicVoiceAvailability()
+    }
+
+    func endAIConversation(_ conversationID: UUID) {
+        guard let sessionID = FlowSessionBridge.readySnapshot()?.sessionId
+            ?? activeSessionId else {
+            return
+        }
+        FlowSessionBridge.writeCommand(
+            FlowCommand(
+                sessionId: sessionID,
+                utteranceId: UUID(),
+                commandSeq: nextCommandSeq(),
+                action: .endAIConversation,
+                localeId: state.localeId,
+                aiConversationID: conversationID
+            )
+        )
+    }
+
    func stopEditRecording() {
        guard currentUtteranceRequest?.isEdit == true else { return }
        pressEnded()
    }

    func cancelCurrentDictation() {
-        guard currentUtteranceRequest?.isEdit != true else { return }
+        guard currentUtteranceRequest?.isEdit != true,
+              currentUtteranceRequest?.isAIQuestion != true else {
+            return
+        }

        recordWhenHostReady = false
        recordAfterHandoff = false
@@ -733,6 +814,50 @@ final class KeyboardFlowCoordinator {
        resetEditTransportState()
    }

+    private func completeAIResult(_ result: FlowResult) {
+        onAIResult(result)
+        FlowSessionBridge.writeAck(
+            FlowAck(
+                sessionId: result.sessionId,
+                utteranceId: result.utteranceId,
+                commandSeq: result.commandSeq,
+                hostGeneration: result.hostGeneration,
+                revision: result.revision
+            )
+        )
+        FlowSessionBridge.setPendingKeyboardUtteranceId(nil)
+        lastConsumedUtteranceId = result.utteranceId
+        lastStoppedUtteranceId = nil
+        resetEditTransportState()
+        state.phase = .idle
+        state.lastTranscript = ""
+        recomputeMicVoiceAvailability()
+    }
+
+    private func completeAIFailure(_ result: FlowResult) {
+        onAIFailure(
+            result.text ?? ExtL10n.string("keyboard.ai.error.requestFailed"),
+            result.utteranceId
+        )
+        FlowSessionBridge.writeAck(
+            FlowAck(
+                sessionId: result.sessionId,
+                utteranceId: result.utteranceId,
+                commandSeq: result.commandSeq,
+                hostGeneration: result.hostGeneration,
+                revision: result.revision,
+                deliveryOutcome: .rejected
+            )
+        )
+        FlowSessionBridge.setPendingKeyboardUtteranceId(nil)
+        lastConsumedUtteranceId = result.utteranceId
+        lastStoppedUtteranceId = nil
+        resetEditTransportState()
+        state.phase = .idle
+        state.lastTranscript = ""
+        recomputeMicVoiceAvailability()
+    }
+
    func pressBegan() {
        _ = startUtterance(.dictation)
    }
@@ -762,6 +887,9 @@ final class KeyboardFlowCoordinator {
        currentUtteranceId = utteranceID
        currentUtteranceRequest = request
        editHostConfirmed = false
+        if request.isAIQuestion {
+            onAIUtterancePrepared(utteranceID)
+        }
        currentStartDeadlineAt = Date().timeIntervalSince1970
            + FlowSessionKeys.utteranceStartBudget

@@ -863,7 +991,10 @@ final class KeyboardFlowCoordinator {
        writeCommand(.stopRecording)
        debug("pressEnded wrote stop command")
        state.phase = .processing
-        if currentUtteranceRequest?.isEdit != true {
+        if currentUtteranceRequest?.isAIQuestion == true,
+           let currentUtteranceId {
+            onAIRecognitionStarted(currentUtteranceId)
+        } else if currentUtteranceRequest?.isEdit != true {
            state.lastTranscript = ExtL10n.string("keyboard.flow.transcribing")
        }
        startFlowResultWatchdog()
@@ -923,6 +1054,16 @@ final class KeyboardFlowCoordinator {
                state.lastTranscript = ""
                return
            }
+            if currentUtteranceRequest?.isAIQuestion == true {
+                onAIFailure(
+                    ExtL10n.string("keyboard.error.manualOpenForFlow"),
+                    currentUtteranceId
+                )
+                resetEditTransportState()
+                state.phase = .idle
+                state.lastTranscript = ""
+                return
+            }
            showManualOpenHint(path: "startflow")
            recomputeMicVoiceAvailability()
            return
@@ -972,13 +1113,16 @@ final class KeyboardFlowCoordinator {
            commandSeq: nextCommandSeq(),
            action: action,
            localeId: state.localeId,
-            fieldContext: action == .stopRecording ? fieldContextProvider() : nil,
+            fieldContext: action == .stopRecording && !request.isAIQuestion
+                ? fieldContextProvider()
+                : nil,
            utteranceMode: mode,
            editSourceText: action == .startRecording
                ? request.editSourceText
                : nil,
            sourceHistoryEntryID: request.sourceHistoryEntryID,
            sourceHistoryEntryRevision: request.sourceHistoryEntryRevision,
+            aiConversationID: request.aiConversationID,
            startDeadlineAt: action == .startRecording ? currentStartDeadlineAt : nil,
            processingDeadlineAt: action == .stopRecording && request.isEdit
                ? Date().timeIntervalSince1970
@@ -1027,6 +1171,10 @@ final class KeyboardFlowCoordinator {
               consumeCancelledEditResultIfNeeded(result) {
                return
            }
+            if let result = matchingResult(),
+               consumeCancelledAIResultIfNeeded(result) {
+                return
+            }
            if let result = matchingResult(), result.status == .final, let text = result.text, !text.isEmpty {
                isAwaitingFlowResult = false
                stopFlowWatchdog()
@@ -1035,6 +1183,10 @@ final class KeyboardFlowCoordinator {
                    onEditResult(result)
                    return
                }
+                if result.resolvedUtteranceMode == .aiQuestion {
+                    completeAIResult(result)
+                    return
+                }
                textInserter.handleFlowTranscript(
                    TranscriptionDelivery(
                        text: text,
@@ -1083,6 +1235,10 @@ final class KeyboardFlowCoordinator {
                    state.phase = .idle
                    return
                }
+                if result.resolvedUtteranceMode == .aiQuestion {
+                    completeAIFailure(result)
+                    return
+                }
                FlowSessionBridge.writeAck(
                    FlowAck(
                        sessionId: result.sessionId,
@@ -1172,6 +1328,30 @@ final class KeyboardFlowCoordinator {
        return true
    }

+    private func consumeCancelledAIResultIfNeeded(_ result: FlowResult) -> Bool {
+        guard cancelledAIUtteranceIDs.contains(result.utteranceId),
+              result.status == .final || isTerminalFailure(result) else {
+            return false
+        }
+        FlowSessionBridge.writeAck(
+            FlowAck(
+                sessionId: result.sessionId,
+                utteranceId: result.utteranceId,
+                commandSeq: result.commandSeq,
+                hostGeneration: result.hostGeneration,
+                revision: result.revision,
+                deliveryOutcome: .rejected
+            )
+        )
+        cancelledAIUtteranceIDs.remove(result.utteranceId)
+        lastConsumedUtteranceId = result.utteranceId
+        resetEditTransportState()
+        state.phase = .idle
+        state.lastTranscript = ""
+        recomputeMicVoiceAvailability()
+        return true
+    }
+
    private func adoptPendingResultIfNeeded() {
        guard !isAwaitingFlowResult, currentUtteranceId == nil,
              let pendingId = FlowSessionBridge.pendingKeyboardUtteranceId(),
@@ -1343,6 +1523,18 @@ final class KeyboardFlowCoordinator {
            )
            return
        }
+        if let id = currentUtteranceId,
+           cancelledAIUtteranceIDs.remove(id) != nil {
+            stopUtteranceCountdown()
+            stopFlowWatchdog()
+            ExtensionScreenWakeLock.release()
+            resetEditTransportState()
+            state.level = 0
+            state.phase = .idle
+            state.lastTranscript = ""
+            recomputeMicVoiceAvailability()
+            return
+        }
        if deliverRawFallbackIfAvailable(reason: "hostDisconnected") {
            return
        }
@@ -1367,6 +1559,14 @@ final class KeyboardFlowCoordinator {
            recomputeMicVoiceAvailability()
            return
        }
+        if currentUtteranceRequest?.isAIQuestion == true {
+            onAIFailure(message, currentUtteranceId)
+            resetEditTransportState()
+            state.phase = .idle
+            state.lastTranscript = ""
+            recomputeMicVoiceAvailability()
+            return
+        }
        state.phase = .error(.flowSessionExpired, message: message)
        scheduleAutoClearError()
        recomputeMicVoiceAvailability()
@@ -1469,6 +1669,12 @@ final class KeyboardFlowCoordinator {
            if currentUtteranceRequest?.isEdit == true {
                onEditFailure(ExtL10n.string("keyboard.edit.error.startTimeout"))
                abortEditRecording()
+            } else if currentUtteranceRequest?.isAIQuestion == true {
+                onAIFailure(
+                    ExtL10n.string("keyboard.ai.error.startTimeout"),
+                    currentUtteranceId
+                )
+                cancelAIRecording()
            } else {
                state.phase = .error(
                    .hostAudioUnavailable,
@@ -1538,6 +1744,10 @@ final class KeyboardFlowCoordinator {
        state.phase = currentUtteranceRequest?.isEdit == true
            ? .requestingPermissions
            : .recording
+        if currentUtteranceRequest?.isAIQuestion == true,
+           let currentUtteranceId {
+            onAIRecordingStarted(currentUtteranceId)
+        }
        recomputeMicVoiceAvailability()
        if let view = wakeLockView() {
            ExtensionScreenWakeLock.acquire(from: view)
@@ -1658,11 +1868,19 @@ final class KeyboardFlowCoordinator {
        guard isFlowRecording || isAwaitingFlowResult else { return }
        switch state.phase {
        case .recording, .processing:
-            if let result = matchingResult(),
-               result.status == .partial || result.status == .rawReady,
+            guard let result = matchingResult() else { return }
+            if result.status == .streaming,
+               result.resolvedUtteranceMode == .aiQuestion {
+                onAIStreamingAnswer(result.text ?? "", result.utteranceId)
+                return
+            }
+            if (result.status == .partial || result.status == .rawReady),
               let partial = result.text,
               !partial.isEmpty {
                state.lastTranscript = partial
+                if result.resolvedUtteranceMode == .aiQuestion {
+                    onAITranscript(partial, result.utteranceId, result.status)
+                }
            }
        default:
            break
@@ -1679,11 +1897,21 @@ final class KeyboardFlowCoordinator {
        let isCancelledDictation = currentUtteranceId.map {
            cancelledDictationUtteranceIDs.contains($0)
        } ?? false
-        let resultTimeout = isCancelledEdit || isCancelledDictation
-            ? FlowSessionKeys.utteranceStartBudget
-            : (currentUtteranceRequest?.isEdit != true
-                ? FlowWatchdog.resultTimeout(engineMode: state.engineMode)
-                : FlowSessionKeys.editLastInputProcessingBudget)
+        let isCancelledAI = currentUtteranceId.map {
+            cancelledAIUtteranceIDs.contains($0)
+        } ?? false
+        let resultTimeout: TimeInterval
+        if isCancelledEdit || isCancelledDictation || isCancelledAI {
+            resultTimeout = FlowSessionKeys.utteranceStartBudget
+        } else if currentUtteranceRequest?.isAIQuestion == true {
+            resultTimeout = FlowSessionKeys.keyboardAIResultTimeout(
+                engineMode: state.engineMode
+            )
+        } else if currentUtteranceRequest?.isEdit == true {
+            resultTimeout = FlowSessionKeys.editLastInputProcessingBudget
+        } else {
+            resultTimeout = FlowWatchdog.resultTimeout(engineMode: state.engineMode)
+        }
        debug("resultWatchdog started timeout=\(Int(resultTimeout))s engine=\(state.engineMode)")
        flowWatchdogTask = Task { @MainActor [weak self] in
            while let self, !Task.isCancelled {
@@ -1700,6 +1928,10 @@ final class KeyboardFlowCoordinator {
                   self.consumeCancelledEditResultIfNeeded(result) {
                    return
                }
+                if let result = self.matchingResult(),
+                   self.consumeCancelledAIResultIfNeeded(result) {
+                    return
+                }
                if let result = self.matchingResult(), result.status == .final, let text = result.text, !text.isEmpty {
                    self.isAwaitingFlowResult = false
                    self.stopFlowWatchdog()
@@ -1707,6 +1939,10 @@ final class KeyboardFlowCoordinator {
                        self.onEditResult(result)
                        return
                    }
+                    if result.resolvedUtteranceMode == .aiQuestion {
+                        self.completeAIResult(result)
+                        return
+                    }
                    self.textInserter.handleFlowTranscript(
                        TranscriptionDelivery(
                            text: text,
@@ -1749,6 +1985,10 @@ final class KeyboardFlowCoordinator {
                        self.state.phase = .idle
                        return
                    }
+                    if result.resolvedUtteranceMode == .aiQuestion {
+                        self.completeAIFailure(result)
+                        return
+                    }
                    FlowSessionBridge.writeAck(
                        FlowAck(
                            sessionId: result.sessionId,
@@ -1823,6 +2063,14 @@ final class KeyboardFlowCoordinator {
                        self.recomputeMicVoiceAvailability()
                        return
                    }
+                    if let id = self.currentUtteranceId,
+                       self.cancelledAIUtteranceIDs.remove(id) != nil {
+                        self.resetEditTransportState()
+                        self.state.phase = .idle
+                        self.state.lastTranscript = ""
+                        self.recomputeMicVoiceAvailability()
+                        return
+                    }
                    if self.currentUtteranceRequest?.isEdit == true {
                        self.isAwaitingFlowResult = false
                        self.stopFlowWatchdog()
@@ -1837,6 +2085,21 @@ final class KeyboardFlowCoordinator {
                        self.state.phase = .idle
                        return
                    }
+                    if self.currentUtteranceRequest?.isAIQuestion == true {
+                        let utteranceID = self.currentUtteranceId
+                        self.isAwaitingFlowResult = false
+                        self.stopFlowWatchdog()
+                        self.writeCommand(.abort)
+                        self.onAIFailure(
+                            ExtL10n.string("keyboard.ai.error.requestTimeout"),
+                            utteranceID
+                        )
+                        self.resetEditTransportState()
+                        self.state.phase = .idle
+                        self.state.lastTranscript = ""
+                        self.recomputeMicVoiceAvailability()
+                        return
+                    }
                    if self.deliverRawFallbackIfAvailable(reason: "resultTimeout") {
                        return
                    }