Add and harden AI provider cache checkpoints#641
Draft
sroussey wants to merge 31 commits into
Draft
Conversation
…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
sroussey
force-pushed
the
ai-provider-cache-checkpoints
branch
from
July 17, 2026 15:59
d76e7f3 to
68458c8
Compare
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>
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
3 participants
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