Keep the KV cache on cancelled llama generations #513
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } | ||
|
Comment on lines
+28
to
+36
Contributor
There was a problem hiding this comment.
|
||
|
|
||
| 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<String, Error> = .success("") | ||
| private(set) var resetCount = 0 | ||
|
|
||
| func generate( | ||
| prompt: String, | ||
| cachedPrefixBytes: Int?, | ||
| options: LlamaGenerationOptions | ||
| ) async throws -> String { | ||
| try generateResult.get() | ||
| } | ||
|
|
||
| func resetPromptCache() { | ||
| resetCount += 1 | ||
| } | ||
| } | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
SuggestionSubsystemContracts.swiftdeclares itself in its file header as defining contracts thatSuggestionCoordinatordepends on, butLlamaRuntimeGeneratingis consumed exclusively byLlamaSuggestionEngine. In isolation this isn't a bug, but as this file grows it becomes less clear which protocols belong to the coordinator boundary and which are internal seams elsewhere. A dedicatedLlamaRuntimeContracts.swift(alongside the existingRuntime/files) or a comment noting this protocol is for the engine layer would keep the boundary explicit for future readers.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!