Summary

Cotabby's Apple Foundation Models path currently constructs a new LanguageModelSession on every keystroke, which throws away Apple's instruction-prefix cache and pays full instruction tokenization on each request. This change holds the session as long as the rendered instructions are unchanged, and asks the engine to prewarm() after focus arrives on a new editable surface so weight loading and prefix tokenization happen before the user types instead of inside the critical path of the first respond.

Stacked on eval/fm-expand-driftset — verify with that PR's expanded eval suite once both land.

Validation

xcodebuild -project Cotabby.xcodeproj -scheme Cotabby -destination 'platform=macOS' build CODE_SIGNING_ALLOWED=NO
→ ** BUILD SUCCEEDED **

xcodebuild -project Cotabby.xcodeproj -scheme Cotabby -destination 'platform=macOS' build-for-testing CODE_SIGNING_ALLOWED=NO
→ ** TEST BUILD SUCCEEDED **

xcodebuild test ... -only-testing:CotabbyTests/FoundationModelPromptRendererTests -only-testing:CotabbyTests/SuggestionEngineRouterTests
→ Executed 9 tests, with 0 failures (0 unexpected)

swiftlint lint --quiet → no violations.

End-to-end FM perf check expected to be done with the eval suite from the parent PR.

Linked issues

Refs the FM-quality investigation.

Risk / rollout notes

• SuggestionGenerating gained prewarm(for:). A default no-op extension keeps every existing conformer (llama engine, router, unavailable engine, the StubSuggestionEngine in tests) source-compatible. Only FM provides a real implementation.
• The cached session is keyed on the rendered instructions string. Any user-driven change that alters instructions (language override, custom rules) flows through FoundationModelPromptRenderer.sessionInstructions(for:), so its content drives invalidation correctly without an extra cache-busting hook.
• resetCachedGenerationContext continues to clear the cache, so a field change or settings edit still produces a fresh session next request. The new prewarm Task awaits the existing cache-reset barrier (awaitCachedGenerationContextResetIfNeeded) so prewarm always runs after the reset and cannot have its session nulled out from under it.
• Prewarm is gated on a real field-change event (hasFocusedElementChanged), not on every focus poll, so the call site does not thrash on noisy AX updates. The engine internally skips work when the session is already cached, so duplicate calls are cheap regardless.

Greptile Summary

This PR introduces two related optimisations to the Apple Foundation Models path: a bounded LanguageModelSession cache (keyed on rendered instructions, guarded by !isResponding and a pristine-transcript check) so the prewarmed session is handed to the first real keystroke request, and a prewarm(for:) hook wired into the focus-change path so weight loading and instruction tokenisation happen before the user types rather than inside respond's critical path.

• Session caching (ensureSession): reuse is correctly limited to one request — the three-way predicate addresses the concurrentRequests and exceededContextWindowSize failure modes flagged in prior reviews. The guard let self fix in prewarmEngineForCurrentField and the cache-reset barrier order are both handled correctly.
• Protocol extension: prewarm(for:) is added to SuggestionGenerating with a default no-op, keeping llama, router, stub, and unavailable engines source-compatible with zero changes.
• Streaming: respond is replaced by streamResponse with a cumulative-snapshot loop, enabling cooperative cancellation mid-decode; the didReceiveSnapshot guard correctly surfaces a zero-output stream as a generationFailed error.

Confidence Score: 4/5

Safe to merge. The session-reuse logic is well-guarded and the prewarm plumbing is correctly sequenced behind the cache-reset barrier.

The three-way ensureSession predicate correctly avoids both the concurrent-stream and transcript-accumulation failure modes. One open question — whether Apple's runtime updates transcript.count when a stream is cancelled mid-flight — means the "consumed once" bound is a documented assumption rather than a verified guarantee. This does not produce incorrect output in any observed path, but the behaviour under rapid keystroke cancellations relies on an undocumented Apple API contract.

Cotabby/Services/Runtime/FoundationModelSuggestionEngine.swift — specifically the ensureSession reuse predicate and the transcript.count == pristineTranscriptCount condition for cancelled streams.

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant Coord as SuggestionCoordinator
    participant Engine as FoundationModelSuggestionEngine
    participant Apple as LanguageModelSession (Apple)

    Note over Coord: User focuses new editable field
    Coord->>Coord: cancelPredictionWork()
    Coord->>Coord: resetCachedGenerationContext() [schedules Task]
    Coord->>Coord: prewarmEngineForCurrentField() [schedules Task]

    Note over Coord: Reset Task runs
    Coord->>Engine: resetCachedGenerationContext()
    Engine->>Engine: "cachedSession = nil"

    Note over Coord: Prewarm Task runs (after barrier)
    Coord->>Coord: awaitCachedGenerationContextResetIfNeeded()
    Coord->>Engine: "prewarm(for: request[gen=0])"
    Engine->>Engine: ensureSession() → new CachedSession
    Engine->>Apple: session.prewarm() [fire-and-forget]
    Apple-->>Apple: loads weights, primes prefix cache

    Note over Coord: User types (first keystroke)
    Coord->>Engine: "generateSuggestion(for: request[gen=N])"
    Engine->>Engine: ensureSession() → reuse cached
    Engine->>Apple: session.streamResponse(to: prompt)
    Apple-->>Engine: partial snapshots (cumulative)
    Engine-->>Coord: SuggestionResult

    Note over Engine: transcript.count > pristineTranscriptCount
    Note over Coord: User types again (second keystroke)
    Coord->>Engine: "generateSuggestion(for: request[gen=N+1])"
    Engine->>Engine: ensureSession() → transcript diverged → new session
    Engine->>Apple: new LanguageModelSession(...)
    Engine->>Apple: session.streamResponse(to: prompt)
    Apple-->>Engine: partial snapshots
    Engine-->>Coord: SuggestionResult

_{Reviews (4): Last reviewed commit: "Trigger CI after base change to main" | Re-trigger Greptile}