diff --git a/Cotabby.xcodeproj/project.pbxproj b/Cotabby.xcodeproj/project.pbxproj index 654f718c..299be512 100644 --- a/Cotabby.xcodeproj/project.pbxproj +++ b/Cotabby.xcodeproj/project.pbxproj @@ -191,6 +191,7 @@ B93AB7E845086F6FBB068369 /* SuggestionRequestFactoryTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = EE94342B888A5A2CCF66BC93 /* SuggestionRequestFactoryTests.swift */; }; BB6325CA50F97B18B9725918 /* SuggestionTextNormalizer.swift in Sources */ = {isa = PBXBuildFile; fileRef = B424E2AC97C99D335B0D5751 /* SuggestionTextNormalizer.swift */; }; BBE22CE4EF43247F8775B25D /* FocusPollBackoff.swift in Sources */ = {isa = PBXBuildFile; fileRef = 09FADF683BE7B3558377FA76 /* FocusPollBackoff.swift */; }; + BE3CB85508055D159C35020A /* LlamaSuggestionEngineCancellationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABCC3FD99B1824A81E665F3 /* LlamaSuggestionEngineCancellationTests.swift */; }; BFCA7FAFDAEBF586AB615567 /* ClipboardRelevanceFilterTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 90B0D133AB77A2503FB08827 /* ClipboardRelevanceFilterTests.swift */; }; C0B833234748E82D3382631A /* emoji.json in Resources */ = {isa = PBXBuildFile; fileRef = C379D77029D6E88C8C1B9AF7 /* emoji.json */; }; C0FE11D76BDF01A5470C554D /* FocusCapabilityFlickerGate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6A44BEC8C23FF227731DD0CD /* FocusCapabilityFlickerGate.swift */; }; @@ -422,6 +423,7 @@ A863F41C0C03D7B4AC5DC002 /* MarkerSelectionSynthesizer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MarkerSelectionSynthesizer.swift; sourceTree = ""; }; A9199B9CEAB320982CA333B8 /* WelcomeTemplateStepView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WelcomeTemplateStepView.swift; sourceTree = ""; }; AA33F5FFAC5B99384E15CE3E /* BundledRuntimeLocator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BundledRuntimeLocator.swift; sourceTree = ""; }; + AABCC3FD99B1824A81E665F3 /* LlamaSuggestionEngineCancellationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LlamaSuggestionEngineCancellationTests.swift; sourceTree = ""; }; AC70775535A3428991025AB8 /* AXHelper.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AXHelper.swift; sourceTree = ""; }; AD752451330486FE270018B0 /* CustomRulesTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomRulesTests.swift; sourceTree = ""; }; AD9573F3504CAE6891DF9B7D /* AppUpdateManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppUpdateManager.swift; sourceTree = ""; }; @@ -764,6 +766,7 @@ 4793D4EA5D36D7E5CC216C27 /* LanguageSupportTests.swift */, 5807E8508D9355D0271A00C5 /* LaunchAtLoginStateTests.swift */, 0CA88BB29BC8727878C99E95 /* LlamaPromptCacheHintTrackerTests.swift */, + AABCC3FD99B1824A81E665F3 /* LlamaSuggestionEngineCancellationTests.swift */, 52BAFA2F989C3C4F7FB892B5 /* MarkerSelectionSynthesizerTests.swift */, 1274F897631B1B3A835D157F /* MidWordContinuationPolicyTests.swift */, FC83D14A7557BC0196E59007 /* MirrorOverlayLayoutTests.swift */, @@ -1281,6 +1284,7 @@ E912D4617AE1376061DF1F00 /* LanguageSupportTests.swift in Sources */, E27E6377D36D4981301568DD /* LaunchAtLoginStateTests.swift in Sources */, E38801433B99E65BD7E45A0E /* LlamaPromptCacheHintTrackerTests.swift in Sources */, + BE3CB85508055D159C35020A /* LlamaSuggestionEngineCancellationTests.swift in Sources */, 87806DE08881D11F2608A13D /* MarkerSelectionSynthesizerTests.swift in Sources */, 7C36DBA762E19C8C31676D44 /* MidWordContinuationPolicyTests.swift in Sources */, 14D77F0B8A195AC2FA8D24A9 /* MirrorOverlayLayoutTests.swift in Sources */, diff --git a/Cotabby/Models/SuggestionSubsystemContracts.swift b/Cotabby/Models/SuggestionSubsystemContracts.swift index bb2c379f..c3116982 100644 --- a/Cotabby/Models/SuggestionSubsystemContracts.swift +++ b/Cotabby/Models/SuggestionSubsystemContracts.swift @@ -82,6 +82,17 @@ extension SuggestionGenerating { func prewarm(for request: SuggestionRequest) async {} } +/// Behavior-shaped view of the llama runtime that `LlamaSuggestionEngine` depends on: run one +/// generation and drop the native KV cache. Extracted so the engine's failure handling — in +/// particular the invariant that a *cancelled* generation must NOT reset the cache (resetting it on +/// every superseded keystroke was the base-model input-lag regression) — can be unit-tested against +/// a fake runtime instead of loading a real model. `LlamaRuntimeManager` is the production conformer. +@MainActor +protocol LlamaRuntimeGenerating: AnyObject { + func generate(prompt: String, cachedPrefixBytes: Int?, options: LlamaGenerationOptions) async throws -> String + func resetPromptCache() +} + @MainActor protocol SuggestionSettingsProviding: AnyObject { var snapshot: SuggestionSettingsSnapshot { get } diff --git a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift index 6d3e0797..be4d9917 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeManager.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeManager.swift @@ -325,3 +325,5 @@ final class LlamaRuntimeManager: ObservableObject { state = .ready("Loaded \(preparedRuntime.resolvedRuntime.modelDisplayName) in-process.") } } + +extension LlamaRuntimeManager: LlamaRuntimeGenerating {} diff --git a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift index 21b0c0fd..0f084dba 100644 --- a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift @@ -10,7 +10,7 @@ import Logging /// That separation matters because prompt strategy changes far more often than model lifecycle code. @MainActor final class LlamaSuggestionEngine { - private let runtimeManager: LlamaRuntimeManager + private let runtimeManager: LlamaRuntimeGenerating private var promptCacheHintTracker = LlamaPromptCacheHintTracker() /// UserDefaults key (no UI) that routes llama generation through the deterministic constrained @@ -22,7 +22,7 @@ final class LlamaSuggestionEngine { UserDefaults.standard.bool(forKey: constrainedDecoderDefaultsKey) } - init(runtimeManager: LlamaRuntimeManager) { + init(runtimeManager: LlamaRuntimeGenerating) { self.runtimeManager = runtimeManager } @@ -102,6 +102,22 @@ final class LlamaSuggestionEngine { } catch is CancellationError { CotabbyLogger.suggestion.debug("Llama generation cancelled", metadata: baseMetadata) throw SuggestionClientError.cancelled + } catch LlamaRuntimeError.cancelled { + // A cancelled generation is NOT a runtime failure, so it must not reset the KV cache. + // `LlamaRuntimeManager.generate` surfaces an outer-Task cancellation as + // `LlamaRuntimeError.cancelled` (its `catch is CancellationError` rethrows it so callers + // share one error vocabulary). Without this branch that case falls through to the generic + // `LlamaRuntimeError` handler below and wipes the native KV sequence on every cancel. + // + // During fast typing nearly every keystroke supersedes the previous in-flight generation, + // so that path fired ~twice a second — each time synchronously destroying the prompt KV on + // the main actor (contending with the keystroke-delivery run loop) and forcing the next + // keystroke to re-decode the whole prompt from scratch. The cooperative cancel inside + // `LlamaRuntimeCore.generate` already unwound cleanly (its KV-trim defer restored + // prompt-only state), so the cache is still valid and reusable. Route this to the same + // quiet path as `CancellationError` and leave the cache intact. + CotabbyLogger.suggestion.debug("Llama generation cancelled (runtime task)", metadata: baseMetadata) + throw SuggestionClientError.cancelled } catch let error as LlamaRuntimeError { CotabbyLogger.suggestion.error( "Llama runtime error, resetting cache: \(error.localizedDescription)", diff --git a/CotabbyTests/LlamaSuggestionEngineCancellationTests.swift b/CotabbyTests/LlamaSuggestionEngineCancellationTests.swift new file mode 100644 index 00000000..47c69a1a --- /dev/null +++ b/CotabbyTests/LlamaSuggestionEngineCancellationTests.swift @@ -0,0 +1,144 @@ +import CoreGraphics +import Foundation +import XCTest +@testable import Cotabby + +/// Regression tests for `LlamaSuggestionEngine`'s failure handling, guarding the input-lag fix: +/// a *cancelled* generation must be treated as a quiet cancellation, NOT as a runtime error that +/// wipes the native KV cache. During fast typing nearly every keystroke supersedes the in-flight +/// generation, so resetting the cache on each cancel (the base-model regression) fired ~twice a +/// second — synchronously destroying the prompt KV on the main actor and forcing a full prompt +/// re-decode on the next keystroke. These tests pin the routing for both cancellation shapes the +/// runtime can surface (`CancellationError` and `LlamaRuntimeError.cancelled`) and confirm genuine +/// runtime errors still reset. +@MainActor +final class LlamaSuggestionEngineCancellationTests: XCTestCase { + + func test_runtimeCancelledError_doesNotResetCache_andThrowsCancelled() async { + // `LlamaRuntimeManager.generate` surfaces an outer-Task cancellation as + // `LlamaRuntimeError.cancelled`. The engine must route that to the quiet cancel path. + let runtime = FakeLlamaRuntime() + runtime.generateResult = .failure(LlamaRuntimeError.cancelled) + let engine = LlamaSuggestionEngine(runtimeManager: runtime) + + await assertThrowsCancelled(engine) + XCTAssertEqual(runtime.resetCount, 0, "A cancelled generation must not reset the KV cache") + } + + func test_pureCancellationError_doesNotResetCache_andThrowsCancelled() async { + // Guards the pre-existing clean path so a future refactor cannot regress it either. + let runtime = FakeLlamaRuntime() + runtime.generateResult = .failure(CancellationError()) + let engine = LlamaSuggestionEngine(runtimeManager: runtime) + + await assertThrowsCancelled(engine) + XCTAssertEqual(runtime.resetCount, 0) + } + + func test_genuineRuntimeError_resetsCache_andThrowsUnavailable() async { + let runtime = FakeLlamaRuntime() + runtime.generateResult = .failure(LlamaRuntimeError.generationFailed("boom")) + let engine = LlamaSuggestionEngine(runtimeManager: runtime) + + do { + _ = try await engine.generateSuggestion(for: makeRequest(prompt: "hello")) + XCTFail("Expected a thrown error") + } catch SuggestionClientError.unavailable { + // Expected: a real runtime failure does reset and surfaces as unavailable. + } catch { + XCTFail("Expected SuggestionClientError.unavailable, got \(error)") + } + XCTAssertEqual(runtime.resetCount, 1, "A genuine runtime error should reset the KV cache exactly once") + } + + func test_successfulGeneration_doesNotResetCache() async throws { + let runtime = FakeLlamaRuntime() + runtime.generateResult = .success("world") + let engine = LlamaSuggestionEngine(runtimeManager: runtime) + + let result = try await engine.generateSuggestion(for: makeRequest(prompt: "hello ")) + + XCTAssertEqual(result.generation, 1) + XCTAssertEqual(runtime.resetCount, 0) + } + + // MARK: - Helpers + + private func assertThrowsCancelled( + _ engine: LlamaSuggestionEngine, + file: StaticString = #filePath, + line: UInt = #line + ) async { + do { + _ = try await engine.generateSuggestion(for: makeRequest(prompt: "hello")) + XCTFail("Expected a thrown error", file: file, line: line) + } catch SuggestionClientError.cancelled { + // Expected quiet cancellation. + } catch { + XCTFail("Expected SuggestionClientError.cancelled, got \(error)", file: file, line: line) + } + } + + private func makeRequest(prompt: String) -> SuggestionRequest { + let snapshot = FocusedInputSnapshot( + applicationName: "TestApp", + bundleIdentifier: "com.example.TestApp", + processIdentifier: 123, + elementIdentifier: "field", + role: "AXTextField", + subrole: nil, + caretRect: .zero, + inputFrameRect: nil, + caretSource: "test", + caretQuality: .exact, + observedCharWidth: nil, + precedingText: prompt, + trailingText: "", + selection: NSRange(location: prompt.count, length: 0), + isSecure: false + ) + let context = FocusedInputContext(snapshot: snapshot, generation: 1) + + return SuggestionRequest( + context: context, + prefixText: prompt, + prompt: prompt, + generation: context.generation, + maxPredictionTokens: 8, + temperature: 0.1, + topK: 20, + topP: 0.7, + minP: 0.08, + repetitionPenalty: 1.05, + randomSeed: 42, + maxSuffixCharacters: 192, + completionLengthInstruction: "Return only the next few words.", + userName: nil, + customRules: [], + languageInstruction: nil, + clipboardContext: nil, + visualContextSummary: nil, + isMultiLineEnabled: false + ) + } +} + +/// Minimal `LlamaRuntimeGenerating` fake that returns a staged result and counts cache resets, +/// so the engine's failure routing can be exercised without loading a real model. +@MainActor +private final class FakeLlamaRuntime: LlamaRuntimeGenerating { + var generateResult: Result = .success("") + private(set) var resetCount = 0 + + func generate( + prompt: String, + cachedPrefixBytes: Int?, + options: LlamaGenerationOptions + ) async throws -> String { + try generateResult.get() + } + + func resetPromptCache() { + resetCount += 1 + } +}