Add and harden AI provider cache checkpoints#641
Open
sroussey wants to merge 38 commits into
Open
Conversation
sroussey
force-pushed
the
ai-provider-cache-checkpoints
branch
from
July 17, 2026 15:59
d76e7f3 to
68458c8
Compare
sroussey
added a commit
that referenced
this pull request
Jul 21, 2026
…eckpoints-verify-xe6adz fix(ai,node-llama-cpp,google-gemini): five HIGH-severity cache-checkpoint correctness fixes (follow-up to #641)
…text Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
…he_control Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
… fallback Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
…atches consuming tokenization
…emit-only generation throws Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
…nt cache checkpoints
…ycle docs, emit-session cleanup Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
packages/workglow grew 157 -> 242 instantiations (+85 absolute) from the new checkpoint type exports flowing through the meta-package re-exports; small expected growth in packages/ai, packages/test, and the three touched providers. Also records providers/duckdb, which was absent from the baseline. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Correctness fixes from the branch code review: - HFT_ToolCalling / HFT_Chat: render the checkpoint prefix into the fed prompt (continuation) and attach prefix KV only under a prompt.startsWith(prefixPrompt) parity guard, falling back to a full re-encode — previously the prefix KV was attached to a prompt that never contained the warm-up rendering, positionally corrupting generation. - renderHftPrefixPrompt: no empty user turn for message-less prefixes, so the parity guard can actually hold for the common system+tools warm-up. - HFT fingerprint tool-session: warm only the shared tools+systemPrompt region instead of the full prompt, so a second tool task sharing the fingerprint no longer attaches a cache poisoned by the first task's turn. - LlamaCpp_TextGeneration: consuming a checkpoint takes sole ownership of the live session (map entry removed, disposed at turn end unless re-keyed for emit) — a live sequence mutates in place, so a second consumer or a kept parent previously saw the first consumer's tokens and two ids could alias (and double-dispose) one native sequence. - resolveCheckpointSession: gate checkpoint/emitCheckpoint on the provider actually serving cache.checkpoint — OpenAI-shaped providers previously ignored the prefix silently (context dropped, handles backed by nothing). - ToolCallingTask emit accumulator: upsert toolCalls by id (mirroring StreamProcessor) instead of replacing — OpenAI-shaped providers emit one single-element array per tool call, so parallel calls lost all but one. - Anthropic checkpoint params/replay: guard on the converted message array; an all-system prefix.messages converts to [] and crashed annotateLastBlock. - HFT emit-only checkpoints now attach an empty cache so the turn's KV is snapshotted under the emitted id instead of discarded. Cleanups: shared validateParentCheckpoint/mergeCheckpointPrefix helpers (CacheCheckpointTask no longer duplicates them), redundant _parent field removed, stale _computedSessionId cleared between reused-instance runs, upfront getJobInput gated to checkpoint runs in TextGenerationTask, disposeSession idempotency contract documented, and the new checkpoint interfaces use explicit '| undefined' optionals. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
OpenAI maps checkpoints onto its automatic prompt cache: the warm-up run-fn sends the prefix once (Responses API, minimal output) and consumers replay the prefix content ahead of their tail via mergeOpenAICheckpointPrefix — the derived prompt_cache_key (model + instructions + tools) aligns warm-up and consumers without coordination. Gemini maps checkpoints to explicit server-side CachedContent with a 1h TTL: the warm-up creates the cache (degrading gracefully below the minimum cacheable size), consumers reference it with tail-only requests when the call adds no system prompt or non-default tool choice (the API rejects systemInstruction/tools alongside cachedContent) and replay the prefix inline otherwise. The id → cache-name store lives in an SDK-free module so both provider shells can delete the cache eagerly on disposeSession (TTL is the backstop). Also advertises cache.checkpoint in capability inference for supporting Anthropic / OpenAI / Gemini model families — previously only inline model configs with hand-declared capabilities passed the task-level gate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
A checkpoint fed to AiChatTask must never make the chat slower than no checkpoint. Previously run-fns inferred immutable-checkpoint semantics from the bare presence of session.prefix, which disabled HFT's per-turn KV snapshotting for the chat's own session — every turn re-encoded the growing conversation. AiSessionContext gains ownedSession: the sessionId is the caller's own mutable session merely seeded from the prefix. AiChatTask sets it; HFT_Chat keys immutability (snapshot target, supersede delete) on prefix && !ownedSession so progressive snapshotting stays alive; LlamaCpp_Chat labels the seeded session progressive and LlamaCpp_TextGeneration never applies take-ownership stealing to an owned session. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DsBvrmfhe1aY63cjztn1WW
Replace the unreliable vi.mock("@google/genai") in the Gemini checkpoint
params test with a runtime test seam (_testOnly.setGeminiClientForTests).
Module-level SDK mocks cannot intercept the provider here: the workspace
ships several @google/genai copies and the provider resolves it from a
bundled dist file the mock never reaches. The seam works identically for
src and dist and captures the requests the run-fns build without a live
call. Surface it through both the /ai and /ai-runtime _testOnly barrels
since dist splits them into separate bundles with independent module state.
Gate the worker-boundary session-dispose case with it.skipIf(!isBun): it
spawns a real .ts worker that needs Bun's native TS-worker execution and
bun:test's mock.module. It runs under `bun test`; the sibling in-process
case covers the dispose path under vitest.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The prior renderer flattened a checkpoint prefix to raw "role: text" strings
via renderLlamaCppPrefixText and preloaded that text — bypassing the model's
chat wrapper and silently dropping tool_use / tool_result blocks. Warm-up
tokens no longer matched what the consumer's generateResponse produced, so
cache reuse was defeated and any prior tool exchange in the prefix was lost
by the missing-state fallback.
Replace with a ChatHistoryItem[] renderer that reuses the existing
convertMessagesToChatHistory pipeline (extracted as a pure helper that skips
the trailing empty-user placeholder). Warm-up and every missing-state
fallback (Chat, TextGeneration, ToolCalling) now call
session.setChatHistory(history) then session.preloadPrompt("",
{ functions }), so the KV state is populated through the same wrapper the
consumer uses. Tools are handed to preloadPrompt via buildChatModelFunctions
so the wrapper embeds their descriptions identically. Deleted
renderLlamaCppPrefixText so no future path can silently drop tool blocks.
Test plan updated: flip the reconstructed-checkpoint preloadPrompts
expectation to [""] and assert the setChatHistory payload; add a
tool_use/tool_result preservation regression, plus dedicated unit tests for
the new render helpers (system-only, tools-only, mixed text+tool_use,
tool_result matching).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01413eP4Awsvn2HSXqr4MFYh
checkpointModelKey returns "" for a missing model_id, and validateParentCheckpoint short-circuited its mismatch check on either side being empty — so two keyless models on the same provider silently shared a fungible checkpoint slot (cross-model contamination) and an emit path could mint a slot with no identity at all. Add requireCheckpointModelKey (throws TaskConfigurationError for a keyless model) and route every mint / validate site through it: the parent-validate site (mismatch comparison is now unconditional as defense in depth), the resolveCheckpointSession pre-dispatch gate (so emit paths fail before createSession mints), and CacheCheckpointTask.prepareCheckpoint's mint site (the returned key is threaded straight into registerCheckpoint). Tests: dedicated requireCheckpointModelKey unit coverage plus a model-key fail-closed L2 describe covering all five paths (warm-up rejection with no run-fn invocation, parent-consume rejection before dispatch, keyless-vs-keyless slot contamination, emit path rejection before createSession, and a regression on the existing keyed-mismatch path). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01413eP4Awsvn2HSXqr4MFYh
…n atomic steal get + delete on the session map bracketed async work in both _TextGeneration and _ToolCalling stream paths; two concurrent consumers of the same immutable checkpoint id both observed the same cached state and both called .generate() on the shared LlamaContextSequence — a live sequence advances in place, so the second consumer's turn was corrupted. withModelInUse is a refcount, not a mutex, so it did not close the window. Add stealLlamaCppSession(id) in LlamaCpp_Runtime — a synchronous Map.get + Map.delete that is race-free on JS's single thread — and route both stream paths through it when consuming an immutable checkpoint (isCheckpoint && !ownedSession). Under contention exactly one caller wins; losers observe undefined and re-encode via the existing missing-state fallback (which acquires its own sequence). The ownership-tracking flag is now decided at the same point as the steal-vs-get split, dropping the downstream flip-flop that redundantly deleted the map entry on a hit. AiChatTask's ownedSession mode is the caller's mutable session (never a checkpoint), so LlamaCpp_Chat continues to use the non-consuming getLlamaCppSession — documented on the lookup site so a future refactor does not "unify" the two paths back into a shared race. Tests: stealLlamaCppSession unit coverage (atomic get+delete on hit, no-op on unknown id); a race test in the ToolCalling checkpoint suite covering the two-consumer scenario (winner reuses warmed sequence, loser re-encodes with its own) plus an ownedSession regression asserting racing ownedSession runs do not evict each other's map entries; and a new TextGenerationCheckpoint fixture mirroring the ToolCalling one so the text-generation path has parity coverage. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01413eP4Awsvn2HSXqr4MFYh
The Gemini cache-checkpoint warm-up wrapped ai.caches.create in a blanket try/catch that logged and swallowed every error, so a caller-initiated abort, a 401/403 auth failure, a 429 quota, and a 5xx transport error all silently degraded to the same "no cache, replay inline" fate as a genuine prefix-too- small 400. It also failed to clean up when a resource had been minted but a downstream step (bookkeeping / write) threw, leaving a server-side CachedContent billing until TTL. Classify the error into abort / degrade / throw. Abort and throw rethrow (with a best-effort delete of any resource whose name we saw); only a 400 / INVALID_ARGUMENT whose message names the too-small-prefix / cache-unsupported condition degrades to inline replay. Track `createdName` across the try / catch so the partial-delete cleanup can fire. Gemini_CacheStore gets a matching `createdAtMs` field on each entry, an `isGeminiCacheEntryStale` helper (default 3_500_000 ms — a hair under the 1h TTL used at write time), and a `deleteGeminiCachedContentLocal` for runtime-only eviction. These land here so PR #641's follow-up (H2, the consume-side fallback) can chain on top. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UC9LGu3iokhuhAB46yt1ee
…achedContent Gemini's explicit CachedContent is TTL-bound (~1h at creation) and can also vanish server-side outside the consumer's control. The text.generation and tool-calling run-fns treated the cache handle as guaranteed to resolve: an entry near-TTL still went out with `cachedContent`, and a NOT_FOUND surfaced as a hard failure to the caller. Add Gemini_CachedContentFallback: a proactive stale check (evicts the runtime-local entry before the request, leaves the server-side to its own TTL) plus a reactive NOT_FOUND catch that evicts (locally + best-effort server delete), rebuilds the request inline (prefix + tail), and retries once. Refactor both run-fns to share `buildCachedRequest` / `buildInlineReplayRequest` locals so the retry rebuilds an identical shape. Streaming contract unchanged (providers still yield text-delta/object-delta then finish; no accumulation). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UC9LGu3iokhuhAB46yt1ee
…fake session
The L1 checkpoint-prefix-render fix routes the missing-state fallback in
LlamaCpp_Chat through session.setChatHistory(history) + preloadPrompt(""),
matching the real LlamaChatSession API. The checkpoint mocks under
ai-provider-nodellama were updated to stub setChatHistory, but the
FakeLlamaChatSession in this ai-provider suite was missed, so the
checkpoint-seeded owned-session case threw "setChatHistory is not a
function". Add the no-op stub to align the fake with the real session.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YMyipztz9RFDp4aZPSNqEV
…main Main added HftSessionBase.cacheKey (pipeline-eviction sweeping); the rebased checkpoint commits created sessions without it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
sroussey
marked this pull request as ready for review
July 22, 2026 06:23
sroussey
force-pushed
the
ai-provider-cache-checkpoints
branch
from
July 22, 2026 06:23
e40442b to
9035d1f
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Summary
Verification
bun x vitest run packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts packages/test/src/test/ai-provider/Gemini_SessionDispose.test.ts— 18 passed, 1 Bun-only test skippedbun test packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts packages/test/src/test/ai-provider/Gemini_SessionDispose.test.ts— 19 passedbun run build:types— 40 packages passed