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
@@ -0,0 +1,162 @@
+// AIKeyboardCoordinator.swift
+// OSGKeyboard · Keyboard Extension
+//
+// Owns the temporary AI-mode UI state. Audio, ASR, and LLM work remain in the
+// shared Flow transport and host app; this coordinator never performs network
+// work and never inserts an answer before explicit user confirmation.
+
+import Foundation
+import OSGKeyboardShared
+
+@MainActor
+final class AIKeyboardCoordinator {
+    private let state: KeyboardState
+    private let flow: KeyboardFlowCoordinator
+    private let insertAnswer: (AIAnswer) -> Bool
+    private let performReturn: () -> Void
+
+    init(
+        state: KeyboardState,
+        flow: KeyboardFlowCoordinator,
+        insertAnswer: @escaping (AIAnswer) -> Bool,
+        performReturn: @escaping () -> Void
+    ) {
+        self.state = state
+        self.flow = flow
+        self.insertAnswer = insertAnswer
+        self.performReturn = performReturn
+    }
+
+    func beginNewPresentation() {
+        endConversationIfNeeded()
+        state.aiSession.enter()
+    }
+
+    func enterIfNeeded() {
+        guard !state.aiSession.isActive else { return }
+        state.aiSession.enter()
+    }
+
+    func leave() {
+        if state.aiSession.isBusy {
+            flow.cancelAIRecording()
+        }
+        endConversationIfNeeded()
+        state.aiSession.leave()
+    }
+
+    func toggleMicrophone() {
+        switch state.aiSession.phase {
+        case .listening:
+            guard let utteranceID = state.aiSession.activeUtteranceID else { return }
+            state.aiSession.beginRecognizing(utteranceID: utteranceID)
+            flow.stopAIRecording()
+        case .idle, .ready, .awaitingSend, .inserted, .sent, .failed:
+            enterIfNeeded()
+            guard let conversationID = state.aiSession.conversationID else { return }
+            let disposition = flow.beginAIRecording(conversationID: conversationID)
+            if case .rejected(let rejection) = disposition {
+                state.aiSession.fail(message(for: rejection), utteranceID: nil)
+            }
+        case .inactive:
+            enterIfNeeded()
+            toggleMicrophone()
+        case .preparing, .recognizing, .generating:
+            break
+        }
+    }
+
+    func cancel() {
+        guard state.aiSession.isBusy else { return }
+        flow.cancelAIRecording()
+        state.aiSession.cancelCurrentWork()
+    }
+
+    func sendLatestAnswer() {
+        if state.aiSession.canInsert, let answer = state.aiSession.answer {
+            guard insertAnswer(answer) else { return }
+            state.aiSession.markAnswerInserted(
+                offersSend: state.returnKeyRole == .send
+            )
+        } else if state.aiSession.canSend {
+            state.aiSession.markAnswerSent()
+            // Let the host consume the inserted answer before issuing Return.
+            Task { @MainActor [weak self] in
+                await Task.yield()
+                self?.performReturn()
+            }
+        }
+    }
+
+    func utterancePrepared(_ utteranceID: UUID) {
+        state.aiSession.beginPreparing(utteranceID: utteranceID)
+    }
+
+    func recordingStarted(_ utteranceID: UUID) {
+        state.aiSession.beginListening(utteranceID: utteranceID)
+    }
+
+    func recognitionStarted(_ utteranceID: UUID) {
+        state.aiSession.beginRecognizing(utteranceID: utteranceID)
+    }
+
+    func receiveTranscript(
+        _ transcript: String,
+        utteranceID: UUID,
+        status: FlowResult.Status
+    ) {
+        state.aiSession.updateTranscript(transcript, utteranceID: utteranceID)
+        if status == .rawReady {
+            state.aiSession.beginGenerating(
+                question: transcript,
+                utteranceID: utteranceID
+            )
+        }
+    }
+
+    func receivePartialAnswer(_ draft: String, utteranceID: UUID) {
+        state.aiSession.receivePartialAnswer(draft, utteranceID: utteranceID)
+    }
+
+    func receive(result: FlowResult) {
+        guard result.resolvedUtteranceMode == .aiQuestion else {
+            return
+        }
+        guard result.aiConversationID == state.aiSession.conversationID,
+              let answer = result.text,
+              !answer.isEmpty else {
+            state.aiSession.fail(
+                ExtL10n.string("keyboard.ai.error.requestFailed"),
+                utteranceID: result.utteranceId
+            )
+            return
+        }
+        state.aiSession.receiveAnswer(answer, utteranceID: result.utteranceId)
+    }
+
+    func fail(_ message: String, utteranceID: UUID?) {
+        state.aiSession.fail(message, utteranceID: utteranceID)
+    }
+
+    private func endConversationIfNeeded() {
+        guard let conversationID = state.aiSession.conversationID else { return }
+        flow.endAIConversation(conversationID)
+    }
+
+    private func message(for rejection: FlowUtteranceStartRejection) -> String {
+        switch rejection {
+        case .onboardingIncomplete:
+            return ExtL10n.string("keyboard.hint.finishSetupInApp")
+        case .missingAPIKey:
+            return ExtL10n.string("keyboard.ai.error.missingAPIKey")
+        case .noFullAccess:
+            return ExtL10n.string("keyboard.error.fullAccessRequired")
+        case .appGroupUnavailable:
+            return ExtL10n.string("keyboard.error.appGroupCommunication")
+        case .hostUnavailable:
+            return ExtL10n.string("keyboard.flow.hostDisconnected")
+        case .pipelineBusy:
+            return ExtL10n.string("keyboard.ai.error.pipelineBusy")
+        }
+    }
+}
@@ -42,10 +42,7 @@ public struct AppGroupPersistor {
        state.handednessPreference = store.handednessPreference
        state.cursorDragNavigationEnabled = store.cursorDragNavigationEnabled
        state.keyboardHapticIntensity = store.keyboardHapticIntensity
-        state.micDisabled = store.isCloudAPIKeyMissingForVoiceInput
-        state.micDisabledHint = store.isCloudAPIKeyMissingForVoiceInput
-            ? ExtL10n.string("keyboard.mic.disabled.missingApiKey")
-            : ""
+        applyAPIKeyAvailability(store: store, into: state)

        #if DEBUG
        // Print a masked view of the live App Group config so we can see
@@ -95,10 +92,26 @@ public struct AppGroupPersistor {
        state.handednessPreference = store.handednessPreference
        state.cursorDragNavigationEnabled = store.cursorDragNavigationEnabled
        state.keyboardHapticIntensity = store.keyboardHapticIntensity
-        state.micDisabled = store.isCloudAPIKeyMissingForVoiceInput
-        state.micDisabledHint = store.isCloudAPIKeyMissingForVoiceInput
-            ? ExtL10n.string("keyboard.mic.disabled.missingApiKey")
-            : ""
+        applyAPIKeyAvailability(store: store, into: state)
+    }
+
+    /// Cloud without ASR/LLM keys blocks the mic. Local ASR still works when
+    /// the polish key is missing — show a soft tip above the mic instead.
+    private func applyAPIKeyAvailability(
+        store: AppGroupStore,
+        into state: KeyboardViewController.State
+    ) {
+        state.aiServiceAvailable = !store.isPolishKeyMissing
+        if store.isCloudAPIKeyMissingForVoiceInput {
+            state.micDisabled = true
+            state.micDisabledHint = ExtL10n.string("keyboard.mic.disabled.missingApiKey")
+        } else if store.isPolishKeyMissing {
+            state.micDisabled = false
+            state.micDisabledHint = ExtL10n.string("keyboard.mic.hint.missingPolishApiKey")
+        } else {
+            state.micDisabled = false
+            state.micDisabledHint = ""
+        }
    }

    /// Persist `mode` to the App Group store.
@@ -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
                    }
@@ -98,6 +98,42 @@ final class KeyboardTextInserter {
        OSGLog.keyboardExt.info("flow insert length=\(trimmed.count, privacy: .public)")
    }

+    /// Insert one explicitly confirmed AI answer and enqueue history/statistics
+    /// only after the field mutation has been issued.
+    @discardableResult
+    func insertAIAnswer(_ answer: AIAnswer) -> Bool {
+        let trimmed = answer.text.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !trimmed.isEmpty else { return false }
+
+        let separator = DictationTextComposer.insertionSeparator(
+            previousContext: contextBeforeInput(),
+            insertion: trimmed
+        )
+        let inserted = separator + trimmed
+        insertText(inserted)
+
+        let mutation = HistoryMutation(
+            action: .append,
+            entryID: answer.id,
+            text: trimmed,
+            engineMode: state.engineMode,
+            source: .ai,
+            usageCategory: .ai
+        )
+        HistoryMutationOutbox.enqueue(mutation)
+        recordLastInsertion(
+            inserted,
+            displayText: trimmed,
+            historyEntryID: answer.id,
+            historyEntryRevision: 0,
+            pendingHistoryMutationID: mutation.id
+        )
+        state.lastTranscript = ""
+        state.level = 0
+        OSGLog.keyboardExt.info("AI answer insert length=\(trimmed.count, privacy: .public)")
+        return true
+    }
+
    /// Roll back the last voice insertion when it is still at the caret.
    func undoLastInsertion() {
        if undoLastEditIfPossible() {