fix(ai,node-llama-cpp): three HIGH-severity checkpoint correctness fixes (follow-up to #641)#642
Open
sroussey wants to merge 35 commits into
Open
fix(ai,node-llama-cpp): three HIGH-severity checkpoint correctness fixes (follow-up to #641)#642sroussey wants to merge 35 commits into
sroussey wants to merge 35 commits into
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
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
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
Three HIGH-severity correctness fixes surfaced during scheduled review of PR #641 (
ai-provider-cache-checkpoints). Each fix is a self-contained commit; each has targeted tests.L1 — LlamaCpp checkpoint prefix rendered through the model's chat template
renderLlamaCppPrefixTextflattened the prefix to raw"role: text"strings and silently droppedtool_use/tool_resultblocks. Warm-up KV tokens no longer matched the consumer's chat-template rendering — cache reuse was defeated and prior tool exchanges were lost. Replaced with aChatHistoryItem[]renderer routed throughsetChatHistory+preloadPrompt("", { functions }). Removed the raw-text renderer so no future path can silently drop tool blocks.L2 — CheckpointRegistry model-key fail-closed
checkpointModelKeyreturned""for a missingmodel_id, andvalidateParentCheckpointshort-circuited its mismatch check on either side being empty — two keyless models on the same provider shared a fungible checkpoint slot. AddedrequireCheckpointModelKey(throwsTaskConfigurationErrorfor keyless models), routed all mint/validate sites through it, and made the mismatch comparison unconditional as defense in depth.L3 — LlamaCpp checkpoint concurrent-consumer race
get+deleteon the session map bracketed async work in both_TextGeneration_Streamand_ToolCalling_Stream; two concurrent consumers of the same checkpoint id both took the same reference and generated on the sharedLlamaContextSequence. AddedstealLlamaCppSession(id)inLlamaCpp_Runtime(synchronousget+delete, race-free on JS's single thread); under contention exactly one consumer wins, losers deterministically re-encode via the existing fallback.Test plan
bun scripts/test.ts provider-nodellama vitest— 57 tests pass acrossLlamaCpp_CheckpointPrefixRender.test.ts(new),LlamaCpp_ToolCallingCheckpoint.test.ts,LlamaCpp_TextGenerationCheckpoint.test.ts(new),LlamaCpp_Runtime.test.ts,LlamaCppProvider.test.tsbunx vitest run packages/test/src/test/ai/CacheCheckpoint.test.ts— 33 tests pass (extended withrequireCheckpointModelKeyunit coverage and a five-casemodel-key fail-closed (L2)describe)bun run build:types— clean across all 40 packagesNote on base
This branch was based on
main(a90e59c) but the fixes patch code introduced in PR #641. To keep the fixes runnable and testable in this session,origin/ai-provider-cache-checkpointswas merged in as the base for the three fix commits — the first commit on the branch is that merge. Fixes here are intended to be cherry-picked or merged into PR #641's branch (rebase / merge order at the maintainer's discretion); each of the three fix commits is self-contained and applies cleanly on top of PR #641.Generated by Claude Code