From 8f9dfccaa028315fe170471357f3313624ae6259 Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Sun, 24 May 2026 23:35:14 -0700 Subject: [PATCH 1/4] Wire Compose Mode end-to-end with full test coverage The compose-interaction-mode branch put the plumbing in place but left the actual Tab handling stubbed with a "pipeline not installed" disabled state. This makes Compose Mode functional. Tab routing - First Tab in Compose Mode kicks off `runComposeGeneration`, which walks the AX context, builds a `ComposeRequest`, and routes it through the engine (llama only). - Second Tab on a ready draft calls `SuggestionInserter.typeDraft`, which chunks the synthetic Unicode insertion and re-checks focus identity between chunks. - Esc, navigation, and text mutation events cancel any active session or in-flight work. `composeTypingShouldContinue` short-circuits the inserter when the session is cleared or focus shifts. Focus + lifecycle - A focus change while the preview is visible drops the session immediately rather than silently keeping a stale draft alive. - `schedulePrediction` is a no-op in Compose Mode instead of emitting a disabled state; the user typing in their field should not surface a misleading "predictions disabled" message. - Mode flips in either direction reset interaction state through the existing settings-change path; the work controller cancels in-flight generation. Model gate - `RuntimeBootstrapModel.prepareForComposeMode` auto-switches to `tabby-depth-1` if installed and records the user's previous autocomplete model so `restoreAutocompleteModel` puts it back on the flip back. Subscribed from `CotabbyAppEnvironment` so the mode setting drives the runtime without `RuntimeBootstrapModel` needing to know about settings. - Menu bar + Settings show "Compose requires the Open Source engine" and "Compose requires gemma-3n-E4B-it-Q4_K_M.gguf" banners when those conditions aren't met; model picker is disabled while Compose is active so the user cannot fight the auto-switch. Testability - Introduced `ComposeContextCollecting` protocol so the coordinator can be exercised against an in-memory fake AX walker. - Returned type is `ComposeContextCollectionResult` (top-level) instead of `ComposeContextCollector.Result` so fakes don't depend on the concrete collector. Tests (36 new, 310 total, 0 failures) - ComposePromptRendererTests: prompt shape, optional sections, empty placeholder, final instruction position. - ComposeTextNormalizerTests: prompt echo, chat tokens, markdown fences, leading labels, wrapping quotes, typed prefix echo, paragraph preservation. - ComposeRequestFactoryTests: clipping, clipboard gating, sampling minimums, prompt-preview parity, user profile carry-through. - SuggestionCoordinatorComposeTests: first-Tab triggers generation, second-Tab types draft, Esc/navigation/text mutation cancel session, focus change clears session, mode change clears session, empty draft returns to idle, engine failure surfaces failed state. Includes full in-file fake suite for every coordinator collaborator. --- Cotabby.xcodeproj/project.pbxproj | 16 + .../SuggestionCoordinator+Compose.swift | 401 ++++++++++++++ .../SuggestionCoordinator+Input.swift | 21 +- .../SuggestionCoordinator+Prediction.swift | 10 +- .../Coordinators/SuggestionCoordinator.swift | 4 +- Cotabby/App/Core/CotabbyAppEnvironment.swift | 20 + Cotabby/Models/RuntimeBootstrapModel.swift | 39 ++ Cotabby/Models/SuggestionEngineModels.swift | 1 - .../Models/SuggestionSubsystemContracts.swift | 33 ++ .../Context/ComposeContextCollector.swift | 13 +- Cotabby/UI/MenuBarView.swift | 20 + Cotabby/UI/SettingsView.swift | 18 + CotabbyTests/ComposePromptRendererTests.swift | 115 ++++ CotabbyTests/ComposeRequestFactoryTests.swift | 146 +++++ CotabbyTests/ComposeTextNormalizerTests.swift | 143 +++++ .../SuggestionCoordinatorComposeTests.swift | 502 ++++++++++++++++++ 16 files changed, 1483 insertions(+), 19 deletions(-) create mode 100644 Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift create mode 100644 CotabbyTests/ComposePromptRendererTests.swift create mode 100644 CotabbyTests/ComposeRequestFactoryTests.swift create mode 100644 CotabbyTests/ComposeTextNormalizerTests.swift create mode 100644 CotabbyTests/SuggestionCoordinatorComposeTests.swift diff --git a/Cotabby.xcodeproj/project.pbxproj b/Cotabby.xcodeproj/project.pbxproj index 69c93a1..368d2a4 100644 --- a/Cotabby.xcodeproj/project.pbxproj +++ b/Cotabby.xcodeproj/project.pbxproj @@ -36,6 +36,10 @@ F10000012FA0000100EEE001 /* TextDirectionDetectorTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F10000112FA0000100EEE011 /* TextDirectionDetectorTests.swift */; }; G10000012FB0000100FFF001 /* WordCountFormatterTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = G10000112FB0000100FFF011 /* WordCountFormatterTests.swift */; }; H10000012FC0000100GGG001 /* ComposeContextNormalizerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H10000112FC0000100GGG011 /* ComposeContextNormalizerTests.swift */; }; + H20000012FC0000100GGG002 /* ComposePromptRendererTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H20000112FC0000100GGG012 /* ComposePromptRendererTests.swift */; }; + H30000012FC0000100GGG003 /* ComposeTextNormalizerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H30000112FC0000100GGG013 /* ComposeTextNormalizerTests.swift */; }; + H40000012FC0000100GGG004 /* ComposeRequestFactoryTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H40000112FC0000100GGG014 /* ComposeRequestFactoryTests.swift */; }; + H50000012FC0000100GGG005 /* SuggestionCoordinatorComposeTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H50000112FC0000100GGG015 /* SuggestionCoordinatorComposeTests.swift */; }; /* End PBXBuildFile section */ /* Begin PBXContainerItemProxy section */ @@ -75,6 +79,10 @@ F9D35DB9E86506B9FAE1CFE9 /* ModelFileValidatorTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ModelFileValidatorTests.swift; sourceTree = ""; }; G10000112FB0000100FFF011 /* WordCountFormatterTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = WordCountFormatterTests.swift; sourceTree = ""; }; H10000112FC0000100GGG011 /* ComposeContextNormalizerTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ComposeContextNormalizerTests.swift; sourceTree = ""; }; + H20000112FC0000100GGG012 /* ComposePromptRendererTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ComposePromptRendererTests.swift; sourceTree = ""; }; + H30000112FC0000100GGG013 /* ComposeTextNormalizerTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ComposeTextNormalizerTests.swift; sourceTree = ""; }; + H40000112FC0000100GGG014 /* ComposeRequestFactoryTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ComposeRequestFactoryTests.swift; sourceTree = ""; }; + H50000112FC0000100GGG015 /* SuggestionCoordinatorComposeTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionCoordinatorComposeTests.swift; sourceTree = ""; }; /* End PBXFileReference section */ @@ -153,6 +161,10 @@ G10000122FB0000100FFF012 /* ClipboardContentDistillerTests.swift */, G10000112FB0000100FFF011 /* WordCountFormatterTests.swift */, H10000112FC0000100GGG011 /* ComposeContextNormalizerTests.swift */, + H20000112FC0000100GGG012 /* ComposePromptRendererTests.swift */, + H30000112FC0000100GGG013 /* ComposeTextNormalizerTests.swift */, + H40000112FC0000100GGG014 /* ComposeRequestFactoryTests.swift */, + H50000112FC0000100GGG015 /* SuggestionCoordinatorComposeTests.swift */, ); path = CotabbyTests; sourceTree = ""; @@ -297,6 +309,10 @@ G10000022FB0000100FFF002 /* ClipboardContentDistillerTests.swift in Sources */, G10000012FB0000100FFF001 /* WordCountFormatterTests.swift in Sources */, H10000012FC0000100GGG001 /* ComposeContextNormalizerTests.swift in Sources */, + H20000012FC0000100GGG002 /* ComposePromptRendererTests.swift in Sources */, + H30000012FC0000100GGG003 /* ComposeTextNormalizerTests.swift in Sources */, + H40000012FC0000100GGG004 /* ComposeRequestFactoryTests.swift in Sources */, + H50000012FC0000100GGG005 /* SuggestionCoordinatorComposeTests.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift new file mode 100644 index 0000000..2e5f182 --- /dev/null +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift @@ -0,0 +1,401 @@ +import CoreGraphics +import Foundation +import Logging + +/// File overview: +/// Compose Mode entry points for `SuggestionCoordinator`. +/// This is the deliberate two-step Tab flow: first Tab gathers context, generates a full draft, +/// and shows a preview; second Tab types the draft into the focused field. Autocomplete continues +/// to live in the sibling extension files and is untouched by this path. +/// +/// State invariants this file protects: +/// - Only one Compose generation is in flight per `currentWorkID`; later Tabs cancel earlier work. +/// - The `ActiveComposeSession` is never accepted against a focused field whose process/identity +/// has changed since the draft was generated. +/// - Synthetic typing is rearmed against `InputSuppressionController` per chunk via the +/// `SuggestionInserter.typeDraft` contract. +extension SuggestionCoordinator { + // MARK: - Tab Routing + + /// Routes a Tab/Escape/typing event while Compose Mode is active. The autocomplete pipeline + /// is intentionally bypassed so typing in the field does not trigger inline generation. + func handleComposeInputEvent(_ event: CapturedInputEvent) -> Bool { + switch event.kind { + case .acceptance, .fullAcceptance: + if interactionState.activeComposeSession != nil { + return acceptComposeDraft() + } + return startComposeGeneration() + + case .dismissal: + cancelComposeWork(reason: "Compose cancelled by Escape.") + return false + + case .navigation: + // Arrow keys, page navigation, etc. — drop any in-flight draft because the field + // context the user originally asked Tabby to draft against has moved. + if interactionState.activeComposeSession != nil || isAnyComposeWorkInFlight { + cancelComposeWork(reason: "Compose cancelled because the caret moved.") + } + return false + + case .textMutation, .shortcutMutation: + // Typing or paste during preview invalidates the draft. Compose is "ask, review, type"; + // mid-stream edits mean the user has stopped reviewing. + if interactionState.activeComposeSession != nil || isAnyComposeWorkInFlight { + cancelComposeWork(reason: "Compose cancelled because the focused text changed.") + } + return false + + case .other: + return false + } + } + + // MARK: - Generation + + /// First-Tab handler: kick off Compose generation against the current focused field. + /// Returns `true` to consume the Tab so the host app does not receive it. + @discardableResult + func startComposeGeneration() -> Bool { + guard permissionManager.inputMonitoringGranted else { + return passTabThrough(reason: "Input Monitoring permission is required before Cotabby can draft a Compose response.") + } + + let snapshot = focusModel.snapshot + guard case .supported = snapshot.capability, let rawContext = snapshot.context else { + return passTabThrough(reason: snapshot.capability.summary) + } + + if let disabledReason = SuggestionAvailabilityEvaluator.disabledReason( + globallyEnabled: settingsSnapshot.isGloballyEnabled, + disabledAppBundleIdentifiers: settingsSnapshot.disabledAppBundleIdentifiers, + interactionMode: settingsSnapshot.selectedInteractionMode, + inputMonitoringGranted: permissionManager.inputMonitoringGranted, + screenRecordingGranted: permissionManager.screenRecordingGranted, + focusSnapshot: snapshot + ) { + return passTabThrough(reason: disabledReason) + } + + let context = interactionState.materializeContext(from: rawContext) + + // Reuse the debounced-work plumbing with a zero delay so cancellation and stale-work guards + // are identical to the autocomplete path. Compose has no real debounce — Tab is explicit. + let workID = workController.replaceDebouncedWork(delayMilliseconds: 0) { [weak self] workID in + await self?.runComposeGeneration(for: context, workID: workID) + } + latestGenerationNumber = context.generation + state = .generating + logStage( + "compose-generating", + workID: workID, + generation: context.generation, + message: "Gathering Compose context for \(context.elementIdentifier) in \(context.applicationName)." + ) + return true + } + + // swiftlint:disable:next cyclomatic_complexity + private func runComposeGeneration(for context: FocusedInputContext, workID: UInt64) async { + guard workController.isCurrent(workID) else { return } + await awaitCachedGenerationContextResetIfNeeded() + guard workController.isCurrent(workID) else { return } + + let collected: ComposeContextCollectionResult + do { + collected = try await composeContextCollector.collect(for: context) + } catch is CancellationError { + return + } catch { + guard workController.isCurrent(workID) else { return } + await applyComposeFailure(error.localizedDescription, workID: workID) + return + } + guard workController.isCurrent(workID) else { return } + + let clipboardContext: String? = { + guard settingsSnapshot.isClipboardContextEnabled else { return nil } + return clipboardContextProvider.currentContext() + }() + let visualContextSummary = visualContextCoordinator.excerpt(for: context) + + let buildResult = ComposeRequestFactory.buildRequest( + context: context, + settings: settingsSnapshot, + configuration: configuration, + surroundingContext: collected.text, + clipboardContext: clipboardContext, + visualContextSummary: visualContextSummary + ) + latestPromptPreview = buildResult.promptPreview + latestRawModelOutput = nil + let request = buildResult.request + + workController.replaceGenerationWork(for: workID) { [weak self] in + guard let self else { return } + do { + let result = try await self.suggestionEngine.generateCompose(for: request) + guard !Task.isCancelled, self.workController.isCurrent(workID) else { return } + await self.applyComposeResult(result, workID: workID) + } catch SuggestionClientError.cancelled { + return + } catch { + guard self.workController.isCurrent(workID) else { return } + await self.applyComposeFailure(error.localizedDescription, workID: workID) + } + } + } + + /// Stale-result guard mirrors the autocomplete path: we re-read focus before showing anything, + /// and bail if the field's generation or process no longer matches what we asked the model for. + private func applyComposeResult(_ result: ComposeResult, workID: UInt64) async { + guard workController.isCurrent(workID) else { return } + + focusModel.refreshNow() + let snapshot = focusModel.snapshot + + guard case .supported = snapshot.capability, let rawContext = snapshot.context else { + disablePredictions(reason: snapshot.capability.summary) + return + } + let liveContext = interactionState.materializeContext(from: rawContext) + + guard liveContext.generation == result.generation else { + latestRawModelOutput = SuggestionDebugLogger.debugPreview(result.rawText) + logStage( + "compose-stale-drop", + workID: workID, + generation: result.generation, + message: "Dropped stale Compose draft because live generation is \(liveContext.generation).", + rawOutput: result.rawText, + normalizedOutput: result.text + ) + hideOverlay(reason: "Overlay hidden because the focused field changed before the draft was ready.") + return + } + + latestRawModelOutput = SuggestionDebugLogger.debugPreview(result.rawText) + latestLatencyMilliseconds = Int(result.latency * 1000) + latestGenerationNumber = liveContext.generation + + let trimmed = result.text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + interactionState.clearSuggestion() + hideOverlay(reason: "Overlay hidden because Compose returned an empty draft.") + state = .idle + logStage( + "compose-empty", + workID: workID, + generation: result.generation, + message: "Compose draft was empty after normalization.", + rawOutput: result.rawText, + normalizedOutput: result.text + ) + return + } + + let session = interactionState.startComposeSession( + fullText: result.text, + liveContext: liveContext, + latency: result.latency + ) + latestSuggestionPreview = session.fullText + latestFullSuggestionPreview = session.fullText + latestRemainingSuggestionPreview = session.fullText + latestAcceptedCharacterCount = 0 + latestRemainingCharacterCount = session.fullText.count + state = .ready(text: session.fullText, latency: session.latency) + + presentComposePreview( + text: session.fullText, + at: liveContext.caretRect, + inputFrameRect: liveContext.inputFrameRect, + caretQuality: liveContext.caretQuality, + observedCharWidth: liveContext.observedCharWidth, + isRightToLeft: TextDirectionDetector.isRightToLeft(liveContext.precedingText) + ) + logStage( + "compose-ready", + workID: workID, + generation: result.generation, + message: "Compose draft ready for review.", + rawOutput: result.rawText, + normalizedOutput: result.text + ) + } + + private func applyComposeFailure(_ message: String, workID: UInt64) async { + guard workController.isCurrent(workID) else { return } + interactionState.clearSuggestion() + hideOverlay(reason: "Overlay hidden because Compose generation failed.") + state = .failed(message) + logStage( + "compose-failed", + workID: workID, + generation: latestGenerationNumber, + message: message + ) + } + + // MARK: - Acceptance + + /// Second-Tab handler: type the active Compose draft into the focused field via + /// `SuggestionInserter.typeDraft`. Each chunk re-checks focus identity before posting. + @discardableResult + func acceptComposeDraft() -> Bool { + guard let session = interactionState.activeComposeSession else { + return passTabThrough(reason: "Key passed through because no Compose draft was ready.") + } + + focusModel.refreshNow() + let snapshot = focusModel.snapshot + guard case .supported = snapshot.capability, let rawContext = snapshot.context else { + cancelComposeWork(reason: snapshot.capability.summary) + return passTabThrough(reason: snapshot.capability.summary) + } + let liveContext = interactionState.materializeContext(from: rawContext) + guard liveContext.identity == session.baseContext.identity, + liveContext.processIdentifier == session.baseContext.processIdentifier else { + cancelComposeWork(reason: "Compose cancelled because the focused field changed before typing began.") + return passTabThrough(reason: "Key passed through because the focused field changed.") + } + guard liveContext.selection.length == 0 else { + cancelComposeWork(reason: "Compose cancelled because text is selected.") + return passTabThrough(reason: "Key passed through because text is selected.") + } + + state = .typing + recordAcceptedWords(from: session.fullText) + logStage( + "compose-accepting", + workID: currentWorkID, + generation: liveContext.generation, + message: "Typing Compose draft (\(session.fullText.count) characters) into \(liveContext.applicationName).", + normalizedOutput: session.fullText + ) + + let typingSession = session + Task { @MainActor [weak self] in + guard let self else { return } + let didType = await self.suggestionInserter.typeDraft(typingSession.fullText) { [weak self] in + self?.composeTypingShouldContinue(matching: typingSession) ?? false + } + + // The session may have already been cleared by a cancellation path (focus change, + // mode change, Esc). If so, don't reset state again. + guard let active = self.interactionState.activeComposeSession, active == typingSession else { + return + } + + if didType { + self.interactionState.clearComposeSession(typingSession) + self.hideOverlay(reason: "Overlay hidden after Compose draft was typed into the field.") + self.state = .idle + self.latestAcceptanceAction = "Compose draft typed into the field." + self.logStage( + "compose-typed", + workID: self.currentWorkID, + generation: typingSession.baseContext.generation, + message: "Compose draft fully typed into the focused field.", + normalizedOutput: typingSession.fullText + ) + } else { + let message = self.suggestionInserter.lastErrorMessage + ?? "Compose typing stopped before the draft was complete." + self.interactionState.clearComposeSession(typingSession) + self.hideOverlay(reason: "Overlay hidden because Compose typing did not complete.") + self.state = .failed(message) + self.logStage( + "compose-type-aborted", + workID: self.currentWorkID, + generation: typingSession.baseContext.generation, + message: message, + normalizedOutput: typingSession.fullText + ) + } + } + return true + } + + /// Focus-identity guard checked between every synthetic key chunk. Bailing here lets the + /// inserter stop posting mid-draft when the user switches apps or fields. + private func composeTypingShouldContinue(matching session: ActiveComposeSession) -> Bool { + guard let active = interactionState.activeComposeSession, active == session else { + return false + } + let snapshot = focusModel.snapshot + guard case .supported = snapshot.capability, let rawContext = snapshot.context else { + return false + } + return rawContext.processIdentifier == session.baseContext.processIdentifier + && rawContext.elementIdentifier == session.baseContext.elementIdentifier + && rawContext.focusChangeSequence == session.baseContext.focusChangeSequence + } + + // MARK: - Cancellation + + /// Cancels any in-flight Compose work and clears the active session. Safe to call when nothing + /// is in flight; the underlying controllers are no-op in that case. + func cancelComposeWork(reason: String) { + let hadActiveSession = interactionState.activeComposeSession != nil + let hadInflightWork = isAnyComposeWorkInFlight + guard hadActiveSession || hadInflightWork else { return } + + workController.cancelAll() + if let session = interactionState.activeComposeSession { + interactionState.clearComposeSession(session) + } else { + interactionState.clearSuggestion() + } + hideOverlay(reason: reason) + state = .idle + logStage( + "compose-cancelled", + workID: currentWorkID, + generation: latestGenerationNumber, + message: reason + ) + } + + /// True when a Compose generation or typing task could still emit results. We treat any state + /// other than `.idle` / `.disabled` as "could still emit" because the work controller only + /// surfaces task identity, not status. + var isAnyComposeWorkInFlight: Bool { + switch state { + case .generating, .typing: + return true + case .idle, .disabled, .debouncing, .ready, .failed: + return false + } + } + + // MARK: - Overlay + + // swiftlint:disable function_parameter_count + /// Sibling of `presentOverlay` for the multiline Compose preview surface. + private func presentComposePreview( + text: String, + at caretRect: CGRect, + inputFrameRect: CGRect?, + caretQuality: CaretGeometryQuality, + observedCharWidth: CGFloat?, + isRightToLeft: Bool + ) { + let geometry = SuggestionOverlayGeometry( + caretRect: caretRect, + inputFrameRect: inputFrameRect, + caretQuality: caretQuality, + observedCharWidth: observedCharWidth, + isRightToLeft: isRightToLeft + ) + if let message = overlayPresenter.presentComposePreview( + text: text, + geometry: geometry, + previousState: overlayState + ) { + latestOverlayMessage = message + } + } + // swiftlint:enable function_parameter_count +} diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift index 5878db8..8565a9f 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift @@ -67,7 +67,19 @@ extension SuggestionCoordinator { state = .idle } - if interactionState.activeSession != nil { + // Compose's preview is field-scoped: if the user clicked a different field or app while + // the draft was visible, drop the session immediately. Same-field updates are a no-op + // because Compose does not reconcile against typing the way autocomplete does. + if let composeSession = interactionState.activeComposeSession { + let processChanged = composeSession.baseContext.processIdentifier != focusedContext.processIdentifier + let elementChanged = composeSession.baseContext.elementIdentifier != focusedContext.elementIdentifier + if processChanged || elementChanged { + cancelComposeWork(reason: "Compose cancelled because the focused field changed.") + } + return + } + + if interactionState.activeAutocompleteSession != nil { reconcileActiveSession(with: snapshot) return } @@ -98,6 +110,13 @@ extension SuggestionCoordinator { return false } + // Compose has a deliberate two-step Tab flow and never participates in inline-autocomplete + // session reconciliation or per-keystroke prediction debouncing. Fork the event handling + // here so the autocomplete branch below can keep assuming "this is autocomplete state". + if settingsSnapshot.selectedInteractionMode == .compose { + return handleComposeInputEvent(event) + } + if event.kind == .acceptance { return acceptCurrentSuggestion() } diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift index 4e1536f..3b5ff81 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift @@ -8,8 +8,10 @@ extension SuggestionCoordinator { // MARK: - Prediction Pipeline func schedulePrediction() { + // Compose Mode is Tab-driven, not typing-driven. Returning here keeps Tabby quiet while + // the user types into the field without surfacing a misleading "disabled" state — the + // explicit Compose path lives in `SuggestionCoordinator+Compose.swift`. guard settingsSnapshot.selectedInteractionMode == .autocomplete else { - disablePredictionsPreservingVisualContext(reason: composeModePendingReason) return } @@ -38,8 +40,9 @@ extension SuggestionCoordinator { /// Refreshes focus after debounce, materializes a stable context, and starts generation. func generateFromCurrentFocus(workID: UInt64) async { + // Mode may have flipped to Compose while this debounced task was sleeping. Bail silently + // — the Compose path is driven from `handleComposeInputEvent`, not from this debounce. guard settingsSnapshot.selectedInteractionMode == .autocomplete else { - disablePredictionsPreservingVisualContext(reason: composeModePendingReason) return } @@ -453,7 +456,4 @@ extension SuggestionCoordinator { schedulePrediction() } - private var composeModePendingReason: String { - "Compose Mode is selected. Draft generation will be enabled after the Compose request pipeline is installed." - } } diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator.swift b/Cotabby/App/Coordinators/SuggestionCoordinator.swift index 84fea79..7551cdc 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator.swift @@ -44,7 +44,7 @@ final class SuggestionCoordinator: ObservableObject { let clipboardContextProvider: any ClipboardContextProviding let clipboardRelevanceFilter: any ClipboardRelevanceFiltering let visualContextCoordinator: any VisualContextCoordinating - let composeContextCollector: ComposeContextCollector + let composeContextCollector: any ComposeContextCollecting let interactionState: SuggestionInteractionState let workController: SuggestionWorkController let configuration: SuggestionConfiguration @@ -74,7 +74,7 @@ final class SuggestionCoordinator: ObservableObject { clipboardContextProvider: any ClipboardContextProviding, clipboardRelevanceFilter: any ClipboardRelevanceFiltering, visualContextCoordinator: any VisualContextCoordinating, - composeContextCollector: ComposeContextCollector, + composeContextCollector: any ComposeContextCollecting, interactionState: SuggestionInteractionState, workController: SuggestionWorkController, configuration: SuggestionConfiguration, diff --git a/Cotabby/App/Core/CotabbyAppEnvironment.swift b/Cotabby/App/Core/CotabbyAppEnvironment.swift index 282fff3..0ecca66 100644 --- a/Cotabby/App/Core/CotabbyAppEnvironment.swift +++ b/Cotabby/App/Core/CotabbyAppEnvironment.swift @@ -184,6 +184,26 @@ final class CotabbyAppEnvironment { } .store(in: &cancellables) + // Compose Mode requires `tabby-depth-1`. Auto-switch on entry / restore on exit lives here + // (not on `RuntimeBootstrapModel`) because the runtime model intentionally does not know + // about interaction modes. `dropFirst()` ignores the initial publish; `removeDuplicates` + // avoids redundant work when other settings change. + suggestionSettings.$selectedInteractionMode + .removeDuplicates() + .dropFirst() + .sink { [weak runtimeModel] mode in + guard let runtimeModel else { return } + Task { @MainActor in + switch mode { + case .compose: + await runtimeModel.prepareForComposeMode() + case .autocomplete: + await runtimeModel.restoreAutocompleteModel() + } + } + } + .store(in: &cancellables) + // Key code changes reach InputMonitor through closures that read from the model // at event time (set above), so no Combine subscription is needed here. } diff --git a/Cotabby/Models/RuntimeBootstrapModel.swift b/Cotabby/Models/RuntimeBootstrapModel.swift index 3e18a56..c6e40ac 100644 --- a/Cotabby/Models/RuntimeBootstrapModel.swift +++ b/Cotabby/Models/RuntimeBootstrapModel.swift @@ -19,6 +19,10 @@ final class RuntimeBootstrapModel: ObservableObject { private let userDefaults: UserDefaults private var cancellables = Set() private var runtimeTask: Task? + /// The model the user had selected before Compose Mode auto-switched to `tabby-depth-1`. + /// Persisted only in memory because the Compose round-trip is a within-session feature; if the + /// app relaunches in Compose Mode we re-select the required model on the next entry anyway. + private var preComposeAutocompleteModelFilename: String? /// Called immediately before the runtime begins switching models so suggestion state can reset. var onWillReloadModel: (() -> Void)? @@ -131,6 +135,41 @@ final class RuntimeBootstrapModel: ObservableObject { await runtimeTask?.value } + /// Whether the Compose Mode required local model is currently discovered. + /// Read by UI so it can offer a clear "Compose needs tabby-depth-1" message without each view + /// duplicating the catalog lookup. + var isComposeRequiredModelInstalled: Bool { + availableModels.contains { $0.filename == RuntimeModelCatalog.composeRequiredFilename } + } + + /// Switches the runtime to `tabby-depth-1` for Compose Mode, remembering the prior selection so + /// returning to Autocomplete restores it. No-op when the required model is missing, when it is + /// already selected, or when a runtime task is in flight (the next entry will retry). + func prepareForComposeMode() async { + let required = RuntimeModelCatalog.composeRequiredFilename + guard isComposeRequiredModelInstalled else { + TabbyLogger.runtime.info("Compose Mode required model \(required) is not installed; leaving selection unchanged.") + return + } + guard selectedModelFilename != required else { return } + guard runtimeTask == nil else { return } + + preComposeAutocompleteModelFilename = selectedModelFilename + await selectModel(required) + } + + /// Restores the user's pre-Compose model selection. No-op when no prior selection was saved + /// or when the saved model is no longer available. + func restoreAutocompleteModel() async { + guard let previous = preComposeAutocompleteModelFilename else { return } + preComposeAutocompleteModelFilename = nil + guard availableModels.contains(where: { $0.filename == previous }) else { return } + guard selectedModelFilename != previous else { return } + guard runtimeTask == nil else { return } + + await selectModel(previous) + } + /// Cancels pending startup work and forwards shutdown to the underlying runtime manager. func stop() { runtimeTask?.cancel() diff --git a/Cotabby/Models/SuggestionEngineModels.swift b/Cotabby/Models/SuggestionEngineModels.swift index e643fa7..831ce75 100644 --- a/Cotabby/Models/SuggestionEngineModels.swift +++ b/Cotabby/Models/SuggestionEngineModels.swift @@ -106,7 +106,6 @@ struct SuggestionSettingsSnapshot: Equatable, Sendable { let focusPollIntervalMilliseconds: Int let isMultiLineEnabled: Bool - // swiftlint:disable:next function_parameter_count init( isGloballyEnabled: Bool, disabledAppBundleIdentifiers: Set, diff --git a/Cotabby/Models/SuggestionSubsystemContracts.swift b/Cotabby/Models/SuggestionSubsystemContracts.swift index f2ef2a4..dcc2965 100644 --- a/Cotabby/Models/SuggestionSubsystemContracts.swift +++ b/Cotabby/Models/SuggestionSubsystemContracts.swift @@ -102,3 +102,36 @@ protocol VisualContextCoordinating: AnyObject { func cancel(resetState: Bool) func excerpt(for context: FocusedInputContext) -> String? } + +/// Behavior-shaped contract for Compose Mode AX context collection. +/// +/// The coordinator only needs "collect a normalized surrounding context for this focused field"; +/// keeping the contract narrow lets tests substitute a deterministic fake without standing up the +/// real AX tree, while production code still uses the bounded DFS in `ComposeContextCollector`. +@MainActor +protocol ComposeContextCollecting: AnyObject { + func collect(for context: FocusedInputContext) async throws -> ComposeContextCollectionResult +} + +/// What a Compose context collector returns. +/// +/// This sits next to the protocol (not on the concrete collector) so test fakes can construct +/// results without depending on the real collector's nested types. +struct ComposeContextCollectionResult: Equatable, Sendable { + let text: String + let visitedNodeCount: Int + let retainedTextCount: Int + let droppedTextCount: Int + + init( + text: String, + visitedNodeCount: Int = 0, + retainedTextCount: Int = 0, + droppedTextCount: Int = 0 + ) { + self.text = text + self.visitedNodeCount = visitedNodeCount + self.retainedTextCount = retainedTextCount + self.droppedTextCount = droppedTextCount + } +} diff --git a/Cotabby/Services/Context/ComposeContextCollector.swift b/Cotabby/Services/Context/ComposeContextCollector.swift index f99705e..78e8f53 100644 --- a/Cotabby/Services/Context/ComposeContextCollector.swift +++ b/Cotabby/Services/Context/ComposeContextCollector.swift @@ -7,7 +7,7 @@ import Foundation /// side-effectful macOS boundary with app-specific failure modes. The coordinator should ask for a /// bounded, normalized context string; it should not own traversal budgets or Core Foundation reads. @MainActor -final class ComposeContextCollector { +final class ComposeContextCollector: ComposeContextCollecting { struct Limits: Equatable, Sendable { let maxAncestorDepth: Int let maxDFSDepth: Int @@ -24,13 +24,6 @@ final class ComposeContextCollector { ) } - struct Result: Equatable, Sendable { - let text: String - let visitedNodeCount: Int - let retainedTextCount: Int - let droppedTextCount: Int - } - enum CollectionError: LocalizedError, Equatable { case noFocusedElement case staleFocus @@ -74,7 +67,7 @@ final class ComposeContextCollector { self.limits = limits } - func collect(for context: FocusedInputContext) async throws -> Result { + func collect(for context: FocusedInputContext) async throws -> ComposeContextCollectionResult { try Task.checkCancellation() guard let focusedElement = AXHelper.focusedElement() else { @@ -147,7 +140,7 @@ final class ComposeContextCollector { limits: limits.normalizerLimits ) - return Result( + return ComposeContextCollectionResult( text: normalizedText, visitedNodeCount: visitedNodeCount, retainedTextCount: retainedTextCount, diff --git a/Cotabby/UI/MenuBarView.swift b/Cotabby/UI/MenuBarView.swift index 9650440..e07e825 100644 --- a/Cotabby/UI/MenuBarView.swift +++ b/Cotabby/UI/MenuBarView.swift @@ -119,6 +119,21 @@ struct MenuBarView: View { .foregroundStyle(.orange) } + if suggestionSettings.selectedInteractionMode == .compose, + suggestionSettings.selectedEngine != .llamaOpenSource { + Text("Compose requires the Open Source engine.") + .font(.caption) + .foregroundStyle(.orange) + } + + if suggestionSettings.selectedInteractionMode == .compose, + suggestionSettings.selectedEngine == .llamaOpenSource, + !runtimeModel.isComposeRequiredModelInstalled { + Text("Compose requires \(RuntimeModelCatalog.composeRequiredFilename). Add it to the models folder.") + .font(.caption) + .foregroundStyle(.orange) + } + if suggestionSettings.selectedEngine.supportsLocalModelManagement { modelRow } @@ -338,6 +353,11 @@ struct MenuBarView: View { } private var runtimePickerDisabled: Bool { + // Compose Mode pins the runtime to `tabby-depth-1`, so user-driven model selection is + // disabled while Compose is active to prevent half-supported configurations. + if suggestionSettings.selectedInteractionMode == .compose { + return true + } switch runtimeModel.state { case .starting, .loading: return true diff --git a/Cotabby/UI/SettingsView.swift b/Cotabby/UI/SettingsView.swift index a42a576..fa487c0 100644 --- a/Cotabby/UI/SettingsView.swift +++ b/Cotabby/UI/SettingsView.swift @@ -175,6 +175,24 @@ struct SettingsView: View { } } + if suggestionSettings.selectedInteractionMode == .compose, + suggestionSettings.selectedEngine != .llamaOpenSource { + Text("Compose requires the Open Source engine. Switch to it to draft full responses.") + .font(.caption) + .foregroundStyle(.orange) + } + + if suggestionSettings.selectedInteractionMode == .compose, + suggestionSettings.selectedEngine == .llamaOpenSource, + !runtimeModel.isComposeRequiredModelInstalled { + Text( + "Compose requires \(RuntimeModelCatalog.composeRequiredFilename). " + + "Add it to your models folder to enable draft generation." + ) + .font(.caption) + .foregroundStyle(.orange) + } + Picker("Length", selection: selectedWordCountPresetBinding) { ForEach(SuggestionWordCountPreset.allCases) { preset in Text(preset.displayLabel) diff --git a/CotabbyTests/ComposePromptRendererTests.swift b/CotabbyTests/ComposePromptRendererTests.swift new file mode 100644 index 0000000..2e8cf86 --- /dev/null +++ b/CotabbyTests/ComposePromptRendererTests.swift @@ -0,0 +1,115 @@ +import XCTest +@testable import Cotabby + +/// Tests `ComposePromptRenderer` — the pure mapping from a `ComposeRequest` to the prompt string +/// fed into the local llama runtime. +/// +/// These tests lock down which sections appear in the prompt and how the renderer handles empty +/// or optional fields. The prompt shape is the contract the model relies on, so silent reordering +/// here would degrade Compose Mode quality without any compile-time warning. +final class ComposePromptRendererTests: XCTestCase { + func test_prompt_includesTaskInstructionsAndFinalDirective() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest()) + + XCTAssertTrue(prompt.contains("This is Compose Mode, not autocomplete and not chat.")) + XCTAssertTrue(prompt.contains("Return only the final typeable draft.")) + // The "Write the full draft now." instruction must remain at the tail so the model is + // primed to produce the draft directly after the surrounding context block. + XCTAssertTrue(prompt.hasSuffix("Final instruction:\nWrite the full draft now.")) + } + + func test_prompt_includesApplicationAndTypedPrefix() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest( + applicationName: "GitHub", + typedPrefix: "Thanks for the review — " + )) + + XCTAssertTrue(prompt.contains("App:\nGitHub")) + XCTAssertTrue(prompt.contains("Text already typed in the focused field:\nThanks for the review — ")) + } + + func test_prompt_rendersEmptyPlaceholderWhenTypedPrefixIsBlank() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest(typedPrefix: " \n ")) + + XCTAssertTrue(prompt.contains("Text already typed in the focused field:\n(empty)")) + } + + func test_prompt_includesUserNameAndUserTagsWhenProvided() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest( + userName: "Jacob", + userTags: ["engineer", "macOS"] + )) + + XCTAssertTrue(prompt.contains("User name:\nJacob")) + XCTAssertTrue(prompt.contains("User profile tags:\nengineer, macOS")) + } + + func test_prompt_omitsUserNameWhenEmptyAfterTrimming() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest(userName: " ")) + + XCTAssertFalse(prompt.contains("User name:")) + } + + func test_prompt_omitsUserTagsWhenEmpty() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest(userTags: [])) + + XCTAssertFalse(prompt.contains("User profile tags:")) + } + + func test_prompt_includesTrailingTextOnlyWhenNotBlank() { + let withTrailing = ComposePromptRenderer.prompt(for: composeRequest(trailingText: "...rest of paragraph")) + XCTAssertTrue(withTrailing.contains("Text after the caret:\n...rest of paragraph")) + + let withoutTrailing = ComposePromptRenderer.prompt(for: composeRequest(trailingText: "")) + XCTAssertFalse(withoutTrailing.contains("Text after the caret:")) + } + + func test_prompt_includesClipboardAndVisualContextWhenProvided() { + let prompt = ComposePromptRenderer.prompt(for: composeRequest( + visualContextSummary: "PAGE_HEADER", + clipboardContext: "COPIED_LINK" + )) + + XCTAssertTrue(prompt.contains("Clipboard context:\nCOPIED_LINK")) + XCTAssertTrue(prompt.contains("Visual context summary:\nPAGE_HEADER")) + } + + func test_prompt_includesSurroundingContextEvenWhenEmpty() { + // The surrounding-context section is always present so the model sees consistent prompt + // shape across runs; empty content is rendered as "(empty)". + let prompt = ComposePromptRenderer.prompt(for: composeRequest(surroundingContext: " ")) + + XCTAssertTrue(prompt.contains("Relevant surrounding context:\n(empty)")) + } + + private func composeRequest( + applicationName: String = "TestApp", + typedPrefix: String = "Hello", + trailingText: String = "", + surroundingContext: String = "Some surrounding context.", + visualContextSummary: String? = nil, + clipboardContext: String? = nil, + userName: String? = nil, + userTags: [String]? = nil + ) -> ComposeRequest { + ComposeRequest( + context: CotabbyTestFixtures.focusedInputContext(applicationName: applicationName), + typedPrefix: typedPrefix, + trailingText: trailingText, + surroundingContext: surroundingContext, + visualContextSummary: visualContextSummary, + clipboardContext: clipboardContext, + applicationName: applicationName, + generation: 1, + maxPredictionTokens: 256, + temperature: 0.4, + topK: 40, + topP: 0.9, + minP: 0.05, + repetitionPenalty: 1.1, + randomSeed: nil, + userName: userName, + userTags: userTags + ) + } +} diff --git a/CotabbyTests/ComposeRequestFactoryTests.swift b/CotabbyTests/ComposeRequestFactoryTests.swift new file mode 100644 index 0000000..09b6484 --- /dev/null +++ b/CotabbyTests/ComposeRequestFactoryTests.swift @@ -0,0 +1,146 @@ +import XCTest +@testable import Cotabby + +/// Tests `ComposeRequestFactory` — the pure builder that turns a focused-input context plus settings +/// into a `ComposeRequest` with sane clipping and sampling defaults. +/// +/// Compose Mode uses larger token budgets and looser sampling than autocomplete; these tests lock +/// in those minimums so future tuning cannot accidentally make Compose behave like autocomplete. +final class ComposeRequestFactoryTests: XCTestCase { + func test_buildRequest_clipsLongTypedPrefixWithEllipsis() { + let oversizedPrefix = String(repeating: "a", count: 5_000) + let context = CotabbyTestFixtures.focusedInputContext(precedingText: oversizedPrefix) + + let result = ComposeRequestFactory.buildRequest( + context: context, + settings: settings(), + configuration: .standard, + surroundingContext: "", + clipboardContext: nil + ) + + XCTAssertLessThanOrEqual(result.request.typedPrefix.count, 4_000) + XCTAssertTrue(result.request.typedPrefix.hasSuffix("...")) + } + + func test_buildRequest_clipsTrailingTextAndSurroundingContext() { + let trailing = String(repeating: "b", count: 2_000) + let surrounding = String(repeating: "c", count: 12_000) + let context = CotabbyTestFixtures.focusedInputContext( + precedingText: "Hi", + trailingText: trailing + ) + + let result = ComposeRequestFactory.buildRequest( + context: context, + settings: settings(), + configuration: .standard, + surroundingContext: surrounding, + clipboardContext: nil + ) + + XCTAssertLessThanOrEqual(result.request.trailingText.count, 1_000) + XCTAssertLessThanOrEqual(result.request.surroundingContext.count, 8_000) + } + + func test_buildRequest_omitsClipboardWhenDisabledEvenIfProvided() { + let result = ComposeRequestFactory.buildRequest( + context: CotabbyTestFixtures.focusedInputContext(), + settings: settings(isClipboardContextEnabled: false), + configuration: .standard, + surroundingContext: "Context.", + clipboardContext: "Should be ignored." + ) + + XCTAssertNil(result.request.clipboardContext) + } + + func test_buildRequest_includesClipboardWhenEnabled() { + let result = ComposeRequestFactory.buildRequest( + context: CotabbyTestFixtures.focusedInputContext(), + settings: settings(isClipboardContextEnabled: true), + configuration: .standard, + surroundingContext: "Context.", + clipboardContext: "https://example.com" + ) + + XCTAssertEqual(result.request.clipboardContext, "https://example.com") + } + + func test_buildRequest_appliesComposeSamplingMinimumsAboveAutocomplete() { + // Autocomplete uses temperature 0.1 and 8 tokens. Compose should bump both up because a + // multi-sentence draft needs more headroom than an inline tail. + let result = ComposeRequestFactory.buildRequest( + context: CotabbyTestFixtures.focusedInputContext(), + settings: settings(), + configuration: .standard, + surroundingContext: "", + clipboardContext: nil + ) + + XCTAssertGreaterThanOrEqual(result.request.maxPredictionTokens, 256) + XCTAssertGreaterThanOrEqual(result.request.temperature, 0.35) + XCTAssertGreaterThanOrEqual(result.request.topK, 40) + XCTAssertGreaterThanOrEqual(result.request.topP, 0.9) + XCTAssertGreaterThanOrEqual(result.request.repetitionPenalty, 1.08) + } + + func test_buildRequest_carriesUserNameAndUserTagsThroughToTheRequest() { + let snapshot = SuggestionSettingsSnapshot( + isGloballyEnabled: true, + disabledAppBundleIdentifiers: [], + selectedInteractionMode: .compose, + selectedEngine: .llamaOpenSource, + selectedWordCountPreset: .sevenToTwelve, + isClipboardContextEnabled: true, + userName: "Jacob", + userTags: ["engineer", "macOS"], + debounceMilliseconds: 50, + focusPollIntervalMilliseconds: 50, + isMultiLineEnabled: false + ) + + let result = ComposeRequestFactory.buildRequest( + context: CotabbyTestFixtures.focusedInputContext(), + settings: snapshot, + configuration: .standard, + surroundingContext: "Context", + clipboardContext: nil + ) + + XCTAssertEqual(result.request.userName, "Jacob") + XCTAssertEqual(result.request.userTags, ["engineer", "macOS"]) + } + + func test_buildRequest_includesPromptPreviewMatchingRenderer() { + let result = ComposeRequestFactory.buildRequest( + context: CotabbyTestFixtures.focusedInputContext(applicationName: "GitHub"), + settings: settings(), + configuration: .standard, + surroundingContext: "Surrounding context.", + clipboardContext: nil + ) + + XCTAssertEqual(result.promptPreview, ComposePromptRenderer.prompt(for: result.request)) + } + + private func settings( + isClipboardContextEnabled: Bool = true, + userName: String = "", + userTags: [String] = [] + ) -> SuggestionSettingsSnapshot { + SuggestionSettingsSnapshot( + isGloballyEnabled: true, + disabledAppBundleIdentifiers: [], + selectedInteractionMode: .compose, + selectedEngine: .llamaOpenSource, + selectedWordCountPreset: .sevenToTwelve, + isClipboardContextEnabled: isClipboardContextEnabled, + userName: userName, + userTags: userTags, + debounceMilliseconds: 50, + focusPollIntervalMilliseconds: 50, + isMultiLineEnabled: false + ) + } +} diff --git a/CotabbyTests/ComposeTextNormalizerTests.swift b/CotabbyTests/ComposeTextNormalizerTests.swift new file mode 100644 index 0000000..c79122c --- /dev/null +++ b/CotabbyTests/ComposeTextNormalizerTests.swift @@ -0,0 +1,143 @@ +import XCTest +@testable import Cotabby + +/// Tests `ComposeTextNormalizer` — the last-mile cleanup that turns raw llama output into the text +/// Cotabby actually types into the focused field. +/// +/// These tests exist because the existing `SuggestionTextNormalizer` aggressively truncates to a +/// short inline tail; Compose has a different contract (preserve paragraphs, strip wrappers). +final class ComposeTextNormalizerTests: XCTestCase { + func test_normalize_stripsLeadingPromptEcho() { + let prompt = "PROMPT_SECTION\nFinal instruction:\nWrite the full draft now." + let raw = prompt + "\nHello team,\n\nThanks for the review." + + let normalized = ComposeTextNormalizer.normalize(raw, prompt: prompt, request: request()) + + XCTAssertEqual(normalized, "Hello team,\n\nThanks for the review.") + } + + func test_normalize_stripsChatTemplateSpecialTokens() { + let normalized = ComposeTextNormalizer.normalize( + "<|im_start|>Hello.<|im_end|>", + prompt: "", + request: request() + ) + + XCTAssertEqual(normalized, "Hello.") + } + + func test_normalize_stripsMarkdownFences() { + let raw = """ + ``` + Hello team, + + Thanks for the review. + ``` + """ + + let normalized = ComposeTextNormalizer.normalize(raw, prompt: "", request: request()) + + XCTAssertEqual(normalized, "Hello team,\n\nThanks for the review.") + } + + func test_normalize_stripsLeadingLabelsCaseInsensitively() { + let cases: [(String, String)] = [ + ("Final answer: Hello.", "Hello."), + ("Draft: A short draft.", "A short draft."), + ("comment: lowercase label", "lowercase label"), + ("Reply: With a label.", "With a label.") + ] + + for (raw, expected) in cases { + let normalized = ComposeTextNormalizer.normalize(raw, prompt: "", request: request()) + XCTAssertEqual(normalized, expected, "expected \(expected) for \(raw)") + } + } + + func test_normalize_stripsWrappingQuotesOnly() { + let doubled = ComposeTextNormalizer.normalize( + "\"Hello team.\"", + prompt: "", + request: request() + ) + XCTAssertEqual(doubled, "Hello team.") + + // Inline quotes in the middle of a draft must not be stripped — they are content. + let preserved = ComposeTextNormalizer.normalize( + "Hello \"team\" again.", + prompt: "", + request: request() + ) + XCTAssertEqual(preserved, "Hello \"team\" again.") + } + + func test_normalize_stripsTypedPrefixEchoWhenPresent() { + let typedPrefix = "Thanks for the review — " + let raw = "Thanks for the review — this looks good to ship." + + let normalized = ComposeTextNormalizer.normalize( + raw, + prompt: "", + request: request(typedPrefix: typedPrefix) + ) + + XCTAssertEqual(normalized, "this looks good to ship.") + } + + func test_normalize_preservesParagraphBoundariesAndCollapsesRunsOfBlankLines() { + let raw = """ + Paragraph one. + + + + Paragraph two. + + + Paragraph three. + """ + + let normalized = ComposeTextNormalizer.normalize(raw, prompt: "", request: request()) + + XCTAssertEqual(normalized, "Paragraph one.\n\nParagraph two.\n\nParagraph three.") + } + + func test_normalize_trimsLeadingAndTrailingNewlinesWithoutTouchingContent() { + let normalized = ComposeTextNormalizer.normalize( + "\n\n Paragraph one.\n\nParagraph two.\n\n", + prompt: "", + request: request() + ) + + XCTAssertEqual(normalized, "Paragraph one.\n\nParagraph two.") + } + + func test_normalize_isNoOpWhenPromptEchoIsAbsent() { + let raw = "Just a clean draft." + + let normalized = ComposeTextNormalizer.normalize(raw, prompt: "DIFFERENT_PROMPT", request: request()) + + XCTAssertEqual(normalized, "Just a clean draft.") + } + + private func request(typedPrefix: String = "") -> ComposeRequest { + ComposeRequest( + context: CotabbyTestFixtures.focusedInputContext(), + typedPrefix: typedPrefix, + trailingText: "", + surroundingContext: "", + visualContextSummary: nil, + clipboardContext: nil, + applicationName: "TestApp", + generation: 1, + maxPredictionTokens: 256, + temperature: 0.4, + topK: 40, + topP: 0.9, + minP: 0.05, + repetitionPenalty: 1.1, + randomSeed: nil, + userName: nil, + userTags: nil + ) + } +} diff --git a/CotabbyTests/SuggestionCoordinatorComposeTests.swift b/CotabbyTests/SuggestionCoordinatorComposeTests.swift new file mode 100644 index 0000000..d4c6c0c --- /dev/null +++ b/CotabbyTests/SuggestionCoordinatorComposeTests.swift @@ -0,0 +1,502 @@ +import Combine +import CoreGraphics +import Foundation +import XCTest +@testable import Cotabby + +/// Tests `SuggestionCoordinator+Compose` end to end with fakes for every collaborator. +/// +/// Most coordinator tests in the autocomplete path live inside `SuggestionInteractionState` +/// helpers; Compose's orchestration is new enough that those state helpers cannot prove the +/// behavior reviewers actually care about: first Tab generates, second Tab types, focus changes +/// cancel mid-flight. The fakes below are all kept in this file so the contract under test stays +/// readable as one unit. +@MainActor +final class SuggestionCoordinatorComposeTests: XCTestCase { + private static var retainedCoordinators: [SuggestionCoordinator] = [] + + override func tearDown() async throws { + Self.retainedCoordinators.removeAll() + try await super.tearDown() + } + + // MARK: - First Tab: generation + + func test_firstTab_inComposeMode_callsGenerateComposeAndShowsPreview() async { + let generateExpectation = expectation(description: "generateCompose called") + let showPreviewExpectation = expectation(description: "showComposePreview called") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(composeResult(text: "Hello team.")), + onGenerateCalled: { generateExpectation.fulfill() }, + onShowComposePreview: { showPreviewExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [generateExpectation, showPreviewExpectation], timeout: 1.0) + + XCTAssertEqual(env.engine.composeCallCount, 1) + XCTAssertEqual(env.overlayController.composePreviewText, "Hello team.") + XCTAssertNotNil(env.coordinator.interactionState.activeComposeSession) + } + + func test_firstTab_inComposeMode_returnsTrueToConsumeTab() { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(composeResult(text: "draft")) + ) + + let consumed = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + + XCTAssertTrue(consumed) + } + + func test_firstTab_inAutocompleteMode_doesNotInvokeComposePath() async { + let env = makeEnvironment( + mode: .autocomplete, + engineBehavior: .success(composeResult(text: "should not run")) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + // Yield to let any spurious Task progress. + await Task.yield() + await Task.yield() + + XCTAssertEqual(env.engine.composeCallCount, 0) + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + // MARK: - Second Tab: acceptance + + func test_secondTab_withActiveDraft_callsTypeDraftAndClearsSession() async { + let generateExpectation = expectation(description: "generateCompose called") + let typeExpectation = expectation(description: "typeDraft called") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(composeResult(text: "Typed draft.")), + onGenerateCalled: { generateExpectation.fulfill() }, + onTypeDraftCalled: { typeExpectation.fulfill() } + ) + + // First Tab — generate. + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [generateExpectation], timeout: 1.0) + + // Second Tab — accept. + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [typeExpectation], timeout: 1.0) + + XCTAssertEqual(env.inserter.typedDrafts, ["Typed draft."]) + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + // MARK: - Cancellation + + func test_escape_inComposeMode_clearsActiveSession() async { + let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .dismissal)) + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + XCTAssertEqual(env.inserter.typedDrafts, []) + } + + func test_textMutationDuringPreview_clearsActiveSession() async { + let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + + _ = env.coordinator.handleInputEvent( + CotabbyTestFixtures.inputEvent(kind: .textMutation, characters: "x") + ) + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + func test_navigationDuringPreview_clearsActiveSession() async { + let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .navigation)) + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + func test_focusChange_duringPreview_clearsActiveSession() async { + let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + + env.focusModel.publish(snapshot: focusSnapshot(processIdentifier: 999, elementIdentifier: "different")) + await Task.yield() + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + func test_modeChangeToAutocomplete_clearsActiveSession() async { + let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + + env.settings.publish(snapshot: snapshot(mode: .autocomplete)) + await Task.yield() + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + // MARK: - Failure paths + + func test_emptyDraft_returnsToIdleAndHidesOverlay() async { + let generateExpectation = expectation(description: "generateCompose called") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(composeResult(text: " \n ")), + onGenerateCalled: { generateExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [generateExpectation], timeout: 1.0) + // Empty-draft handling is fully synchronous after `generateCompose` resolves; yield once + // so the awaiting `await applyComposeResult` runs. + await Task.yield() + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + if case .idle = env.coordinator.state { + // expected + } else { + XCTFail("Expected idle state after empty compose result, got \(env.coordinator.state)") + } + } + + func test_engineFailure_surfacesFailedState() async { + let generateExpectation = expectation(description: "generateCompose called") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .failure(SuggestionClientError.unavailable("Compose Mode requires tabby-depth-1.")), + onGenerateCalled: { generateExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [generateExpectation], timeout: 1.0) + await Task.yield() + + if case .failed(let message) = env.coordinator.state { + XCTAssertTrue(message.contains("tabby-depth-1")) + } else { + XCTFail("Expected failed state after engine failure, got \(env.coordinator.state)") + } + } + + // MARK: - Test environment + + private struct ComposeTestEnvironment { + let coordinator: SuggestionCoordinator + let permissions: FakeSuggestionPermissions + let focusModel: FakeFocusModel + let inputMonitor: FakeInputMonitor + let overlayController: FakeOverlayController + let inserter: FakeSuggestionInserter + let engine: FakeSuggestionEngine + let settings: FakeSuggestionSettings + let clipboard: FakeClipboardContextProvider + let visualContext: FakeVisualContextCoordinator + let composeCollector: FakeComposeContextCollector + } + + private func makeEnvironment( + mode: SuggestionInteractionMode, + engineBehavior: FakeSuggestionEngine.Behavior, + onGenerateCalled: (() -> Void)? = nil, + onShowComposePreview: (() -> Void)? = nil, + onTypeDraftCalled: (() -> Void)? = nil + ) -> ComposeTestEnvironment { + let permissions = FakeSuggestionPermissions() + let focusModel = FakeFocusModel(initialSnapshot: focusSnapshot()) + let inputMonitor = FakeInputMonitor() + let overlayController = FakeOverlayController() + overlayController.onShowComposePreview = onShowComposePreview + let inserter = FakeSuggestionInserter() + inserter.onTypeDraftCalled = onTypeDraftCalled + let engine = FakeSuggestionEngine(behavior: engineBehavior) + engine.onComposeCalled = onGenerateCalled + let settings = FakeSuggestionSettings(initialSnapshot: snapshot(mode: mode)) + let clipboard = FakeClipboardContextProvider() + let visualContext = FakeVisualContextCoordinator() + let composeCollector = FakeComposeContextCollector() + + let interactionState = SuggestionInteractionState() + let workController = SuggestionWorkController() + let coordinator = SuggestionCoordinator( + permissionManager: permissions, + focusModel: focusModel, + inputMonitor: inputMonitor, + overlayController: overlayController, + suggestionInserter: inserter, + suggestionEngine: engine, + suggestionSettings: settings, + clipboardContextProvider: clipboard, + clipboardRelevanceFilter: FakeClipboardRelevanceFilter(), + visualContextCoordinator: visualContext, + composeContextCollector: composeCollector, + interactionState: interactionState, + workController: workController, + configuration: .standard, + userDefaults: isolatedUserDefaults() + ) + Self.retainedCoordinators.append(coordinator) + return ComposeTestEnvironment( + coordinator: coordinator, + permissions: permissions, + focusModel: focusModel, + inputMonitor: inputMonitor, + overlayController: overlayController, + inserter: inserter, + engine: engine, + settings: settings, + clipboard: clipboard, + visualContext: visualContext, + composeCollector: composeCollector + ) + } + + /// Spins up an environment, drives a first Tab through it, and waits for the compose preview + /// to be ready. Cancellation tests use this so they only have to assert post-conditions. + private func makeEnvironmentWithReadyDraft(draftText: String) async -> ComposeTestEnvironment { + let showPreviewExpectation = expectation(description: "showComposePreview called") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(composeResult(text: draftText)), + onShowComposePreview: { showPreviewExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [showPreviewExpectation], timeout: 1.0) + return env + } + + private func snapshot(mode: SuggestionInteractionMode) -> SuggestionSettingsSnapshot { + SuggestionSettingsSnapshot( + isGloballyEnabled: true, + disabledAppBundleIdentifiers: [], + selectedInteractionMode: mode, + selectedEngine: .llamaOpenSource, + selectedWordCountPreset: .sevenToTwelve, + isClipboardContextEnabled: false, + userName: "Tester", + userTags: [], + debounceMilliseconds: 50, + focusPollIntervalMilliseconds: 50, + isMultiLineEnabled: false + ) + } + + private func focusSnapshot( + processIdentifier: Int32 = 123, + elementIdentifier: String = "field" + ) -> FocusSnapshot { + let inputSnapshot = CotabbyTestFixtures.focusedInputSnapshot( + processIdentifier: processIdentifier, + elementIdentifier: elementIdentifier + ) + return FocusSnapshot( + applicationName: inputSnapshot.applicationName, + bundleIdentifier: inputSnapshot.bundleIdentifier, + capability: .supported, + context: inputSnapshot, + inspection: nil + ) + } + + private func composeResult(text: String) -> ComposeResult { + ComposeResult(generation: 1, rawText: text, text: text, latency: 0.05) + } + + private func isolatedUserDefaults() -> UserDefaults { + let suiteName = "SuggestionCoordinatorComposeTests-\(UUID().uuidString)" + guard let userDefaults = UserDefaults(suiteName: suiteName) else { + return .standard + } + userDefaults.removePersistentDomain(forName: suiteName) + return userDefaults + } +} + +// MARK: - Fakes + +@MainActor +private final class FakeSuggestionPermissions: SuggestionPermissionProviding { + var inputMonitoringGranted = true + var screenRecordingGranted = true + private let inputMonitoringSubject = PassthroughSubject() + private let screenRecordingSubject = PassthroughSubject() + + var inputMonitoringGrantedPublisher: AnyPublisher { + inputMonitoringSubject.eraseToAnyPublisher() + } + var screenRecordingGrantedPublisher: AnyPublisher { + screenRecordingSubject.eraseToAnyPublisher() + } +} + +@MainActor +private final class FakeFocusModel: SuggestionFocusProviding { + private(set) var snapshot: FocusSnapshot + private let subject: CurrentValueSubject + + init(initialSnapshot: FocusSnapshot) { + snapshot = initialSnapshot + subject = CurrentValueSubject(initialSnapshot) + } + + var snapshotPublisher: AnyPublisher { + subject.eraseToAnyPublisher() + } + + func refreshNow() {} + + func publish(snapshot: FocusSnapshot) { + self.snapshot = snapshot + subject.send(snapshot) + } +} + +@MainActor +private final class FakeInputMonitor: SuggestionInputMonitoring { + var onEvent: ((CapturedInputEvent) -> Bool)? + var onSuppressedSyntheticInput: (() -> Void)? +} + +@MainActor +private final class FakeOverlayController: SuggestionOverlayControlling { + var state: OverlayState = .hidden(reason: "test idle") + var onStateChange: ((OverlayState) -> Void)? + var onShowComposePreview: (() -> Void)? + private(set) var composePreviewText: String? + private(set) var hideReasons: [String] = [] + + func showSuggestion(_ text: String, geometry: SuggestionOverlayGeometry) { + state = .visible(text: text, geometry: geometry) + onStateChange?(state) + } + + func showComposePreview(_ text: String, geometry: SuggestionOverlayGeometry) { + composePreviewText = text + state = .composePreview(text: text, geometry: geometry) + onStateChange?(state) + onShowComposePreview?() + } + + func hide(reason: String) { + hideReasons.append(reason) + state = .hidden(reason: reason) + onStateChange?(state) + } +} + +@MainActor +private final class FakeSuggestionInserter: SuggestionInserting { + var lastErrorMessage: String? + var onTypeDraftCalled: (() -> Void)? + private(set) var typedDrafts: [String] = [] + + func insert(_ suggestion: String) -> Bool { true } + + func typeDraft(_ draft: String, shouldContinue: @escaping @MainActor () -> Bool) async -> Bool { + typedDrafts.append(draft) + onTypeDraftCalled?() + return true + } +} + +@MainActor +private final class FakeSuggestionEngine: SuggestionGenerating { + enum Behavior { + case success(ComposeResult) + case failure(Error) + } + + private let behavior: Behavior + var onComposeCalled: (() -> Void)? + private(set) var composeCallCount = 0 + + init(behavior: Behavior) { + self.behavior = behavior + } + + func generateSuggestion(for request: SuggestionRequest) async throws -> SuggestionResult { + throw SuggestionClientError.unavailable("Autocomplete not exercised here.") + } + + func generateCompose(for request: ComposeRequest) async throws -> ComposeResult { + composeCallCount += 1 + onComposeCalled?() + switch behavior { + case .success(let result): + return ComposeResult( + generation: request.generation, + rawText: result.rawText, + text: result.text, + latency: result.latency + ) + case .failure(let error): + throw error + } + } + + func resetCachedGenerationContext() async {} +} + +@MainActor +private final class FakeSuggestionSettings: SuggestionSettingsProviding { + private(set) var snapshot: SuggestionSettingsSnapshot + private let subject: CurrentValueSubject + + init(initialSnapshot: SuggestionSettingsSnapshot) { + snapshot = initialSnapshot + subject = CurrentValueSubject(initialSnapshot) + } + + var snapshotPublisher: AnyPublisher { + subject.eraseToAnyPublisher() + } + + func publish(snapshot: SuggestionSettingsSnapshot) { + self.snapshot = snapshot + subject.send(snapshot) + } +} + +@MainActor +private final class FakeClipboardContextProvider: ClipboardContextProviding { + var contextToReturn: String? + var currentChangeCount: Int = 0 + func currentContext() -> String? { contextToReturn } +} + +@MainActor +private final class FakeClipboardRelevanceFilter: ClipboardRelevanceFiltering { + func filter(clipboard: String?, pasteboardChangeCount: Int, precedingText: String) -> String? { + clipboard + } +} + +@MainActor +private final class FakeVisualContextCoordinator: VisualContextCoordinating { + var status: VisualContextStatus = .idle + var latestExcerpt: String? + var onStateChange: ((VisualContextStatus, String?) -> Void)? + var onInjectedContextReady: ((FocusedInputIdentity) -> Void)? + + func startSessionIfNeeded(for snapshotContext: FocusedInputSnapshot) {} + func cancel(resetState: Bool) {} + func excerpt(for context: FocusedInputContext) -> String? { nil } +} + +@MainActor +private final class FakeComposeContextCollector: ComposeContextCollecting { + var textToReturn: String = "Surrounding context." + private(set) var collectCallCount = 0 + + func collect(for context: FocusedInputContext) async throws -> ComposeContextCollectionResult { + collectCallCount += 1 + return ComposeContextCollectionResult( + text: textToReturn, + visitedNodeCount: 1, + retainedTextCount: 1, + droppedTextCount: 0 + ) + } +} From 894e8eb1f50e3688416d1b3963989738bf8ab291 Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Mon, 25 May 2026 00:23:15 -0700 Subject: [PATCH 2/4] Drop the tabby-depth-1 requirement; let any local model drive Compose MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RFC reserved Compose Mode for one specific model. In practice the gate ended up being three layers of plumbing for a constraint we cannot even download — `tabby-depth-1` (gemma-3n-E4B-it-Q4_K_M.gguf) isn't in the model catalog yet. Stripping the gate lets any local llama model draft a Compose response while we figure out which model actually ships. Removed: - `LlamaSuggestionEngine.generateCompose` model guard. - `LlamaRuntimeManager.selectedModelSupportsCompose`. - `RuntimeModelCatalog.composeRequiredFilename` and `supportsCompose`. - `RuntimeBootstrapModel.prepareForComposeMode`, `restoreAutocompleteModel`, `isComposeRequiredModelInstalled`, and the pre-Compose model memory. - The mode-change subscription in `CotabbyAppEnvironment` that drove the auto-switch. - The "Compose requires tabby-depth-1" banner in menu bar and Settings. - The model-picker disable that pinned the user to one filename in Compose Mode. Kept the engine-level gate: Compose still routes to llama only, so the "Compose requires the Open Source engine" banner stays. Apple Intelligence and MLX would need their own `generateCompose` work. 310 tests, 0 failures. --- Cotabby/App/Core/CotabbyAppEnvironment.swift | 20 ---------- Cotabby/Models/LlamaRuntimeModels.swift | 6 --- Cotabby/Models/RuntimeBootstrapModel.swift | 39 ------------------- .../Runtime/LlamaRuntimeManager.swift | 4 -- .../Runtime/LlamaSuggestionEngine.swift | 6 --- Cotabby/UI/MenuBarView.swift | 13 ------- Cotabby/UI/SettingsView.swift | 11 ------ 7 files changed, 99 deletions(-) diff --git a/Cotabby/App/Core/CotabbyAppEnvironment.swift b/Cotabby/App/Core/CotabbyAppEnvironment.swift index 0ecca66..282fff3 100644 --- a/Cotabby/App/Core/CotabbyAppEnvironment.swift +++ b/Cotabby/App/Core/CotabbyAppEnvironment.swift @@ -184,26 +184,6 @@ final class CotabbyAppEnvironment { } .store(in: &cancellables) - // Compose Mode requires `tabby-depth-1`. Auto-switch on entry / restore on exit lives here - // (not on `RuntimeBootstrapModel`) because the runtime model intentionally does not know - // about interaction modes. `dropFirst()` ignores the initial publish; `removeDuplicates` - // avoids redundant work when other settings change. - suggestionSettings.$selectedInteractionMode - .removeDuplicates() - .dropFirst() - .sink { [weak runtimeModel] mode in - guard let runtimeModel else { return } - Task { @MainActor in - switch mode { - case .compose: - await runtimeModel.prepareForComposeMode() - case .autocomplete: - await runtimeModel.restoreAutocompleteModel() - } - } - } - .store(in: &cancellables) - // Key code changes reach InputMonitor through closures that read from the model // at event time (set above), so no Combine subscription is needed here. } diff --git a/Cotabby/Models/LlamaRuntimeModels.swift b/Cotabby/Models/LlamaRuntimeModels.swift index dc2d878..c8c8ba1 100644 --- a/Cotabby/Models/LlamaRuntimeModels.swift +++ b/Cotabby/Models/LlamaRuntimeModels.swift @@ -138,12 +138,6 @@ struct RemoteModelFile: Equatable, Hashable, Sendable { } enum RuntimeModelCatalog { - static let composeRequiredFilename = "gemma-3n-E4B-it-Q4_K_M.gguf" - - static func supportsCompose(filename: String?) -> Bool { - filename == composeRequiredFilename - } - static func displayName(for filename: String) -> String { switch filename { case "Qwen3-0.6B-Q4_K_M.gguf": diff --git a/Cotabby/Models/RuntimeBootstrapModel.swift b/Cotabby/Models/RuntimeBootstrapModel.swift index c6e40ac..3e18a56 100644 --- a/Cotabby/Models/RuntimeBootstrapModel.swift +++ b/Cotabby/Models/RuntimeBootstrapModel.swift @@ -19,10 +19,6 @@ final class RuntimeBootstrapModel: ObservableObject { private let userDefaults: UserDefaults private var cancellables = Set() private var runtimeTask: Task? - /// The model the user had selected before Compose Mode auto-switched to `tabby-depth-1`. - /// Persisted only in memory because the Compose round-trip is a within-session feature; if the - /// app relaunches in Compose Mode we re-select the required model on the next entry anyway. - private var preComposeAutocompleteModelFilename: String? /// Called immediately before the runtime begins switching models so suggestion state can reset. var onWillReloadModel: (() -> Void)? @@ -135,41 +131,6 @@ final class RuntimeBootstrapModel: ObservableObject { await runtimeTask?.value } - /// Whether the Compose Mode required local model is currently discovered. - /// Read by UI so it can offer a clear "Compose needs tabby-depth-1" message without each view - /// duplicating the catalog lookup. - var isComposeRequiredModelInstalled: Bool { - availableModels.contains { $0.filename == RuntimeModelCatalog.composeRequiredFilename } - } - - /// Switches the runtime to `tabby-depth-1` for Compose Mode, remembering the prior selection so - /// returning to Autocomplete restores it. No-op when the required model is missing, when it is - /// already selected, or when a runtime task is in flight (the next entry will retry). - func prepareForComposeMode() async { - let required = RuntimeModelCatalog.composeRequiredFilename - guard isComposeRequiredModelInstalled else { - TabbyLogger.runtime.info("Compose Mode required model \(required) is not installed; leaving selection unchanged.") - return - } - guard selectedModelFilename != required else { return } - guard runtimeTask == nil else { return } - - preComposeAutocompleteModelFilename = selectedModelFilename - await selectModel(required) - } - - /// Restores the user's pre-Compose model selection. No-op when no prior selection was saved - /// or when the saved model is no longer available. - func restoreAutocompleteModel() async { - guard let previous = preComposeAutocompleteModelFilename else { return } - preComposeAutocompleteModelFilename = nil - guard availableModels.contains(where: { $0.filename == previous }) else { return } - guard selectedModelFilename != previous else { return } - guard runtimeTask == nil else { return } - - await selectModel(previous) - } - /// Cancels pending startup work and forwards shutdown to the underlying runtime manager. func stop() { runtimeTask?.cancel() diff --git a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift index 81c26fb..6be8d85 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift @@ -23,10 +23,6 @@ final class LlamaRuntimeManager: ObservableObject { private var cachedRuntime: PreparedLlamaRuntime? private var selectedModelFilename: String? - var selectedModelSupportsCompose: Bool { - RuntimeModelCatalog.supportsCompose(filename: selectedModelFilename) - } - convenience init() { self.init( configuration: .default, diff --git a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift index e93350c..1e4d8b1 100644 --- a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift @@ -75,12 +75,6 @@ final class LlamaSuggestionEngine { } func generateCompose(for request: ComposeRequest) async throws -> ComposeResult { - guard runtimeManager.selectedModelSupportsCompose else { - throw SuggestionClientError.unavailable( - "Compose Mode requires tabby-depth-1. Select or download \(RuntimeModelCatalog.composeRequiredFilename)." - ) - } - do { let startTime = Date() let prompt = ComposePromptRenderer.prompt(for: request) diff --git a/Cotabby/UI/MenuBarView.swift b/Cotabby/UI/MenuBarView.swift index e07e825..d7a9a15 100644 --- a/Cotabby/UI/MenuBarView.swift +++ b/Cotabby/UI/MenuBarView.swift @@ -126,14 +126,6 @@ struct MenuBarView: View { .foregroundStyle(.orange) } - if suggestionSettings.selectedInteractionMode == .compose, - suggestionSettings.selectedEngine == .llamaOpenSource, - !runtimeModel.isComposeRequiredModelInstalled { - Text("Compose requires \(RuntimeModelCatalog.composeRequiredFilename). Add it to the models folder.") - .font(.caption) - .foregroundStyle(.orange) - } - if suggestionSettings.selectedEngine.supportsLocalModelManagement { modelRow } @@ -353,11 +345,6 @@ struct MenuBarView: View { } private var runtimePickerDisabled: Bool { - // Compose Mode pins the runtime to `tabby-depth-1`, so user-driven model selection is - // disabled while Compose is active to prevent half-supported configurations. - if suggestionSettings.selectedInteractionMode == .compose { - return true - } switch runtimeModel.state { case .starting, .loading: return true diff --git a/Cotabby/UI/SettingsView.swift b/Cotabby/UI/SettingsView.swift index fa487c0..7794ad6 100644 --- a/Cotabby/UI/SettingsView.swift +++ b/Cotabby/UI/SettingsView.swift @@ -182,17 +182,6 @@ struct SettingsView: View { .foregroundStyle(.orange) } - if suggestionSettings.selectedInteractionMode == .compose, - suggestionSettings.selectedEngine == .llamaOpenSource, - !runtimeModel.isComposeRequiredModelInstalled { - Text( - "Compose requires \(RuntimeModelCatalog.composeRequiredFilename). " - + "Add it to your models folder to enable draft generation." - ) - .font(.caption) - .foregroundStyle(.orange) - } - Picker("Length", selection: selectedWordCountPresetBinding) { ForEach(SuggestionWordCountPreset.allCases) { preset in Text(preset.displayLabel) From cb4f934ae0f9e1abe79db3d980af53df0a27657c Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Mon, 25 May 2026 00:58:33 -0700 Subject: [PATCH 3/4] Stream Compose tokens straight into the focused field Old flow: press Tab in Compose Mode, wait several seconds, an overlay popup eventually shows the whole draft, press Tab again to type it. There was no visual feedback while the model was generating, so it felt like nothing was happening. New flow: press Tab once, the model starts sampling, and each piece is typed into the focused field as it arrives. Esc / focus change / typing cancels the stream. Already-typed characters stay where they are. How it works - `LlamaRuntimeCore.summarizeStreaming` runs the same ephemeral-sequence sampling as `summarize` but invokes an `onToken` callback per piece instead of accumulating into a single returned string. - `LlamaRuntimeManager.streamUncached` wraps that in an `AsyncThrowingStream` after the runtime is prepared. The stream's `onTermination` cancels the detached sampling task, so workController.cancelAll() stops generation cleanly. - `SuggestionGenerating.generateComposeStreaming` is added to the protocol with a one-shot default (call `generateCompose`, emit once). `LlamaSuggestionEngine` overrides it for real streaming. The router forwards llama-only the same way it does for `generateCompose`. - `SuggestionCoordinator+Compose` consumes the stream on the main actor and calls `suggestionInserter.insert(piece)` per piece. The active `ActiveComposeSession` accumulates `fullText` so logs and diagnostics describe the full draft. - Tab during streaming is absorbed without restarting (no double draft). Esc, focus change, typing, mode change all run `cancelComposeWork`, which cancels the work controller, clears the session, and lets the stream terminate. Coordinator-level guards - `composeStreamShouldContinue` checks the focused field identity before posting each piece. If focus shifts mid-stream, the loop breaks and no further synthetic keys reach the wrong app. - Tabby's own synthetic key events are already absorbed by `InputSuppressionController` before they reach `handleInputEvent`, so the stream's own characters never trigger the textMutation-cancels-compose path. Removed - The two-step Tab `acceptComposeDraft` flow. - `presentComposePreview` overlay path inside the coordinator (the overlay controller's `showComposePreview` stays for future preview modes). Tests - `SuggestionCoordinatorComposeTests` rewritten for streaming: - first Tab streams pieces into the field via the inserter - second Tab during streaming is absorbed (no duplicate stream) - Esc/navigation/text-mutation/focus-change/mode-change cancel cleanly and partial pieces stay typed - engine throw surfaces a `.failed` state - The fake engine uses a `.blocked(initialPieces:)` behavior that yields a few pieces then parks until the underlying task is cancelled, so cancellation paths can be exercised deterministically. 309 tests, 0 failures. --- .../SuggestionCoordinator+Compose.swift | 354 +++++++----------- .../Models/SuggestionSubsystemContracts.swift | 18 + .../Services/Runtime/LlamaRuntimeCore.swift | 67 ++++ .../Runtime/LlamaRuntimeManager.swift | 36 ++ .../Runtime/LlamaSuggestionEngine.swift | 19 + .../Runtime/SuggestionEngineRouter.swift | 15 + .../SuggestionInteractionState.swift | 18 + .../SuggestionCoordinatorComposeTests.swift | 259 +++++++------ 8 files changed, 444 insertions(+), 342 deletions(-) diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift index 2e5f182..1a9739e 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift @@ -4,16 +4,15 @@ import Logging /// File overview: /// Compose Mode entry points for `SuggestionCoordinator`. -/// This is the deliberate two-step Tab flow: first Tab gathers context, generates a full draft, -/// and shows a preview; second Tab types the draft into the focused field. Autocomplete continues -/// to live in the sibling extension files and is untouched by this path. /// -/// State invariants this file protects: -/// - Only one Compose generation is in flight per `currentWorkID`; later Tabs cancel earlier work. -/// - The `ActiveComposeSession` is never accepted against a focused field whose process/identity -/// has changed since the draft was generated. -/// - Synthetic typing is rearmed against `InputSuppressionController` per chunk via the -/// `SuggestionInserter.typeDraft` contract. +/// Interaction model (single-Tab streaming): +/// - First Tab → gather AX context, build request, open a streaming generation against llama. +/// Each sampled piece is typed straight into the focused field via `SuggestionInserter.insert`. +/// - Escape, focus change, app/global disable, or the user typing → cancel the stream. Already- +/// typed characters stay in the field; the cancellation simply stops the next piece from +/// landing. +/// - Subsequent Tabs while streaming are absorbed so the user does not pile a second draft onto +/// the first. extension SuggestionCoordinator { // MARK: - Tab Routing @@ -22,8 +21,10 @@ extension SuggestionCoordinator { func handleComposeInputEvent(_ event: CapturedInputEvent) -> Bool { switch event.kind { case .acceptance, .fullAcceptance: - if interactionState.activeComposeSession != nil { - return acceptComposeDraft() + // Already streaming? Swallow the Tab so it does not start a second stream and does + // not reach the host app while we are typing into it. + if interactionState.activeComposeSession != nil || isAnyComposeWorkInFlight { + return true } return startComposeGeneration() @@ -32,16 +33,17 @@ extension SuggestionCoordinator { return false case .navigation: - // Arrow keys, page navigation, etc. — drop any in-flight draft because the field - // context the user originally asked Tabby to draft against has moved. + // Arrow keys, page navigation, etc. — drop in-flight streams because the field the + // user originally asked Tabby to draft into has moved. if interactionState.activeComposeSession != nil || isAnyComposeWorkInFlight { cancelComposeWork(reason: "Compose cancelled because the caret moved.") } return false case .textMutation, .shortcutMutation: - // Typing or paste during preview invalidates the draft. Compose is "ask, review, type"; - // mid-stream edits mean the user has stopped reviewing. + // Real user typing during streaming → stop. Tabby's own synthetic key events are + // absorbed by `InputSuppressionController` before they reach this handler, so the + // stream's own characters never trigger this path. if interactionState.activeComposeSession != nil || isAnyComposeWorkInFlight { cancelComposeWork(reason: "Compose cancelled because the focused text changed.") } @@ -54,8 +56,7 @@ extension SuggestionCoordinator { // MARK: - Generation - /// First-Tab handler: kick off Compose generation against the current focused field. - /// Returns `true` to consume the Tab so the host app does not receive it. + /// Streams a Compose draft into the currently focused field. Returns `true` to consume the Tab. @discardableResult func startComposeGeneration() -> Bool { guard permissionManager.inputMonitoringGranted else { @@ -83,12 +84,13 @@ extension SuggestionCoordinator { // Reuse the debounced-work plumbing with a zero delay so cancellation and stale-work guards // are identical to the autocomplete path. Compose has no real debounce — Tab is explicit. let workID = workController.replaceDebouncedWork(delayMilliseconds: 0) { [weak self] workID in - await self?.runComposeGeneration(for: context, workID: workID) + await self?.runComposeStreaming(for: context, workID: workID) } latestGenerationNumber = context.generation + latestRawModelOutput = nil state = .generating logStage( - "compose-generating", + "compose-streaming-start", workID: workID, generation: context.generation, message: "Gathering Compose context for \(context.elementIdentifier) in \(context.applicationName)." @@ -96,8 +98,7 @@ extension SuggestionCoordinator { return true } - // swiftlint:disable:next cyclomatic_complexity - private func runComposeGeneration(for context: FocusedInputContext, workID: UInt64) async { + private func runComposeStreaming(for context: FocusedInputContext, workID: UInt64) async { guard workController.isCurrent(workID) else { return } await awaitCachedGenerationContextResetIfNeeded() guard workController.isCurrent(workID) else { return } @@ -129,105 +130,134 @@ extension SuggestionCoordinator { visualContextSummary: visualContextSummary ) latestPromptPreview = buildResult.promptPreview - latestRawModelOutput = nil let request = buildResult.request + // The active session represents "we are streaming into this field". The full text is + // appended to as pieces arrive so logs and diagnostics can describe what was typed. + let initialSession = interactionState.startComposeSession( + fullText: "", + liveContext: context, + latency: 0 + ) + state = .typing + logStage( + "compose-streaming-begin", + workID: workID, + generation: context.generation, + message: "Streaming Compose draft into \(context.applicationName).", + prompt: buildResult.promptPreview + ) + workController.replaceGenerationWork(for: workID) { [weak self] in guard let self else { return } - do { - let result = try await self.suggestionEngine.generateCompose(for: request) - guard !Task.isCancelled, self.workController.isCurrent(workID) else { return } - await self.applyComposeResult(result, workID: workID) - } catch SuggestionClientError.cancelled { - return - } catch { - guard self.workController.isCurrent(workID) else { return } - await self.applyComposeFailure(error.localizedDescription, workID: workID) - } - } - } - - /// Stale-result guard mirrors the autocomplete path: we re-read focus before showing anything, - /// and bail if the field's generation or process no longer matches what we asked the model for. - private func applyComposeResult(_ result: ComposeResult, workID: UInt64) async { - guard workController.isCurrent(workID) else { return } - - focusModel.refreshNow() - let snapshot = focusModel.snapshot - - guard case .supported = snapshot.capability, let rawContext = snapshot.context else { - disablePredictions(reason: snapshot.capability.summary) - return - } - let liveContext = interactionState.materializeContext(from: rawContext) - - guard liveContext.generation == result.generation else { - latestRawModelOutput = SuggestionDebugLogger.debugPreview(result.rawText) - logStage( - "compose-stale-drop", + await self.consumeComposeStream( + request: request, workID: workID, - generation: result.generation, - message: "Dropped stale Compose draft because live generation is \(liveContext.generation).", - rawOutput: result.rawText, - normalizedOutput: result.text + initialSession: initialSession ) - hideOverlay(reason: "Overlay hidden because the focused field changed before the draft was ready.") - return } + } - latestRawModelOutput = SuggestionDebugLogger.debugPreview(result.rawText) - latestLatencyMilliseconds = Int(result.latency * 1000) - latestGenerationNumber = liveContext.generation + private func consumeComposeStream( + request: ComposeRequest, + workID: UInt64, + initialSession: ActiveComposeSession + ) async { + let startTime = Date() + var accumulatedText = "" + var session = initialSession - let trimmed = result.text.trimmingCharacters(in: .whitespacesAndNewlines) - guard !trimmed.isEmpty else { - interactionState.clearSuggestion() - hideOverlay(reason: "Overlay hidden because Compose returned an empty draft.") - state = .idle - logStage( - "compose-empty", + do { + let stream = try await suggestionEngine.generateComposeStreaming(for: request) + for try await piece in stream { + guard !Task.isCancelled, workController.isCurrent(workID) else { break } + guard composeStreamShouldContinue(matching: session) else { break } + guard !piece.isEmpty else { continue } + + accumulatedText += piece + latestRawModelOutput = SuggestionDebugLogger.debugPreview(accumulatedText) + _ = suggestionInserter.insert(piece) + session = interactionState.updateComposeSession( + session, + fullText: accumulatedText, + latency: Date().timeIntervalSince(startTime) + ) ?? session + + // Yield once per piece so cancellation tasks queued on the main actor (focus + // changes, Esc) can run between samples instead of getting starved by the loop. + await Task.yield() + } + } catch is CancellationError { + // Treat cancellation as a normal stop — partial text stays in the field. + await finishComposeStream( + accumulated: accumulatedText, + latency: Date().timeIntervalSince(startTime), workID: workID, - generation: result.generation, - message: "Compose draft was empty after normalization.", - rawOutput: result.rawText, - normalizedOutput: result.text + session: session, + outcome: ComposeStreamOutcome( + stage: "compose-streaming-cancelled", + stageMessage: "Compose stream cancelled." + ) ) return + } catch { + await applyComposeFailure(error.localizedDescription, workID: workID) + return } - let session = interactionState.startComposeSession( - fullText: result.text, - liveContext: liveContext, - latency: result.latency - ) - latestSuggestionPreview = session.fullText - latestFullSuggestionPreview = session.fullText - latestRemainingSuggestionPreview = session.fullText - latestAcceptedCharacterCount = 0 - latestRemainingCharacterCount = session.fullText.count - state = .ready(text: session.fullText, latency: session.latency) - - presentComposePreview( - text: session.fullText, - at: liveContext.caretRect, - inputFrameRect: liveContext.inputFrameRect, - caretQuality: liveContext.caretQuality, - observedCharWidth: liveContext.observedCharWidth, - isRightToLeft: TextDirectionDetector.isRightToLeft(liveContext.precedingText) + await finishComposeStream( + accumulated: accumulatedText, + latency: Date().timeIntervalSince(startTime), + workID: workID, + session: session, + outcome: ComposeStreamOutcome( + stage: "compose-streaming-done", + stageMessage: "Compose stream finished." + ) ) + } + + private struct ComposeStreamOutcome { + let stage: String + let stageMessage: String + } + + private func finishComposeStream( + accumulated: String, + latency: TimeInterval, + workID: UInt64, + session: ActiveComposeSession, + outcome: ComposeStreamOutcome + ) async { + guard workController.isCurrent(workID) else { return } + + latestLatencyMilliseconds = Int(latency * 1000) + latestRawModelOutput = SuggestionDebugLogger.debugPreview(accumulated) + latestAcceptanceAction = accumulated.isEmpty + ? "Compose stream produced no text." + : "Compose draft streamed into the field." + + if interactionState.activeComposeSession == session { + interactionState.clearComposeSession(session) + } + hideOverlay(reason: outcome.stageMessage) + state = .idle logStage( - "compose-ready", + outcome.stage, workID: workID, - generation: result.generation, - message: "Compose draft ready for review.", - rawOutput: result.rawText, - normalizedOutput: result.text + generation: session.baseContext.generation, + message: outcome.stageMessage, + normalizedOutput: accumulated ) } private func applyComposeFailure(_ message: String, workID: UInt64) async { guard workController.isCurrent(workID) else { return } - interactionState.clearSuggestion() + if let session = interactionState.activeComposeSession { + interactionState.clearComposeSession(session) + } else { + interactionState.clearSuggestion() + } hideOverlay(reason: "Overlay hidden because Compose generation failed.") state = .failed(message) logStage( @@ -238,92 +268,10 @@ extension SuggestionCoordinator { ) } - // MARK: - Acceptance - - /// Second-Tab handler: type the active Compose draft into the focused field via - /// `SuggestionInserter.typeDraft`. Each chunk re-checks focus identity before posting. - @discardableResult - func acceptComposeDraft() -> Bool { - guard let session = interactionState.activeComposeSession else { - return passTabThrough(reason: "Key passed through because no Compose draft was ready.") - } - - focusModel.refreshNow() - let snapshot = focusModel.snapshot - guard case .supported = snapshot.capability, let rawContext = snapshot.context else { - cancelComposeWork(reason: snapshot.capability.summary) - return passTabThrough(reason: snapshot.capability.summary) - } - let liveContext = interactionState.materializeContext(from: rawContext) - guard liveContext.identity == session.baseContext.identity, - liveContext.processIdentifier == session.baseContext.processIdentifier else { - cancelComposeWork(reason: "Compose cancelled because the focused field changed before typing began.") - return passTabThrough(reason: "Key passed through because the focused field changed.") - } - guard liveContext.selection.length == 0 else { - cancelComposeWork(reason: "Compose cancelled because text is selected.") - return passTabThrough(reason: "Key passed through because text is selected.") - } - - state = .typing - recordAcceptedWords(from: session.fullText) - logStage( - "compose-accepting", - workID: currentWorkID, - generation: liveContext.generation, - message: "Typing Compose draft (\(session.fullText.count) characters) into \(liveContext.applicationName).", - normalizedOutput: session.fullText - ) - - let typingSession = session - Task { @MainActor [weak self] in - guard let self else { return } - let didType = await self.suggestionInserter.typeDraft(typingSession.fullText) { [weak self] in - self?.composeTypingShouldContinue(matching: typingSession) ?? false - } - - // The session may have already been cleared by a cancellation path (focus change, - // mode change, Esc). If so, don't reset state again. - guard let active = self.interactionState.activeComposeSession, active == typingSession else { - return - } - - if didType { - self.interactionState.clearComposeSession(typingSession) - self.hideOverlay(reason: "Overlay hidden after Compose draft was typed into the field.") - self.state = .idle - self.latestAcceptanceAction = "Compose draft typed into the field." - self.logStage( - "compose-typed", - workID: self.currentWorkID, - generation: typingSession.baseContext.generation, - message: "Compose draft fully typed into the focused field.", - normalizedOutput: typingSession.fullText - ) - } else { - let message = self.suggestionInserter.lastErrorMessage - ?? "Compose typing stopped before the draft was complete." - self.interactionState.clearComposeSession(typingSession) - self.hideOverlay(reason: "Overlay hidden because Compose typing did not complete.") - self.state = .failed(message) - self.logStage( - "compose-type-aborted", - workID: self.currentWorkID, - generation: typingSession.baseContext.generation, - message: message, - normalizedOutput: typingSession.fullText - ) - } - } - return true - } - - /// Focus-identity guard checked between every synthetic key chunk. Bailing here lets the - /// inserter stop posting mid-draft when the user switches apps or fields. - private func composeTypingShouldContinue(matching session: ActiveComposeSession) -> Bool { - guard let active = interactionState.activeComposeSession, active == session else { - return false - } + /// Focus-identity guard checked before posting each streamed piece. Returns false when the + /// session has been cleared or the focused field has changed, which halts the for-await loop. + private func composeStreamShouldContinue(matching session: ActiveComposeSession) -> Bool { + guard interactionState.activeComposeSession == session else { return false } let snapshot = focusModel.snapshot guard case .supported = snapshot.capability, let rawContext = snapshot.context else { return false @@ -335,8 +283,9 @@ extension SuggestionCoordinator { // MARK: - Cancellation - /// Cancels any in-flight Compose work and clears the active session. Safe to call when nothing - /// is in flight; the underlying controllers are no-op in that case. + /// Cancels any in-flight Compose work and clears the active session. Already-typed characters + /// remain in the focused field — Compose's stream is fire-and-forget per piece, so we cannot + /// (and should not) try to undo what the host app has already accepted. func cancelComposeWork(reason: String) { let hadActiveSession = interactionState.activeComposeSession != nil let hadInflightWork = isAnyComposeWorkInFlight @@ -358,9 +307,7 @@ extension SuggestionCoordinator { ) } - /// True when a Compose generation or typing task could still emit results. We treat any state - /// other than `.idle` / `.disabled` as "could still emit" because the work controller only - /// surfaces task identity, not status. + /// True when a Compose generation or streaming task could still emit output. var isAnyComposeWorkInFlight: Bool { switch state { case .generating, .typing: @@ -369,33 +316,4 @@ extension SuggestionCoordinator { return false } } - - // MARK: - Overlay - - // swiftlint:disable function_parameter_count - /// Sibling of `presentOverlay` for the multiline Compose preview surface. - private func presentComposePreview( - text: String, - at caretRect: CGRect, - inputFrameRect: CGRect?, - caretQuality: CaretGeometryQuality, - observedCharWidth: CGFloat?, - isRightToLeft: Bool - ) { - let geometry = SuggestionOverlayGeometry( - caretRect: caretRect, - inputFrameRect: inputFrameRect, - caretQuality: caretQuality, - observedCharWidth: observedCharWidth, - isRightToLeft: isRightToLeft - ) - if let message = overlayPresenter.presentComposePreview( - text: text, - geometry: geometry, - previousState: overlayState - ) { - latestOverlayMessage = message - } - } - // swiftlint:enable function_parameter_count } diff --git a/Cotabby/Models/SuggestionSubsystemContracts.swift b/Cotabby/Models/SuggestionSubsystemContracts.swift index dcc2965..f4e1a2d 100644 --- a/Cotabby/Models/SuggestionSubsystemContracts.swift +++ b/Cotabby/Models/SuggestionSubsystemContracts.swift @@ -40,11 +40,29 @@ protocol SuggestionInputMonitoring: AnyObject { protocol SuggestionGenerating: AnyObject { func generateSuggestion(for request: SuggestionRequest) async throws -> SuggestionResult func generateCompose(for request: ComposeRequest) async throws -> ComposeResult + /// Streaming variant: yields each generated piece through an async stream so the coordinator + /// can type tokens into the focused field as the model produces them. Engines that cannot + /// stream natively get the one-shot fallback in the extension below. + func generateComposeStreaming(for request: ComposeRequest) async throws -> AsyncThrowingStream /// Clears backend-local continuation state when the focused editing context is no longer /// continuous. Stateless engines may implement this as a no-op. func resetCachedGenerationContext() async } +extension SuggestionGenerating { + /// One-shot fallback: run `generateCompose`, then emit the full draft as a single chunk. + /// Engines that can stream natively (Llama) override this to actually emit per token. + func generateComposeStreaming(for request: ComposeRequest) async throws -> AsyncThrowingStream { + let result = try await generateCompose(for: request) + return AsyncThrowingStream { continuation in + if !result.text.isEmpty { + continuation.yield(result.text) + } + continuation.finish() + } + } +} + @MainActor protocol SuggestionSettingsProviding: AnyObject { var snapshot: SuggestionSettingsSnapshot { get } diff --git a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift index fb8eed5..1e5d4b5 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift @@ -227,6 +227,73 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { return generatedText } + /// Streaming sibling of `summarize`. Same ephemeral-sequence semantics (autocomplete KV cache + /// stays clean), but emits each sampled piece through `onToken` before the loop continues so + /// callers can observe generation in real time. Returns the accumulated text for diagnostics. + /// + /// `onToken` is `@Sendable` and runs on the detached sampling thread; callers that need to + /// touch UI must hop back to the main actor. + func summarizeStreaming( + prompt: String, + options: LlamaGenerationOptions, + onToken: @Sendable (String) -> Void + ) throws -> String { + guard let preparedRuntime else { + throw LlamaRuntimeError.unavailable("The llama model is not loaded.") + } + + lifecycleCondition.lock() + guard !isShuttingDown else { + lifecycleCondition.unlock() + throw LlamaRuntimeError.unavailable("The runtime is shutting down.") + } + activeOperationCount += 1 + lifecycleCondition.unlock() + + defer { + lifecycleCondition.lock() + activeOperationCount -= 1 + lifecycleCondition.broadcast() + lifecycleCondition.unlock() + } + + let allPromptTokens = tokenize(prompt) + guard !allPromptTokens.isEmpty else { + throw LlamaRuntimeError.generationFailed("Tokenization returned no prompt tokens.") + } + + let maxPromptTokens = max(1, preparedRuntime.contextWindowTokens - options.maxPredictionTokens) + let promptTokens = allPromptTokens.count > maxPromptTokens + ? Array(allPromptTokens.suffix(maxPromptTokens)) + : allPromptTokens + + let config = Self.samplingConfig(from: options) + let seqID = engine.createSequence(config) + guard seqID >= 0 else { + throw LlamaRuntimeError.generationFailed("Unable to create streaming sequence.") + } + defer { engine.destroySequence(seqID) } + + var tokens = promptTokens + let status = engine.decodePrompt(seqID, &tokens, Int32(tokens.count), 0) + guard status == .ok else { + throw LlamaRuntimeError.generationFailed("Streaming prompt decoding failed.") + } + + var generatedText = "" + for _ in 0 ..< options.maxPredictionTokens { + if Task.isCancelled { break } + let result = engine.sampleNext(seqID) + if result.is_eos || result.was_cancelled { break } + let piece = Self.extractPiece(result) + guard !piece.isEmpty else { continue } + generatedText += piece + onToken(piece) + } + + return generatedText + } + // MARK: - Cache and lifecycle /// Drops the reusable autocomplete sequence while keeping the loaded model alive. diff --git a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift index 6be8d85..a5db439 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift @@ -179,6 +179,42 @@ final class LlamaRuntimeManager: ObservableObject { } } + /// Streaming variant of `generateUncached`: kicks off generation on a detached thread and + /// yields each sampled piece through an `AsyncThrowingStream`. Cancelling the consuming task + /// terminates the stream and cancels the underlying detached sampling task. + /// + /// `preparedRuntime()` runs on the main actor before we hand control to the stream so the + /// caller can `await` once and then iterate; that keeps the AsyncStream's setup cheap. + func streamUncached( + prompt: String, + options: LlamaGenerationOptions + ) async throws -> AsyncThrowingStream { + _ = try await preparedRuntime() + + let core = self.core + return AsyncThrowingStream { continuation in + let task = Task.detached { + do { + _ = try core.summarizeStreaming( + prompt: prompt, + options: options, + onToken: { piece in + continuation.yield(piece) + } + ) + continuation.finish() + } catch is CancellationError { + continuation.finish(throwing: LlamaRuntimeError.cancelled) + } catch let error as LlamaRuntimeError { + continuation.finish(throwing: error) + } catch { + continuation.finish(throwing: LlamaRuntimeError.generationFailed(error.localizedDescription)) + } + } + continuation.onTermination = { _ in task.cancel() } + } + } + /// Clears the native prompt KV cache without unloading the model. func resetPromptCache() { core.resetPromptCache() diff --git a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift index 1e4d8b1..a7c7ed8 100644 --- a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift @@ -114,6 +114,25 @@ final class LlamaSuggestionEngine { } } + /// Streaming variant — yields each sampled piece through an `AsyncThrowingStream` so the + /// coordinator can type into the focused field as the model generates. We deliberately skip + /// `ComposeTextNormalizer` here because that normalizer is whole-text shaped; streaming pieces + /// belong to the user's field immediately and any wrapper cleanup would have to happen against + /// already-typed characters, which is more harm than help. + func generateComposeStreaming(for request: ComposeRequest) async throws -> AsyncThrowingStream { + let prompt = ComposePromptRenderer.prompt(for: request) + let options = LlamaGenerationOptions( + maxPredictionTokens: request.maxPredictionTokens, + temperature: request.temperature, + topK: request.topK, + topP: request.topP, + minP: request.minP, + repetitionPenalty: request.repetitionPenalty, + seed: request.randomSeed + ) + return try await runtimeManager.streamUncached(prompt: prompt, options: options) + } + /// Clears both the Swift-side hint tracker and the native llama KV cache. /// The tracker reset is synchronous because it protects the next request from advertising /// stale reuse; awaiting the runtime reset keeps native KV invalidation ordered before the next diff --git a/Cotabby/Services/Runtime/SuggestionEngineRouter.swift b/Cotabby/Services/Runtime/SuggestionEngineRouter.swift index 376b23c..c916538 100644 --- a/Cotabby/Services/Runtime/SuggestionEngineRouter.swift +++ b/Cotabby/Services/Runtime/SuggestionEngineRouter.swift @@ -45,6 +45,21 @@ final class SuggestionEngineRouter { } } + /// Streaming Compose, llama-only. The protocol's default `generateComposeStreaming` falls back + /// to `generateCompose`, so non-llama engines route through the same "unavailable" path below + /// without us having to fork the routing decision twice. + func generateComposeStreaming(for request: ComposeRequest) async throws -> AsyncThrowingStream { + switch suggestionSettings.selectedEngine { + case .llamaOpenSource: + TabbyLogger.suggestion.debug("Streaming Compose request through open-source llama engine") + return try await llamaEngine.generateComposeStreaming(for: request) + case .appleIntelligence, .mlxSwift: + throw SuggestionClientError.unavailable( + "Compose Mode is only available with the local open-source runtime in this version." + ) + } + } + /// Compose currently only ships behind the local llama runtime. Other engines surface an /// unavailable state instead of forwarding so users see a clear "Compose needs llama" message /// rather than a generic prompt failure from a backend that does not implement the contract. diff --git a/Cotabby/Services/Suggestion/SuggestionInteractionState.swift b/Cotabby/Services/Suggestion/SuggestionInteractionState.swift index 1c3449e..7b66648 100644 --- a/Cotabby/Services/Suggestion/SuggestionInteractionState.swift +++ b/Cotabby/Services/Suggestion/SuggestionInteractionState.swift @@ -278,6 +278,24 @@ final class SuggestionInteractionState { activeSession = nil pendingInsertionConsumedCount = nil } + + /// Returns a new active compose session with the accumulated streaming text. No-op if the + /// previous session is no longer the active one (cancellation may have replaced it). + @discardableResult + func updateComposeSession( + _ previous: ActiveComposeSession, + fullText: String, + latency: TimeInterval + ) -> ActiveComposeSession? { + guard activeComposeSession == previous else { return nil } + let updated = ActiveComposeSession( + baseContext: previous.baseContext, + fullText: fullText, + latency: latency + ) + activeSession = .compose(updated) + return updated + } } /// Wraps reconciliation results with the live buffered context the coordinator needs for UI updates. diff --git a/CotabbyTests/SuggestionCoordinatorComposeTests.swift b/CotabbyTests/SuggestionCoordinatorComposeTests.swift index d4c6c0c..7dabc46 100644 --- a/CotabbyTests/SuggestionCoordinatorComposeTests.swift +++ b/CotabbyTests/SuggestionCoordinatorComposeTests.swift @@ -6,11 +6,9 @@ import XCTest /// Tests `SuggestionCoordinator+Compose` end to end with fakes for every collaborator. /// -/// Most coordinator tests in the autocomplete path live inside `SuggestionInteractionState` -/// helpers; Compose's orchestration is new enough that those state helpers cannot prove the -/// behavior reviewers actually care about: first Tab generates, second Tab types, focus changes -/// cancel mid-flight. The fakes below are all kept in this file so the contract under test stays -/// readable as one unit. +/// Compose Mode is single-Tab streaming: the first Tab opens a stream and each piece is typed +/// straight into the focused field via `SuggestionInserter.insert`. Esc/focus change/typing +/// cancels the stream; already-typed pieces stay. These tests lock in that contract. @MainActor final class SuggestionCoordinatorComposeTests: XCTestCase { private static var retainedCoordinators: [SuggestionCoordinator] = [] @@ -20,30 +18,30 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { try await super.tearDown() } - // MARK: - First Tab: generation + // MARK: - First Tab: streaming starts - func test_firstTab_inComposeMode_callsGenerateComposeAndShowsPreview() async { - let generateExpectation = expectation(description: "generateCompose called") - let showPreviewExpectation = expectation(description: "showComposePreview called") + func test_firstTab_inComposeMode_streamsPiecesIntoTheField() async { + let pieces = ["Hello", " team,", "\n", "Thanks."] + let finishedExpectation = expectation(description: "stream finished") let env = makeEnvironment( mode: .compose, - engineBehavior: .success(composeResult(text: "Hello team.")), - onGenerateCalled: { generateExpectation.fulfill() }, - onShowComposePreview: { showPreviewExpectation.fulfill() } + engineBehavior: .success(pieces: pieces), + onStreamFinished: { finishedExpectation.fulfill() } ) _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [generateExpectation, showPreviewExpectation], timeout: 1.0) + await fulfillment(of: [finishedExpectation], timeout: 1.0) + await Task.yield() - XCTAssertEqual(env.engine.composeCallCount, 1) - XCTAssertEqual(env.overlayController.composePreviewText, "Hello team.") - XCTAssertNotNil(env.coordinator.interactionState.activeComposeSession) + XCTAssertEqual(env.engine.composeStreamCallCount, 1) + XCTAssertEqual(env.inserter.insertedPieces, pieces) + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } func test_firstTab_inComposeMode_returnsTrueToConsumeTab() { let env = makeEnvironment( mode: .compose, - engineBehavior: .success(composeResult(text: "draft")) + engineBehavior: .success(pieces: ["only"]) ) let consumed = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) @@ -51,133 +49,139 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { XCTAssertTrue(consumed) } - func test_firstTab_inAutocompleteMode_doesNotInvokeComposePath() async { + func test_firstTab_inAutocompleteMode_doesNotInvokeComposeStream() async { let env = makeEnvironment( mode: .autocomplete, - engineBehavior: .success(composeResult(text: "should not run")) + engineBehavior: .success(pieces: ["should not run"]) ) _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - // Yield to let any spurious Task progress. await Task.yield() await Task.yield() - XCTAssertEqual(env.engine.composeCallCount, 0) - XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + XCTAssertEqual(env.engine.composeStreamCallCount, 0) + XCTAssertEqual(env.inserter.insertedPieces, []) } - // MARK: - Second Tab: acceptance + // MARK: - Subsequent Tabs while streaming - func test_secondTab_withActiveDraft_callsTypeDraftAndClearsSession() async { - let generateExpectation = expectation(description: "generateCompose called") - let typeExpectation = expectation(description: "typeDraft called") + func test_secondTab_whileStreaming_isAbsorbedWithoutRestartingTheStream() async { let env = makeEnvironment( mode: .compose, - engineBehavior: .success(composeResult(text: "Typed draft.")), - onGenerateCalled: { generateExpectation.fulfill() }, - onTypeDraftCalled: { typeExpectation.fulfill() } + engineBehavior: .blocked(initialPieces: ["partial"]) ) - // First Tab — generate. _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [generateExpectation], timeout: 1.0) + for _ in 0..<10 { await Task.yield() } - // Second Tab — accept. - _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [typeExpectation], timeout: 1.0) + let consumedSecondTab = env.coordinator.handleInputEvent( + CotabbyTestFixtures.inputEvent(kind: .acceptance) + ) - XCTAssertEqual(env.inserter.typedDrafts, ["Typed draft."]) - XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + XCTAssertTrue(consumedSecondTab, "second Tab during streaming should be absorbed, not passed through") + XCTAssertEqual(env.engine.composeStreamCallCount, 1, "second Tab must not start a second stream") } // MARK: - Cancellation - func test_escape_inComposeMode_clearsActiveSession() async { - let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + func test_escape_duringStreaming_cancelsAndKeepsTypedText() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["Hel", "lo"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<15 { await Task.yield() } _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .dismissal)) + for _ in 0..<10 { await Task.yield() } + XCTAssertEqual(env.inserter.insertedPieces, ["Hel", "lo"], "already-typed pieces stay in the field") XCTAssertNil(env.coordinator.interactionState.activeComposeSession) - XCTAssertEqual(env.inserter.typedDrafts, []) } - func test_textMutationDuringPreview_clearsActiveSession() async { - let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + func test_textMutationDuringStreaming_cancelsTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["streamed"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<15 { await Task.yield() } _ = env.coordinator.handleInputEvent( CotabbyTestFixtures.inputEvent(kind: .textMutation, characters: "x") ) + for _ in 0..<10 { await Task.yield() } XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } - func test_navigationDuringPreview_clearsActiveSession() async { - let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + func test_navigationDuringStreaming_cancelsTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["typed"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<15 { await Task.yield() } _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .navigation)) + for _ in 0..<10 { await Task.yield() } XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } - func test_focusChange_duringPreview_clearsActiveSession() async { - let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") - - env.focusModel.publish(snapshot: focusSnapshot(processIdentifier: 999, elementIdentifier: "different")) - await Task.yield() - - XCTAssertNil(env.coordinator.interactionState.activeComposeSession) - } + func test_focusChangeDuringStreaming_cancelsTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["before"]) + ) - func test_modeChangeToAutocomplete_clearsActiveSession() async { - let env = await makeEnvironmentWithReadyDraft(draftText: "Hello.") + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<15 { await Task.yield() } - env.settings.publish(snapshot: snapshot(mode: .autocomplete)) - await Task.yield() + env.focusModel.publish(snapshot: focusSnapshot(processIdentifier: 999, elementIdentifier: "different")) + for _ in 0..<10 { await Task.yield() } XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } - // MARK: - Failure paths - - func test_emptyDraft_returnsToIdleAndHidesOverlay() async { - let generateExpectation = expectation(description: "generateCompose called") + func test_modeChangeToAutocomplete_duringStreaming_cancelsTheStream() async { let env = makeEnvironment( mode: .compose, - engineBehavior: .success(composeResult(text: " \n ")), - onGenerateCalled: { generateExpectation.fulfill() } + engineBehavior: .blocked(initialPieces: ["before"]) ) _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [generateExpectation], timeout: 1.0) - // Empty-draft handling is fully synchronous after `generateCompose` resolves; yield once - // so the awaiting `await applyComposeResult` runs. - await Task.yield() + for _ in 0..<15 { await Task.yield() } + + env.settings.publish(snapshot: snapshot(mode: .autocomplete)) + for _ in 0..<10 { await Task.yield() } XCTAssertNil(env.coordinator.interactionState.activeComposeSession) - if case .idle = env.coordinator.state { - // expected - } else { - XCTFail("Expected idle state after empty compose result, got \(env.coordinator.state)") - } } - func test_engineFailure_surfacesFailedState() async { - let generateExpectation = expectation(description: "generateCompose called") + // MARK: - Failure paths + + func test_engineFailure_surfacesFailedStateAndClearsSession() async { + let finishedExpectation = expectation(description: "stream finished") let env = makeEnvironment( mode: .compose, - engineBehavior: .failure(SuggestionClientError.unavailable("Compose Mode requires tabby-depth-1.")), - onGenerateCalled: { generateExpectation.fulfill() } + engineBehavior: .failure(SuggestionClientError.unavailable("No local model loaded.")), + onStreamFinished: { finishedExpectation.fulfill() } ) _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [generateExpectation], timeout: 1.0) + await fulfillment(of: [finishedExpectation], timeout: 1.0) await Task.yield() if case .failed(let message) = env.coordinator.state { - XCTAssertTrue(message.contains("tabby-depth-1")) + XCTAssertTrue(message.contains("local model")) } else { XCTFail("Expected failed state after engine failure, got \(env.coordinator.state)") } + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } // MARK: - Test environment @@ -199,19 +203,15 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { private func makeEnvironment( mode: SuggestionInteractionMode, engineBehavior: FakeSuggestionEngine.Behavior, - onGenerateCalled: (() -> Void)? = nil, - onShowComposePreview: (() -> Void)? = nil, - onTypeDraftCalled: (() -> Void)? = nil + onStreamFinished: (() -> Void)? = nil ) -> ComposeTestEnvironment { let permissions = FakeSuggestionPermissions() let focusModel = FakeFocusModel(initialSnapshot: focusSnapshot()) let inputMonitor = FakeInputMonitor() let overlayController = FakeOverlayController() - overlayController.onShowComposePreview = onShowComposePreview let inserter = FakeSuggestionInserter() - inserter.onTypeDraftCalled = onTypeDraftCalled let engine = FakeSuggestionEngine(behavior: engineBehavior) - engine.onComposeCalled = onGenerateCalled + engine.onStreamFinished = onStreamFinished let settings = FakeSuggestionSettings(initialSnapshot: snapshot(mode: mode)) let clipboard = FakeClipboardContextProvider() let visualContext = FakeVisualContextCoordinator() @@ -252,21 +252,6 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { ) } - /// Spins up an environment, drives a first Tab through it, and waits for the compose preview - /// to be ready. Cancellation tests use this so they only have to assert post-conditions. - private func makeEnvironmentWithReadyDraft(draftText: String) async -> ComposeTestEnvironment { - let showPreviewExpectation = expectation(description: "showComposePreview called") - let env = makeEnvironment( - mode: .compose, - engineBehavior: .success(composeResult(text: draftText)), - onShowComposePreview: { showPreviewExpectation.fulfill() } - ) - - _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) - await fulfillment(of: [showPreviewExpectation], timeout: 1.0) - return env - } - private func snapshot(mode: SuggestionInteractionMode) -> SuggestionSettingsSnapshot { SuggestionSettingsSnapshot( isGloballyEnabled: true, @@ -300,10 +285,6 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { ) } - private func composeResult(text: String) -> ComposeResult { - ComposeResult(generation: 1, rawText: text, text: text, latency: 0.05) - } - private func isolatedUserDefaults() -> UserDefaults { let suiteName = "SuggestionCoordinatorComposeTests-\(UUID().uuidString)" guard let userDefaults = UserDefaults(suiteName: suiteName) else { @@ -363,7 +344,6 @@ private final class FakeInputMonitor: SuggestionInputMonitoring { private final class FakeOverlayController: SuggestionOverlayControlling { var state: OverlayState = .hidden(reason: "test idle") var onStateChange: ((OverlayState) -> Void)? - var onShowComposePreview: (() -> Void)? private(set) var composePreviewText: String? private(set) var hideReasons: [String] = [] @@ -376,7 +356,6 @@ private final class FakeOverlayController: SuggestionOverlayControlling { composePreviewText = text state = .composePreview(text: text, geometry: geometry) onStateChange?(state) - onShowComposePreview?() } func hide(reason: String) { @@ -389,28 +368,33 @@ private final class FakeOverlayController: SuggestionOverlayControlling { @MainActor private final class FakeSuggestionInserter: SuggestionInserting { var lastErrorMessage: String? - var onTypeDraftCalled: (() -> Void)? - private(set) var typedDrafts: [String] = [] + private(set) var insertedPieces: [String] = [] - func insert(_ suggestion: String) -> Bool { true } + func insert(_ suggestion: String) -> Bool { + insertedPieces.append(suggestion) + return true + } func typeDraft(_ draft: String, shouldContinue: @escaping @MainActor () -> Bool) async -> Bool { - typedDrafts.append(draft) - onTypeDraftCalled?() - return true + true } } @MainActor private final class FakeSuggestionEngine: SuggestionGenerating { enum Behavior { - case success(ComposeResult) + /// Yields all pieces then finishes immediately. + case success(pieces: [String]) + /// Yields `initialPieces`, then parks the stream until the underlying task is cancelled. + /// Lets tests trigger cancellation events (Esc, focus change, etc.) and verify cleanup. + case blocked(initialPieces: [String]) + /// Throws immediately. Tests use `onStreamFinished` to wait for the failure to propagate. case failure(Error) } private let behavior: Behavior - var onComposeCalled: (() -> Void)? - private(set) var composeCallCount = 0 + var onStreamFinished: (() -> Void)? + private(set) var composeStreamCallCount = 0 init(behavior: Behavior) { self.behavior = behavior @@ -421,18 +405,45 @@ private final class FakeSuggestionEngine: SuggestionGenerating { } func generateCompose(for request: ComposeRequest) async throws -> ComposeResult { - composeCallCount += 1 - onComposeCalled?() - switch behavior { - case .success(let result): - return ComposeResult( - generation: request.generation, - rawText: result.rawText, - text: result.text, - latency: result.latency - ) - case .failure(let error): - throw error + throw SuggestionClientError.unavailable("Compose Mode streams here; one-shot path unused.") + } + + func generateComposeStreaming(for request: ComposeRequest) async throws -> AsyncThrowingStream { + composeStreamCallCount += 1 + let behavior = self.behavior + let onStreamFinished = self.onStreamFinished + + return AsyncThrowingStream { continuation in + let task = Task.detached { + switch behavior { + case .success(let pieces): + for piece in pieces { + if Task.isCancelled { break } + continuation.yield(piece) + } + continuation.finish() + onStreamFinished?() + + case .blocked(let initialPieces): + for piece in initialPieces { + if Task.isCancelled { break } + continuation.yield(piece) + } + // Park until the underlying detached task is cancelled by the consumer's + // termination handler. Sleeping in small slices keeps `Task.isCancelled` + // responsive without burning CPU. + while !Task.isCancelled { + try? await Task.sleep(nanoseconds: 10_000_000) + } + continuation.finish() + onStreamFinished?() + + case .failure(let error): + continuation.finish(throwing: error) + onStreamFinished?() + } + } + continuation.onTermination = { _ in task.cancel() } } } From bc476d9cda47162012c89477e45eb32de50e533d Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Mon, 25 May 2026 01:27:47 -0700 Subject: [PATCH 4/4] =?UTF-8?q?Show=20a=20"Drafting=E2=80=A6"=20indicator?= =?UTF-8?q?=20at=20the=20caret=20while=20Compose=20warms=20up?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The context walk plus model warm-up can take a beat before the first streamed token lands. With nothing on screen during that window, Tab felt like it did nothing. This adds a small "Drafting…" pill at the caret the instant Tab registers; it disappears as soon as the first token starts landing in the field. - New `OverlayState.composeProgress(label:geometry:)` plus `SuggestionOverlayControlling.showComposeProgress`. - `OverlayController` renders a compact `ComposeProgressView` (spinner + label) in a capsule, sized as a status chip rather than the larger preview box. - `SuggestionOverlayPresenter.presentComposeProgress` is idempotent so re-presenting the same label/geometry does not churn the overlay log. - Coordinator shows the pill when `startComposeGeneration` flips to `.generating`, and hides it on the first non-empty streamed piece. Cancellation / completion already hide the overlay. Tests: indicator appears before the first token and the overlay ends hidden after the stream completes. 311 tests, 0 failures. --- .../SuggestionCoordinator+Compose.swift | 29 +++++++++ Cotabby/Models/SuggestionModels.swift | 13 +++- .../Models/SuggestionSubsystemContracts.swift | 1 + .../SuggestionOverlayPresenter.swift | 14 +++++ Cotabby/Services/UI/OverlayController.swift | 61 +++++++++++++++++++ .../SuggestionCoordinatorComposeTests.swift | 47 ++++++++++++++ CotabbyTests/SuggestionStateHelperTests.swift | 8 +++ 7 files changed, 170 insertions(+), 3 deletions(-) diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift index 1a9739e..4be2d9a 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift @@ -89,6 +89,9 @@ extension SuggestionCoordinator { latestGenerationNumber = context.generation latestRawModelOutput = nil state = .generating + // Instant feedback: the context walk + model warm-up can take a beat before the first + // token lands, so show a "Drafting…" pill at the caret the moment Tab registers. + presentComposeProgress(label: "Drafting…", for: context) logStage( "compose-streaming-start", workID: workID, @@ -174,6 +177,12 @@ extension SuggestionCoordinator { guard composeStreamShouldContinue(matching: session) else { break } guard !piece.isEmpty else { continue } + // First real token: tear down the "Drafting…" pill — the field is now showing + // the draft itself, so the status indicator would just overlap real content. + if accumulatedText.isEmpty { + hideOverlay(reason: "Compose draft started streaming into the field.") + } + accumulatedText += piece latestRawModelOutput = SuggestionDebugLogger.debugPreview(accumulatedText) _ = suggestionInserter.insert(piece) @@ -316,4 +325,24 @@ extension SuggestionCoordinator { return false } } + + // MARK: - Overlay + + /// Shows the transient "Drafting…" pill at the caret while we wait on the first token. + private func presentComposeProgress(label: String, for context: FocusedInputContext) { + let geometry = SuggestionOverlayGeometry( + caretRect: context.caretRect, + inputFrameRect: context.inputFrameRect, + caretQuality: context.caretQuality, + observedCharWidth: context.observedCharWidth, + isRightToLeft: TextDirectionDetector.isRightToLeft(context.precedingText) + ) + if let message = overlayPresenter.presentComposeProgress( + label: label, + geometry: geometry, + previousState: overlayState + ) { + latestOverlayMessage = message + } + } } diff --git a/Cotabby/Models/SuggestionModels.swift b/Cotabby/Models/SuggestionModels.swift index 739897e..fc2fc15 100644 --- a/Cotabby/Models/SuggestionModels.swift +++ b/Cotabby/Models/SuggestionModels.swift @@ -403,12 +403,16 @@ enum OverlayState: Equatable { case hidden(reason: String) case visible(text: String, geometry: SuggestionOverlayGeometry) case composePreview(text: String, geometry: SuggestionOverlayGeometry) + /// Transient "Drafting…" indicator shown near the caret while Compose is gathering context and + /// waiting on the first streamed token. It disappears as soon as text starts landing in the + /// field, so it exists only to prove to the user that the Tab press registered. + case composeProgress(label: String, geometry: SuggestionOverlayGeometry) var shortLabel: String { switch self { case .hidden: return "Hidden" - case .visible, .composePreview: + case .visible, .composePreview, .composeProgress: return "Visible" } } @@ -424,12 +428,15 @@ enum OverlayState: Equatable { case let .composePreview(text, geometry): return "Showing Compose preview with \(text.count) characters near " + "(\(Int(geometry.caretRect.minX)), \(Int(geometry.caretRect.minY)))." + case let .composeProgress(label, geometry): + return "Showing Compose \"\(label)\" indicator near " + + "(\(Int(geometry.caretRect.minX)), \(Int(geometry.caretRect.minY)))." } } var isVisible: Bool { switch self { - case .visible, .composePreview: + case .visible, .composePreview, .composeProgress: return true case .hidden: return false @@ -440,7 +447,7 @@ enum OverlayState: Equatable { switch self { case let .visible(text, _), let .composePreview(text, _): return text - case .hidden: + case .composeProgress, .hidden: return nil } } diff --git a/Cotabby/Models/SuggestionSubsystemContracts.swift b/Cotabby/Models/SuggestionSubsystemContracts.swift index f4e1a2d..883457c 100644 --- a/Cotabby/Models/SuggestionSubsystemContracts.swift +++ b/Cotabby/Models/SuggestionSubsystemContracts.swift @@ -106,6 +106,7 @@ protocol SuggestionOverlayControlling: AnyObject { func showSuggestion(_ text: String, geometry: SuggestionOverlayGeometry) func showComposePreview(_ text: String, geometry: SuggestionOverlayGeometry) + func showComposeProgress(_ label: String, geometry: SuggestionOverlayGeometry) func hide(reason: String) } diff --git a/Cotabby/Services/Suggestion/SuggestionOverlayPresenter.swift b/Cotabby/Services/Suggestion/SuggestionOverlayPresenter.swift index fe968dc..d1a741c 100644 --- a/Cotabby/Services/Suggestion/SuggestionOverlayPresenter.swift +++ b/Cotabby/Services/Suggestion/SuggestionOverlayPresenter.swift @@ -75,6 +75,20 @@ struct SuggestionOverlayPresenter { return "Displayed Compose draft preview near the caret." } + /// Shows the transient "Drafting…" indicator. Idempotent: re-presenting the same label and + /// geometry returns `nil` so the coordinator does not log a redundant overlay change. + func presentComposeProgress( + label: String, + geometry: SuggestionOverlayGeometry, + previousState: OverlayState + ) -> String? { + guard previousState != .composeProgress(label: label, geometry: geometry) else { + return nil + } + overlayController.showComposeProgress(label, geometry: geometry) + return "Displayed Compose \"\(label)\" indicator near the caret." + } + func hide(reason: String) -> String { overlayController.hide(reason: reason) return reason diff --git a/Cotabby/Services/UI/OverlayController.swift b/Cotabby/Services/UI/OverlayController.swift index bd131f0..e6509af 100644 --- a/Cotabby/Services/UI/OverlayController.swift +++ b/Cotabby/Services/UI/OverlayController.swift @@ -147,6 +147,43 @@ final class OverlayController: SuggestionOverlayControlling { state = .composePreview(text: previewText, geometry: geometry) } + /// Shows the small "Drafting…" pill near the caret while Compose waits on its first token. + /// Sized tighter than the preview box so it reads as a status chip, not content. + func showComposeProgress(_ label: String, geometry: SuggestionOverlayGeometry) { + let contentView: NSHostingView + let rootView = AnyView(ComposeProgressView(label: label)) + if let existing = hostingView { + existing.rootView = rootView + contentView = existing + } else { + let fresh = NSHostingView(rootView: rootView) + hostingView = fresh + panel.contentView = fresh + contentView = fresh + } + contentView.layoutSubtreeIfNeeded() + + let visibleFrame = targetScreenVisibleFrame(for: geometry.caretRect) + let contentSize = contentView.fittingSize + let width = min(max(contentSize.width, 96), min(220, visibleFrame.width - 32)) + let height = min(max(contentSize.height, 28), 48) + let originX = min( + max(geometry.caretRect.maxX + 8, visibleFrame.minX + 16), + visibleFrame.maxX - width - 16 + ) + let preferredOriginY = geometry.caretRect.minY - height - 8 + let originY = preferredOriginY >= visibleFrame.minY + 16 + ? preferredOriginY + : min(geometry.caretRect.maxY + 8, visibleFrame.maxY - height - 16) + + panel.setFrame( + CGRect(x: originX, y: originY, width: width, height: height).integral, + display: true + ) + panel.orderFrontRegardless() + state = .composeProgress(label: label, geometry: geometry) + } + /// Hides the floating panel and records why the overlay is no longer visible. func hide(reason: String) { panel.orderOut(nil) @@ -275,6 +312,30 @@ private struct ComposePreviewView: View { } } +/// Compact "Drafting…" status pill shown while Compose waits on its first streamed token. +private struct ComposeProgressView: View { + let label: String + + var body: some View { + HStack(spacing: 7) { + ProgressView() + .controlSize(.small) + .scaleEffect(0.7) + + Text(label) + .font(.system(size: 12, weight: .medium, design: .rounded)) + .foregroundStyle(.secondary) + .lineLimit(1) + .fixedSize(horizontal: true, vertical: true) + } + .padding(.horizontal, 11) + .padding(.vertical, 6) + .background(.regularMaterial, in: Capsule()) + .overlay(Capsule().stroke(.quaternary, lineWidth: 1)) + .fixedSize(horizontal: true, vertical: true) + } +} + /// Visual hint that teaches the user which key accepts the suggestion. private struct GhostTabKeycap: View { @Environment(\.colorScheme) var colorScheme diff --git a/CotabbyTests/SuggestionCoordinatorComposeTests.swift b/CotabbyTests/SuggestionCoordinatorComposeTests.swift index 7dabc46..add941e 100644 --- a/CotabbyTests/SuggestionCoordinatorComposeTests.swift +++ b/CotabbyTests/SuggestionCoordinatorComposeTests.swift @@ -38,6 +38,45 @@ final class SuggestionCoordinatorComposeTests: XCTestCase { XCTAssertNil(env.coordinator.interactionState.activeComposeSession) } + func test_firstTab_showsDraftingIndicatorThenHidesItOnceTokensArrive() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: []) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<5 { await Task.yield() } + + // Before any token, the "Drafting…" pill is visible. + XCTAssertEqual(env.overlayController.composeProgressLabels, ["Drafting…"]) + if case .composeProgress = env.overlayController.state { + // expected + } else { + XCTFail("Expected composeProgress overlay state, got \(env.overlayController.state)") + } + } + + func test_draftingIndicatorIsHiddenOnceTheFirstTokenLands() async { + let finishedExpectation = expectation(description: "stream finished") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .success(pieces: ["Hello"]), + onStreamFinished: { finishedExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [finishedExpectation], timeout: 1.0) + await Task.yield() + + XCTAssertEqual(env.overlayController.composeProgressLabels, ["Drafting…"]) + // After streaming completes the overlay is hidden, not stuck on the indicator. + if case .hidden = env.overlayController.state { + // expected + } else { + XCTFail("Expected hidden overlay after stream, got \(env.overlayController.state)") + } + } + func test_firstTab_inComposeMode_returnsTrueToConsumeTab() { let env = makeEnvironment( mode: .compose, @@ -358,6 +397,14 @@ private final class FakeOverlayController: SuggestionOverlayControlling { onStateChange?(state) } + private(set) var composeProgressLabels: [String] = [] + + func showComposeProgress(_ label: String, geometry: SuggestionOverlayGeometry) { + composeProgressLabels.append(label) + state = .composeProgress(label: label, geometry: geometry) + onStateChange?(state) + } + func hide(reason: String) { hideReasons.append(reason) state = .hidden(reason: reason) diff --git a/CotabbyTests/SuggestionStateHelperTests.swift b/CotabbyTests/SuggestionStateHelperTests.swift index e1ae9bd..433c247 100644 --- a/CotabbyTests/SuggestionStateHelperTests.swift +++ b/CotabbyTests/SuggestionStateHelperTests.swift @@ -567,6 +567,14 @@ private final class FakeOverlayController: SuggestionOverlayControlling { onStateChange?(state) } + func showComposeProgress( + _ label: String, + geometry: SuggestionOverlayGeometry + ) { + state = .composeProgress(label: label, geometry: geometry) + onStateChange?(state) + } + func hide(reason: String) { hideReasons.append(reason) state = .hidden(reason: reason)