Repository navigation
fix(injection): keep injected context out of the prompt cache prefix - #3642
opieter-aws wants to merge 2 commits into
Conversation
github-actions
Bot
added
python
GitHub Actions
Error
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
|
Assessment: Comment Solid, well-reasoned fix — the root cause (per-call text landing inside the cached prefix) is correctly addressed by moving volatile content behind a shared Review themes
Nice work on the measurement rigor (nonce-per-arm A/B) and the thorough write-up of the trade-off. |
a8dc17f to
767b13f
Compare
GitHub Actions
Inactive
|
Re-review of the latest push ( Follow-up items resolved
The only thing still open is not a code issue: the behavior/billing contract call (ContextInjector implying prompt caching without |
Documentation Preview ReadyYour documentation preview has been successfully deployed! Changed pages: Updated at: 2026-08-05T13:29:31.419Z |
|
Putting back into draft - integration test is pointing at some token counting issues with caching, need to take a look there |
GitHub Actions
Failure
| last_assistant_idx = i | ||
| for index, message in reversed(list(enumerate(messages))): | ||
| usage = message.get("metadata", {}).get("usage") if message.get("role") == "assistant" else None | ||
| if usage and not (usage.get("cacheReadInputTokens") or usage.get("cacheWriteInputTokens")): |
There was a problem hiding this comment.
Issue (Suggestion): The new cache-skip branch is the crux of this fix, but it's only exercised by the Bedrock-gated integration tests (skipped by default in CI). The fast TestEstimateInputTokens suite (test_event_loop.py) has sibling cases for cold-start, baseline-only, and baseline+delta, but none where an assistant usage carries cacheReadInputTokens/cacheWriteInputTokens. So the not (...) skip path can regress without any always-on test failing.
Suggestion: Add a unit case here (and the mirror in TS agent.test.ts _estimateInputTokens): a recent assistant message with cacheReadInputTokens set should be skipped, so the estimator either anchors on an earlier uncached assistant or falls through to the cold-start count path. That pins the branch independently of live Bedrock. The integ tests are great for the end-to-end contract; this just guards the logic cheaply.
|
Re-review of Assessment of the token-counting fix
No blockers from me on the code. Still in draft, so flagging lightly — and the earlier open item stands: the ContextInjector-implies-caching billing/behavior contract still warrants an |
|
@strandly-the-agent review please |
mkmeral
left a comment
mkmeral
left a comment
There was a problem hiding this comment.
overall for caching, i'd like to have higher integ test count to catch issues (or maybe merge the logic to some other test). aside from that looks good to me
|
|
||
| content = list(message["content"]) | ||
| if content and not any("cachePoint" in block for block in content): | ||
| content.append({"cachePoint": {"type": "default"}}) |
There was a problem hiding this comment.
rest of the model providers translate based on allowlist right? so new message should be ignored, if the model provider doesn't support?
| if msg.get("role") == "user": | ||
| last_user_idx = msg_idx | ||
|
|
||
| # Strip cache points everywhere except the last user message, where auto mode would |
There was a problem hiding this comment.
should we have an integ test for this?
strandly-the-agent
left a comment
strandly-the-agent
left a comment
There was a problem hiding this comment.
TL;DR — recommend changes before merge. The cache-prefix diagnosis is right and the nonce-per-arm A/B is the kind of evidence I wish every perf PR had. But the boundary is emitted for every provider with no gate, and I reproduced two regressions from that. One gate fixes both.
What I actually ran (branch pr3642 @ 8a08eea4, merge-base 0ab31d7c1):
- 🔴
Agent(model=OllamaModel(...), context_manager="agentic")→TypeError: content_type=<cachePoint>atollama.py:158, turn 1, no plugin, no special input. Same for gemini / writer / llamaapi / llamacpp / openai-responses; mistral silently drops the block. The identical script onmainreaches the network instead. - 🔴 A last user turn of
[text, document(csv)]+ any injected content ships[text, document, cachePoint, text]— the adjacencybedrock.py:456-472exists to prevent. It also happens with nocache_configat all, where_inject_cache_pointnever runs to relocate anything. - ✅ Suites green, pre-existing failures only: py 4408 passed / 2 failed (otel) / 28 errors (mantle-config); ts 3928 passed / 2 failed (otel env).
test_bedrock.py185/185. - ✅ I applied the honor-branch patch I suggest below and re-ran: invalid ordering gone, 185/185 still green — but the no-
cache_configpath stays broken, which is why I think the gate matters more than the patch. - ⚪ Not verified: the live Bedrock
ValidationExceptionitself (no creds here). That rests on this repo's own assertions attest_bedrock.py:3319+.
3 inline comments. mkmeral's question on content.py:312 is answered in the first one — short version: no, most providers don't ignore it, they raise.
Questions (the ones I can't settle for you)
1. Should a context-shaping feature turn on prompt caching? (blocking, and it's your call, not mine)
models/model.py:138 documents CacheConfig(strategy="auto") as "automatically detect model support and inject cachePoint" — that's the switch users know about, and types/content.py:311-312 now injects one without it. Verified: Bedrock + ContextInjector with zero cache_config goes from 0 cache points on main to 1 here; AnthropicModel converts unconditionally (anthropic.py:206-207, no cache config exists on that provider at all). There's no opt-out — ContextInjector.__init__ takes render_content, name, trigger.
Two bits of prior art worth weighing, both pointing the same way: team/DECISIONS.md:120 says pay-for-play doesn't cover "changes that affect default behavior", and in #449 Unshure declined auto-injecting a message cachePoint even when cache_tools was set, on the grounds that injected message cache points surprise customers. Your "multi-turn always wins" analysis may well carry the day — but if it does, I'd expect that to land as a documented default on CacheConfig plus a release note, rather than as a side effect of enabling an unrelated feature.
The capability gate you floated in the PR body is the option that closes all three of these at once (billing, the 6 crashing providers, the document-adjacency path with no cache_config), and it's cheaper than it looks: TS already has the predicate (_shouldEnableCaching(), bedrock.ts:476-496), py is an extract-method from bedrock.py:580-591, and both integ tests that assert the cache win already set cache_config=CacheConfig(strategy="auto") (test_context_injector_caching.py:66, :127) — so the measured win survives with no test edits.
2. Re your question on bedrock.py:426 (mkmeral) — yes to an integ test, but I'd sequence it after the fix below, otherwise it encodes the current ordering.
3. (non-blocking) The body describes commit 767b13fe5 only; 8a08eea43 ("Fix token counting with cache points") changes _estimate_input_tokens, which feeds proactive compression — worth a paragraph so a future bisector chasing a compression regression has a reason to look here. Related: that gate fixes a bug that was already reachable on main for anyone using cache_config (the new integ test at test_agentic_context_integration.py:115-155 uses no ContextInjector), so "pre-existing, but this PR makes it common" is probably the framing you want.
4. (non-blocking) Also missing the typescript label for 13 TS files, and the api/needs-review label the earlier reviews asked for — mentioning once, not re-litigating.
Appendix — non-blocking (6)
- ⚪ 4-cache-point ceiling. The strip loop now skips the last user message and the honor branch returns after the first
cachePoint, so extras there survive. Getting past Claude's cap of 4 needs ≥3 hand-placed points plus system + tool caching — plausible, not mainline. Worth knowing the margin is gone: two hand-placed now lands exactly at 4. - ⚪ Boundary-reuse degenerate shapes. A message seeded with
content: []gets no boundary, and Bedrock auto mode then appends one after the ephemeral text (cache write every call, never a read); a leading hand-placedcachePointsuppresses the real boundary so nothing is cacheable. Both need hand-seeded messages — contrived, noting for completeness. - ⚪ No delimiter between the ask and the appended text. Providers that join text blocks fuse them: mistral yields
"What is the capital of France? <context-status>…</context-status> INJECTED".<context-status>/<memory>self-delimit; aRenderContentCallbackreturning plain text doesn't, and nothing in its docstring (_message_injection.py:29-38) says it now needs to. A\n\nprefix in the primitive would settle it. - ⚪
ephemeralis overloaded, inverted.anthropic.py:207writescache_control={"type": "ephemeral"}for content that is cached; the new primitive uses "ephemeral" for content that must not be cached, three lines from that code. It's private/@internal, so_append_per_call_contentis a free rename if you agree (team/DECISIONS.md:168). - ⚪
AnthropicModelIndexError—formatted_contents[-1](anthropic.py:203-207) blows up if every earlier block in the message was filtered out (location-sourced image + an appended boundary). Reproduces onmainwith a hand-placed cache point, so it's pre-existing, not yours — I'll file it separately. TS avoids it with a lookahead (anthropic.ts:377-388), which is probably the shape to port. - ⚪ TS coverage is thinner than py's here. No ts mirror for the stale-boundary-stripping test, and the new integ test at
test/integ/agent.test.ts:159-178constructs the model with nocacheConfig, socacheRead/cacheWriteare trivially0and it passes whether or not8a08eea43's fix exists. The python integ test does turn caching on — that asymmetry is worth closing givenmkmeral's ask for more caching coverage. - ⚪ Trailing whitespace
content.py:290,:295; stray word inmessages.ts:967("injected context count").
Everything above is one agent's pass — worth a human check before you act on it, especially the API question.
|
|
||
| content = list(message["content"]) | ||
| if content and not any("cachePoint" in block for block in content): | ||
| content.append({"cachePoint": {"type": "default"}}) |
There was a problem hiding this comment.
🔴 The boundary goes to every provider, and 6 of them raise on it.
mkmeral asked on this line whether providers translate on an allowlist and would just ignore an unsupported block. I checked all of them — mostly no:
| what happens | providers |
|---|---|
| converts it | bedrock bedrock.py:637, anthropic anthropic.py:206 |
| warns + skips | openai openai.py:413 (this PR), litellm, sagemaker (inherit it) |
raises TypeError |
gemini gemini.py:226, ollama ollama.py:158, writer writer.py:123, llamaapi llamaapi.py:130, llamacpp llamacpp.py:246, openai-responses openai_responses.py:728 |
| silently drops | mistral mistral.py:226-232 (if/elif, no else, no warning) |
Reachable on turn 1 with no plugin and no unusual input, because create_token_usage_middleware (agent.py:414-419) has no provider check and _estimate_input_tokens runs on every model call:
Agent(model=OllamaModel(host="http://localhost:11434", model_id="llama3"), context_manager="agentic")
agent("What is the capital of France?")
→ TypeError: content_type=<cachePoint> | unsupported type (ollama.py:158)
Same via ContextInjector and MemoryManager injection. The identical script on main gets to the network, so it's a regression, not a latent gap. TS is fine — its adapters have a generic default: warn+skip — which I suspect is where the PR body's "matching the Gemini/Vercel providers" came from; py Gemini raises.
Two ways out:
- Gate the emission on a model capability —
Model.wants_cache_boundary(), defaultFalse,BedrockModelreturning_should_enable_caching(), threaded in as a keyword arg here. TS already has the predicate (bedrock.ts:476-496); py is an extract-method frombedrock.py:580-591. This also fixes the Bedrock document-adjacency finding and the no-cache_configbilling question in one move. - Or make every provider tolerate it — e.g. a
drop_unsupported_cache_points()helper alongside_has_location_sourceinmodels/_validation.py, wired into the 6 raisers plus Mistral's missingelse. AcachePointcarries no content the model would miss, so warn-and-skip everywhere seems like the right SDK-wide rule; raising punishes the user for a marker the framework inserted.
Whichever you pick, context-injector.mdx:123 ("Providers without cache-point control ignore the marker") needs to become true — it's currently false for those 6. And there's no test anywhere on this: grep -c cachePoint is 0 in test_gemini.py, test_ollama.py, test_writer.py, test_llamaapi.py, test_llamacpp.py, test_mistral.py, test_openai_responses.py.
| placed = next((block for block in content if "cachePoint" in block), None) | ||
| if placed is not None: | ||
| if cache_config and cache_config.ttl and "ttl" not in placed["cachePoint"]: | ||
| placed["cachePoint"]["ttl"] = cache_config.ttl | ||
| logger.debug("msg_idx=<%s> | honored caller-placed cache point", last_user_idx) | ||
| return |
There was a problem hiding this comment.
🔴 Honoring the boundary skips the non-PDF-document rule this method already enforces below.
The return here lands before the relocation logic at bedrock.py:456-472 — the guard added by beddc3f7d / #2001 precisely because Bedrock rejects a cachePoint directly preceded by a non-PDF document. _append_ephemeral_content always appends at the tail, so a [text, document(csv)] turn plus any injected content produces exactly that shape. Verified on this head:
injector output: ['text', 'document', 'cachePoint', 'text']
after _inject_cache_point (auto): ['text', 'document', 'cachePoint', 'text'] ← cachePoint at idx 2
control, no pre-existing boundary: ['text', 'cachePoint', 'document'] ← correctly relocated
Reached by agent([{"text": "Analyze this file"}, {"document": {"format": "csv", ...}}]) with ContextInjector / MemoryManager / agentic mode — an ordinary "analyze this file" flow. Same in strands-ts/src/models/bedrock.ts:882-891.
This patch fixes it (I applied it and ran test_bedrock.py: 185/185, invalid ordering gone):
| placed = next((block for block in content if "cachePoint" in block), None) | |
| if placed is not None: | |
| if cache_config and cache_config.ttl and "ttl" not in placed["cachePoint"]: | |
| placed["cachePoint"]["ttl"] = cache_config.ttl | |
| logger.debug("msg_idx=<%s> | honored caller-placed cache point", last_user_idx) | |
| return | |
| placed = next((block for block in content if "cachePoint" in block), None) | |
| if placed is not None: | |
| if cache_config and cache_config.ttl and "ttl" not in placed["cachePoint"]: | |
| placed["cachePoint"]["ttl"] = cache_config.ttl | |
| # Bedrock rejects a cache point directly preceded by a non-PDF document, so a | |
| # caller-placed point follows the same relocation rule as one added below. | |
| first_non_pdf_doc_idx = next( | |
| ( | |
| block_idx | |
| for block_idx, block in enumerate(content) | |
| if "document" in block and block["document"].get("format", "") != "pdf" | |
| ), | |
| None, | |
| ) | |
| if first_non_pdf_doc_idx is not None and first_non_pdf_doc_idx < content.index(placed): | |
| content.remove(placed) | |
| if first_non_pdf_doc_idx == 0: | |
| # A leading non-PDF document leaves no prefix to cache | |
| logger.debug("msg_idx=<%s> | dropped cache point for leading document", last_user_idx) | |
| return | |
| content.insert(first_non_pdf_doc_idx, placed) | |
| logger.debug("msg_idx=<%s> | relocated caller-placed cache point", last_user_idx) | |
| return | |
| logger.debug("msg_idx=<%s> | honored caller-placed cache point", last_user_idx) | |
| return |
But the patch isn't sufficient, and this is the part I'd weigh when deciding on the capability gate: with no cache_config at all, _inject_cache_point never runs, yet the injector's boundary still reaches the wire — the invalid ['text', 'document', 'cachePoint', 'text'] payload happens anyway, with the patch applied. Verified:
BedrockModel(model_id=...) # no cache_config
format_request([injected_message]) → ['text', 'document', 'cachePoint', 'text']
So either the boundary is only emitted where the model asked for caching (the gate), or the placement rule has to move somewhere that runs unconditionally. A test that pins it: test_inject_cache_point_honored_boundary_does_not_land_after_a_non_pdf_document, asserting "cachePoint" not in content[doc_idx + 1] — it fails against the current code, which is the evidence. Worth mirroring in TS.
| last_assistant_idx = i | ||
| for index, message in reversed(list(enumerate(messages))): | ||
| usage = message.get("metadata", {}).get("usage") if message.get("role") == "assistant" else None | ||
| if usage and not (usage.get("cacheReadInputTokens") or usage.get("cacheWriteInputTokens")): |
There was a problem hiding this comment.
🟡 The gate is provider-blind, and for the OpenAI family it makes the estimate worse.
Different point from the existing thread on this line (that one asks for a unit test — still worth doing; this is about the logic).
The premise — inputTokens excludes the cached portion — holds for Bedrock/Anthropic, and the gate is right there. It's inverted for the OpenAI family: prompt_tokens includes cached tokens and cacheReadInputTokens is a subset of it (openai.py:596-597, openai_responses.py:888-896, litellm.py:287-295, gemini.py:479-487). So for those providers the gate throws away the one accurate number available and falls through to count_tokens, which has no image/document branch (model.py:39-86 — deliberately, per its own comment). On a vision-heavy conversation:
provider-reported prompt size: 26,000
before this commit: 26,408 (1.6% off)
after: 17 (99.9% off)
proactive compression at 0.85 of a 30k window: fires before → does not fire after
Reachable without any of this PR's features: OpenAI caches server-side automatically above ~1024 prompt tokens, so any multi-turn OpenAIModel conversation reports cacheReadInputTokens and lands in the fallthrough. The consequence is bounded — a wrong number feeding compression and the <context-status> text the model reads, not a hard failure — hence 🟡.
Second-order, worth a sentence in the comment either way: once a boundary is emitted by default, every turn carries cacheWrite or cacheRead, so the baseline+delta optimization never engages again for the life of a cache-enabled conversation — you're doing a full recount per model call. Cheap today (local heuristic), a network call per attempt with use_native_token_count=True.
Suggest keying the skip on whether the provider's inputTokens is net-of-cache rather than on the presence of cache counters — the comment already concedes providers disagree, and that disagreement is knowable per provider.
Description
ContextInjectorprepended its per-call text into the latest user message, ahead of the stable conversation. Because that text is rebuilt on every call and the provider's cache point landed after it, the cached prefix diverged exactly at the injection site — so every turn re-ingested the whole conversation at full price. Measured against Bedrock, turn 2 of a two-turn conversation reused zero cached tokens.Injected text is now appended behind a
cachePointboundary that marks where the reusable prefix ends, and Bedrock's auto mode honors a boundary already present instead of relocating it to the end. Placement belongs to whoever appends the volatile content, so the agentic token-usage middleware shares the same primitive — its live<context-status>count previously stranded itself inside the cached prefix even with the injector fixed.OpenAIModelnow warns and skips acachePointinside a message rather than raisingTypeError, matching its own system-prompt path and the Gemini/Vercel providers. Without it, a provider-agnostic caller cannot mark ephemeral content without special-casing OpenAI.Related Issues
N/A
Documentation PR
Included
Type of Change
Bug fix
Testing
Verified end to end against real Bedrock (
claude-sonnet-4-5), A/B against the pre-fix placement with a unique nonce per arm so no stale cache entry could flatter the result:hatch run prepareRisks / Callouts
This changes billing shape for users who never enabled caching. The boundary is honored by Bedrock and Anthropic even with no
cache_configset, soContextInjectornow implies prompt caching on those providers. Bounds, measured:The only loss case is a single-turn conversation, multi-turn always wins. I considered gating on a
Modelcapability so the boundary is only emitted where the user opted in, but that forfeits the win for everyone who hasn't setcache_config— most users. Flagging it as the reviewable trade-off: if "no silent cost change" is a hard rule regardless of magnitude, the gate is the alternative.Bedrock's
_inject_cache_pointbehavior change is broader than the injector needs. It now honors a caller-placed cache point in the last user message (still stripping stale ones elsewhere, so the 4-cache-point cap can't be exhausted). That affects anyone hand-placing cache points, and an existing test asserting the old strip-everything behavior was updated.Checklist
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.