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..4be2d9a --- /dev/null +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Compose.swift @@ -0,0 +1,348 @@ +import CoreGraphics +import Foundation +import Logging + +/// File overview: +/// Compose Mode entry points for `SuggestionCoordinator`. +/// +/// 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 + + /// 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: + // 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() + + case .dismissal: + cancelComposeWork(reason: "Compose cancelled by Escape.") + return false + + case .navigation: + // 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: + // 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.") + } + return false + + case .other: + return false + } + } + + // MARK: - Generation + + /// Streams a Compose draft into the currently focused field. Returns `true` to consume the Tab. + @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?.runComposeStreaming(for: context, workID: workID) + } + 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, + generation: context.generation, + message: "Gathering Compose context for \(context.elementIdentifier) in \(context.applicationName)." + ) + return true + } + + private func runComposeStreaming(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 + 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 } + await self.consumeComposeStream( + request: request, + workID: workID, + initialSession: initialSession + ) + } + } + + private func consumeComposeStream( + request: ComposeRequest, + workID: UInt64, + initialSession: ActiveComposeSession + ) async { + let startTime = Date() + var accumulatedText = "" + var session = initialSession + + 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 } + + // 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) + 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, + session: session, + outcome: ComposeStreamOutcome( + stage: "compose-streaming-cancelled", + stageMessage: "Compose stream cancelled." + ) + ) + return + } catch { + await applyComposeFailure(error.localizedDescription, workID: workID) + return + } + + 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( + outcome.stage, + workID: workID, + generation: session.baseContext.generation, + message: outcome.stageMessage, + normalizedOutput: accumulated + ) + } + + private func applyComposeFailure(_ message: String, workID: UInt64) async { + guard workController.isCurrent(workID) else { return } + if let session = interactionState.activeComposeSession { + interactionState.clearComposeSession(session) + } else { + interactionState.clearSuggestion() + } + hideOverlay(reason: "Overlay hidden because Compose generation failed.") + state = .failed(message) + logStage( + "compose-failed", + workID: workID, + generation: latestGenerationNumber, + message: message + ) + } + + /// 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 + } + 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. 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 + 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 streaming task could still emit output. + var isAnyComposeWorkInFlight: Bool { + switch state { + case .generating, .typing: + return true + case .idle, .disabled, .debouncing, .ready, .failed: + 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/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/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/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/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 f2ef2a4..883457c 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 } @@ -88,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) } @@ -102,3 +121,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/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 81c26fb..a5db439 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, @@ -183,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 e93350c..a7c7ed8 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) @@ -120,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/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/Cotabby/UI/MenuBarView.swift b/Cotabby/UI/MenuBarView.swift index 9650440..d7a9a15 100644 --- a/Cotabby/UI/MenuBarView.swift +++ b/Cotabby/UI/MenuBarView.swift @@ -119,6 +119,13 @@ 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.selectedEngine.supportsLocalModelManagement { modelRow } diff --git a/Cotabby/UI/SettingsView.swift b/Cotabby/UI/SettingsView.swift index a42a576..7794ad6 100644 --- a/Cotabby/UI/SettingsView.swift +++ b/Cotabby/UI/SettingsView.swift @@ -175,6 +175,13 @@ 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) + } + 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..add941e --- /dev/null +++ b/CotabbyTests/SuggestionCoordinatorComposeTests.swift @@ -0,0 +1,560 @@ +import Combine +import CoreGraphics +import Foundation +import XCTest +@testable import Cotabby + +/// Tests `SuggestionCoordinator+Compose` end to end with fakes for every collaborator. +/// +/// 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] = [] + + override func tearDown() async throws { + Self.retainedCoordinators.removeAll() + try await super.tearDown() + } + + // MARK: - First Tab: streaming starts + + 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(pieces: pieces), + onStreamFinished: { finishedExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [finishedExpectation], timeout: 1.0) + await Task.yield() + + XCTAssertEqual(env.engine.composeStreamCallCount, 1) + XCTAssertEqual(env.inserter.insertedPieces, pieces) + 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, + engineBehavior: .success(pieces: ["only"]) + ) + + let consumed = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + + XCTAssertTrue(consumed) + } + + func test_firstTab_inAutocompleteMode_doesNotInvokeComposeStream() async { + let env = makeEnvironment( + mode: .autocomplete, + engineBehavior: .success(pieces: ["should not run"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await Task.yield() + await Task.yield() + + XCTAssertEqual(env.engine.composeStreamCallCount, 0) + XCTAssertEqual(env.inserter.insertedPieces, []) + } + + // MARK: - Subsequent Tabs while streaming + + func test_secondTab_whileStreaming_isAbsorbedWithoutRestartingTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["partial"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<10 { await Task.yield() } + + let consumedSecondTab = env.coordinator.handleInputEvent( + CotabbyTestFixtures.inputEvent(kind: .acceptance) + ) + + 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_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) + } + + 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_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_focusChangeDuringStreaming_cancelsTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["before"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + for _ in 0..<15 { await Task.yield() } + + env.focusModel.publish(snapshot: focusSnapshot(processIdentifier: 999, elementIdentifier: "different")) + for _ in 0..<10 { await Task.yield() } + + XCTAssertNil(env.coordinator.interactionState.activeComposeSession) + } + + func test_modeChangeToAutocomplete_duringStreaming_cancelsTheStream() async { + let env = makeEnvironment( + mode: .compose, + engineBehavior: .blocked(initialPieces: ["before"]) + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + 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) + } + + // MARK: - Failure paths + + func test_engineFailure_surfacesFailedStateAndClearsSession() async { + let finishedExpectation = expectation(description: "stream finished") + let env = makeEnvironment( + mode: .compose, + engineBehavior: .failure(SuggestionClientError.unavailable("No local model loaded.")), + onStreamFinished: { finishedExpectation.fulfill() } + ) + + _ = env.coordinator.handleInputEvent(CotabbyTestFixtures.inputEvent(kind: .acceptance)) + await fulfillment(of: [finishedExpectation], timeout: 1.0) + await Task.yield() + + if case .failed(let message) = env.coordinator.state { + 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 + + 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, + onStreamFinished: (() -> Void)? = nil + ) -> ComposeTestEnvironment { + let permissions = FakeSuggestionPermissions() + let focusModel = FakeFocusModel(initialSnapshot: focusSnapshot()) + let inputMonitor = FakeInputMonitor() + let overlayController = FakeOverlayController() + let inserter = FakeSuggestionInserter() + let engine = FakeSuggestionEngine(behavior: engineBehavior) + engine.onStreamFinished = onStreamFinished + 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 + ) + } + + 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 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)? + 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) + } + + 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) + onStateChange?(state) + } +} + +@MainActor +private final class FakeSuggestionInserter: SuggestionInserting { + var lastErrorMessage: String? + private(set) var insertedPieces: [String] = [] + + func insert(_ suggestion: String) -> Bool { + insertedPieces.append(suggestion) + return true + } + + func typeDraft(_ draft: String, shouldContinue: @escaping @MainActor () -> Bool) async -> Bool { + true + } +} + +@MainActor +private final class FakeSuggestionEngine: SuggestionGenerating { + enum Behavior { + /// 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 onStreamFinished: (() -> Void)? + private(set) var composeStreamCallCount = 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 { + 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() } + } + } + + 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 + ) + } +} 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)