ISS-83: Add Prompt Caching Read/Write Tokens - #87
Open
jrswab wants to merge 3 commits into
Open
Conversation
Add CacheConfig field to provider.Request to enable prompt caching. Track CacheReadTokens and CacheWriteTokens in provider.Response. Display cache token metrics in runner output for verbose mode. Include cache token fields in JSON output. Supports Anthropic, Bedrock, and OpenAI providers. All providers update usage tracking for cache tokens.
📝 WalkthroughWalkthroughAdds request-level CacheConfig and propagates cache read/write token accounting from provider wire formats (Anthropic/Bedrock/OpenAI) through streaming and the runner into final Result JSON. ChangesCache Token Accounting
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsTimed out fetching pipeline failures after 30000ms Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Owner
Author
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
pkg/runner/run.go (2)
550-559:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winVerbose cache line prints even when both values are 0.
For providers that don't support caching (OpenAI, Gemini, Bedrock without caching), every verbose run will print
Cache: 0 read, 0 written. The equivalent token/budget lines already guard ontracker.Max() > 0to suppress irrelevant output.🛠️ Proposed fix
-_, _ = fmt.Fprintf(stderr, "Cache: %d read, %d written\n", resp.CacheReadTokens, resp.CacheWriteTokens) +if resp.CacheReadTokens > 0 || resp.CacheWriteTokens > 0 { + _, _ = fmt.Fprintf(stderr, "Cache: %d read, %d written\n", resp.CacheReadTokens, resp.CacheWriteTokens) +}Apply the same guard on Line 686 in the conversation-loop path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/runner/run.go` around lines 550 - 559, The verbose output prints "Cache: 0 read, 0 written" for providers without caching; change the verbose printing logic so the cache line is only emitted when budget/tracker is applicable—wrap the cache print that uses resp.CacheReadTokens and resp.CacheWriteTokens in the same guard used for token/budget (check tracker.Max() > 0) in both the current opts.Verbose block and the corresponding conversation-loop branch so the Cache line is suppressed for providers without caching.
371-378:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift
CacheConfig: truehardcoded opts every user into prompt caching without a way to disable it.Prompt caching changes the Anthropic wire format (system becomes a content-block array instead of a plain string) and has billing implications —
cache_creation_input_tokensare priced at 25% more than regular input tokens. There's currently no agent TOML key, CLI flag, or global config option to disable it. Existing users will silently start paying cache-write prices on every request.Consider either:
- Adding a
cache_enabled(defaulttrue) TOML key and wiring it throughOptions→provider.Request.CacheConfig, or- At minimum, documenting this as an always-on behaviour change in the changelog/README.
💡 Minimal opt-out wiring sketch
// In agent config TOML struct (agent package) +CacheEnabled *bool `toml:"cache_enabled"` // nil = default (true) // In run.go +cacheEnabled := true +if cfg.CacheEnabled != nil { + cacheEnabled = *cfg.CacheEnabled +} req := &provider.Request{ ... - CacheConfig: true, + CacheConfig: cacheEnabled, }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/runner/run.go` around lines 371 - 378, The run.go code currently hardcodes CacheConfig: true on the provider.Request (in pkg/runner/run.go), which forces prompt caching; add a new boolean configuration field (e.g., CacheEnabled defaulting to true) on the agent/CLI/TOML config struct (wire it under cfg.Params or Options), surface it through any existing Options -> provider.Request plumbing, and set provider.Request.CacheConfig = cfg.Params.CacheEnabled (or equivalent) instead of true; ensure the new TOML/CLI option is documented and defaults to true so users can opt out of Anthropic prompt-cache writes and extra billing.internal/provider/anthropic.go (1)
480-561:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftRead cache tokens from
message_startto match the authoritative source pattern.The
message_startevent includes cache tokens inmessage.usage(per Anthropic's streaming spec), and the code already handlesInputTokensthat way. Cache tokens currently come frommessage_delta, which duplicates them as a convenience — but if Anthropic ever stops that duplication, the counts silently drop to zero.The fix: capture
CacheReadInputTokensandCacheCreationInputTokensfrommessage_start(line 482) in the same scope asinputTokens, remove redundant reads frommessage_delta, and add a test with cache fields in the mockedmessage_startevent.Also note: existing streaming tests don't include cache token fields, leaving this path untested.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/anthropic.go` around lines 480 - 561, When handling the "message_start" event (where you already capture inputTokens), also read and store message.Message.Usage.CacheReadInputTokens and CacheCreationInputTokens into the same local vars (e.g., cacheReadTokens, cacheWriteTokens) instead of relying on message_delta; stop populating cacheReadTokens/cacheWriteTokens from the "message_delta" case and have the "message_delta" return use those previously-captured vars when building the StreamEvent (fields CacheReadTokens/CacheWriteTokens). Update/add the streaming test to include CacheReadInputTokens and CacheCreationInputTokens in the mocked message_start event so this path is exercised.
🧹 Nitpick comments (6)
internal/provider/bedrock_test.go (2)
286-315: ⚡ Quick winAdd an assertion for the tool-config cachePoint branch
buildBedrockRequestnow setsToolConfig.CachePointwhen tools are present and caching is enabled, but this test only verifies the system block path. Adding a tools-present case would lock down that new branch too.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/bedrock_test.go` around lines 286 - 315, The test TestBedrock_Send_CachePointInSystem only asserts cachePoint on the System block but not the new tools branch; add a case that constructs a Request with Tools (e.g., set Request.Tools to a non-empty slice), CacheConfig true, and then call b.Send (same as existing test) and assert that the decoded receivedReq.ToolConfig is non-nil and that receivedReq.ToolConfig.CachePoint.Type == "default" (mirroring the system assertions); this will exercise buildBedrockRequest’s ToolConfig.CachePoint path when tools are present.
286-380: ⚡ Quick winPlease convert these new cache tests to table-driven format
Coverage is good; reshaping these into table-driven cases will align with repo test conventions and keep future cache cases easy to add.
As per coding guidelines,
**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern for all test files.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/bedrock_test.go` around lines 286 - 380, Convert the three standalone tests (TestBedrock_Send_CachePointInSystem, TestBedrock_Send_CacheUsageParsing, TestBedrock_Send_CacheUsageOmitted) into a single table-driven test: define a []struct with fields like name, serverResponse (bedrockResponse used by httptest server), request (the Request passed to b.Send), wantSystemBlocks (count) / wantSystemCachePoint (bool) / wantCacheReadTokens (int) / wantCacheWriteTokens (int), then range over cases and run each case as t.Run(case.name, func(t *testing.T){...}). Inside the loop create the httptest.NewServer returning case.serverResponse, instantiate Bedrock with NewBedrock(..., WithBedrockBaseURL(server.URL), ...), call b.Send with case.request, and assert expectations (receivedReq.System and its CachePoint for the system-cache case and resp.CacheReadTokens/resp.CacheWriteTokens for usage cases) using the case.want* fields; keep existing JSON encode/decode logic and defer server.Close per subtest. Ensure unique identifiers from the diff—NewBedrock, WithBedrockBaseURL, b.Send, and the bedrockResponse/bedrockUsage fields—are used to locate and adapt the code.internal/provider/provider_test.go (1)
190-212: ⚡ Quick winConvert new cache field tests to table-driven style
These assertions are solid, but please group them into a table-driven test to match the repo’s test convention (future additions become much easier).
As per coding guidelines,
**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern for all test files.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/provider_test.go` around lines 190 - 212, Convert the two tests into table-driven style: replace TestResponse_CacheTokens and TestRequest_CacheConfig with a single or separate table-driven tests that define a []struct{name string, resp *Response, req *Request, wantRead int, wantWrite int, wantCacheConfig bool} (or two small tables) and iterate using t.Run for each case, verifying Response.CacheReadTokens and CacheWriteTokens and Request.CacheConfig against want values; locate the types Response and Request and the fields CacheReadTokens, CacheWriteTokens, and CacheConfig to implement the table entries and assertions.internal/provider/openai_test.go (1)
1670-1761: ⚡ Quick winGreat new cache coverage—please table-drive these cases
The new scenarios are useful; converting them into a single table-driven test block would align with project test standards and reduce duplication.
As per coding guidelines,
**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern for all test files.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/openai_test.go` around lines 1670 - 1761, Combine the three tests TestOpenAI_Send_CacheUsageParsing, TestOpenAI_Send_CacheUsageOmitted and TestOpenAI_SendStream_CacheUsageParsing into one table-driven test that iterates over cases with fields like name, isStream (bool), responsePayload (string), and wantCacheReadTokens (int); for each case spin up an httptest.Server that returns the case.responsePayload (for stream cases set Content-Type: text/event-stream and include the SSE messages used in the original test), create the client via NewOpenAI(..., WithOpenAIBaseURL(server.URL)), call o.Send for non-stream cases and o.SendStream for stream cases (remember to Close the stream), then assert resp.CacheReadTokens or event.CacheReadTokens equals wantCacheReadTokens; keep the original expectations and error handling but remove duplicated setup/teardown by running them inside the loop.internal/provider/anthropic.go (1)
207-215: 💤 Low value
convertToAnthropicTools: thecacheEnabledbranch is a no-op.The
if cacheEnabled && i == len(tools)-1block reassignstd.InputSchema = schema, which is identical to what was just set on line 208. The comment says toolinput_schemacaching is unsupported for now — in that case the entireifblock should be removed rather than left as a no-op, as it misleads future readers into thinking something cache-related is happening.🧹 Proposed cleanup
td := anthropicToolDef{ Name: tool.Name, Description: tool.Description, InputSchema: schema, } - if cacheEnabled && i == len(tools)-1 { - td.InputSchema = schema // cache on input_schema is unsupported; skip for now - } result = append(result, td)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/anthropic.go` around lines 207 - 215, The branch in convertToAnthropicTools that checks "if cacheEnabled && i == len(tools)-1" is a no-op because it reassigns td.InputSchema = schema which was just set; remove that entire if block and its misleading comment, or replace it with a clear TODO explaining input_schema caching is unsupported, ensuring you update references to td, cacheEnabled, tools, and i in the same function so no unused-code remains.pkg/runner/result_test.go (1)
228-276: ⚡ Quick winNew cache-token tests should be folded into the existing table-driven test.
Both
TestResultMarshalJSON_CacheTokensandTestResultMarshalJSON_CacheTokensOmittedWhenZeroduplicate the structure ofTestResultMarshalJSON_TableDrivenabove. Per project guidelines, all tests should follow the[]struct{name, input, want}table-driven pattern.♻️ Proposed refactor: extend the table-driven test instead
cases := []struct { name string input Result want struct { - Cost bool - CostVal float64 - Cache bool - CacheVal string + Cost bool + CostVal float64 + Cache bool + CacheVal string + CacheReadTokens int + CacheWriteTokens int } }{ { name: "cost and cache status present", ... }, { name: "zero cost and empty cache omitted", ... }, + { + name: "non-zero cache tokens present", + input: Result{Content: "Hello", InputTokens: 10, OutputTokens: 5, Model: "openai/gpt-4", StopReason: "stop", CacheReadTokens: 42, CacheWriteTokens: 7}, + want: struct{ ... }{CacheReadTokens: 42, CacheWriteTokens: 7}, + }, + { + name: "zero cache tokens omitted", + input: Result{Content: "Hello", InputTokens: 10, OutputTokens: 5, Model: "openai/gpt-4", StopReason: "stop"}, + want: struct{ ... }{}, + }, }As per coding guidelines: "Use table-driven tests with
[]struct{name, input, want}pattern for all test files."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/runner/result_test.go` around lines 228 - 276, The two new tests TestResultMarshalJSON_CacheTokens and TestResultMarshalJSON_CacheTokensOmittedWhenZero duplicate the pattern used by TestResultMarshalJSON_TableDriven; merge their cases into the existing table-driven test by adding two new entries to the []struct{name, input, want} slice used by TestResultMarshalJSON_TableDriven: one case with Result{... CacheReadTokens:42, CacheWriteTokens:7} expecting JSON keys "cache_read_tokens"=42 and "cache_write_tokens"=7, and one case with zero values expecting those keys to be omitted; then remove the two standalone test functions and rely on the extended TestResultMarshalJSON_TableDriven that serializes Result and asserts parsed JSON as before (refer to Result, CacheReadTokens, CacheWriteTokens, and TestResultMarshalJSON_TableDriven).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/provider/anthropic_test.go`:
- Around line 893-1050: The tests cover cache-token parsing for Send but not for
streaming; add a new unit test (e.g.,
TestAnthropic_SendStream_CacheUsageParsing) that uses the existing MockLLMServer
(or httptest.Server) to emit a hand-crafted SSE stream containing a
message_delta (and optionally message_start) JSON payload that includes usage
fields "cache_creation_input_tokens" and "cache_read_input_tokens", call
Anthropic.SendStream and consume until you receive a StreamEventDone, then
assert the returned StreamEvent (or result) contains CacheWriteTokens and
CacheReadTokens matching the mock; reference SendStream in anthropic.go (which
parses message_delta.usage) and reuse naming/fixtures from
TestAnthropic_Send_CacheUsageParsing for consistency.
---
Outside diff comments:
In `@internal/provider/anthropic.go`:
- Around line 480-561: When handling the "message_start" event (where you
already capture inputTokens), also read and store
message.Message.Usage.CacheReadInputTokens and CacheCreationInputTokens into the
same local vars (e.g., cacheReadTokens, cacheWriteTokens) instead of relying on
message_delta; stop populating cacheReadTokens/cacheWriteTokens from the
"message_delta" case and have the "message_delta" return use those
previously-captured vars when building the StreamEvent (fields
CacheReadTokens/CacheWriteTokens). Update/add the streaming test to include
CacheReadInputTokens and CacheCreationInputTokens in the mocked message_start
event so this path is exercised.
In `@pkg/runner/run.go`:
- Around line 550-559: The verbose output prints "Cache: 0 read, 0 written"
for providers without caching; change the verbose printing logic so the cache
line is only emitted when budget/tracker is applicable—wrap the cache print that
uses resp.CacheReadTokens and resp.CacheWriteTokens in the same guard used for
token/budget (check tracker.Max() > 0) in both the current opts.Verbose block
and the corresponding conversation-loop branch so the Cache line is suppressed
for providers without caching.
- Around line 371-378: The run.go code currently hardcodes CacheConfig: true on
the provider.Request (in pkg/runner/run.go), which forces prompt caching; add a
new boolean configuration field (e.g., CacheEnabled defaulting to true) on the
agent/CLI/TOML config struct (wire it under cfg.Params or Options), surface it
through any existing Options -> provider.Request plumbing, and set
provider.Request.CacheConfig = cfg.Params.CacheEnabled (or equivalent) instead
of true; ensure the new TOML/CLI option is documented and defaults to true so
users can opt out of Anthropic prompt-cache writes and extra billing.
---
Nitpick comments:
In `@internal/provider/anthropic.go`:
- Around line 207-215: The branch in convertToAnthropicTools that checks "if
cacheEnabled && i == len(tools)-1" is a no-op because it reassigns
td.InputSchema = schema which was just set; remove that entire if block and its
misleading comment, or replace it with a clear TODO explaining input_schema
caching is unsupported, ensuring you update references to td, cacheEnabled,
tools, and i in the same function so no unused-code remains.
In `@internal/provider/bedrock_test.go`:
- Around line 286-315: The test TestBedrock_Send_CachePointInSystem only asserts
cachePoint on the System block but not the new tools branch; add a case that
constructs a Request with Tools (e.g., set Request.Tools to a non-empty slice),
CacheConfig true, and then call b.Send (same as existing test) and assert that
the decoded receivedReq.ToolConfig is non-nil and that
receivedReq.ToolConfig.CachePoint.Type == "default" (mirroring the system
assertions); this will exercise buildBedrockRequest’s ToolConfig.CachePoint path
when tools are present.
- Around line 286-380: Convert the three standalone tests
(TestBedrock_Send_CachePointInSystem, TestBedrock_Send_CacheUsageParsing,
TestBedrock_Send_CacheUsageOmitted) into a single table-driven test: define a
[]struct with fields like name, serverResponse (bedrockResponse used by httptest
server), request (the Request passed to b.Send), wantSystemBlocks (count) /
wantSystemCachePoint (bool) / wantCacheReadTokens (int) / wantCacheWriteTokens
(int), then range over cases and run each case as t.Run(case.name, func(t
*testing.T){...}). Inside the loop create the httptest.NewServer returning
case.serverResponse, instantiate Bedrock with NewBedrock(...,
WithBedrockBaseURL(server.URL), ...), call b.Send with case.request, and assert
expectations (receivedReq.System and its CachePoint for the system-cache case
and resp.CacheReadTokens/resp.CacheWriteTokens for usage cases) using the
case.want* fields; keep existing JSON encode/decode logic and defer server.Close
per subtest. Ensure unique identifiers from the diff—NewBedrock,
WithBedrockBaseURL, b.Send, and the bedrockResponse/bedrockUsage fields—are used
to locate and adapt the code.
In `@internal/provider/openai_test.go`:
- Around line 1670-1761: Combine the three tests
TestOpenAI_Send_CacheUsageParsing, TestOpenAI_Send_CacheUsageOmitted and
TestOpenAI_SendStream_CacheUsageParsing into one table-driven test that iterates
over cases with fields like name, isStream (bool), responsePayload (string), and
wantCacheReadTokens (int); for each case spin up an httptest.Server that returns
the case.responsePayload (for stream cases set Content-Type: text/event-stream
and include the SSE messages used in the original test), create the client via
NewOpenAI(..., WithOpenAIBaseURL(server.URL)), call o.Send for non-stream cases
and o.SendStream for stream cases (remember to Close the stream), then assert
resp.CacheReadTokens or event.CacheReadTokens equals wantCacheReadTokens; keep
the original expectations and error handling but remove duplicated
setup/teardown by running them inside the loop.
In `@internal/provider/provider_test.go`:
- Around line 190-212: Convert the two tests into table-driven style: replace
TestResponse_CacheTokens and TestRequest_CacheConfig with a single or separate
table-driven tests that define a []struct{name string, resp *Response, req
*Request, wantRead int, wantWrite int, wantCacheConfig bool} (or two small
tables) and iterate using t.Run for each case, verifying
Response.CacheReadTokens and CacheWriteTokens and Request.CacheConfig against
want values; locate the types Response and Request and the fields
CacheReadTokens, CacheWriteTokens, and CacheConfig to implement the table
entries and assertions.
In `@pkg/runner/result_test.go`:
- Around line 228-276: The two new tests TestResultMarshalJSON_CacheTokens and
TestResultMarshalJSON_CacheTokensOmittedWhenZero duplicate the pattern used by
TestResultMarshalJSON_TableDriven; merge their cases into the existing
table-driven test by adding two new entries to the []struct{name, input, want}
slice used by TestResultMarshalJSON_TableDriven: one case with Result{...
CacheReadTokens:42, CacheWriteTokens:7} expecting JSON keys
"cache_read_tokens"=42 and "cache_write_tokens"=7, and one case with zero values
expecting those keys to be omitted; then remove the two standalone test
functions and rely on the extended TestResultMarshalJSON_TableDriven that
serializes Result and asserts parsed JSON as before (refer to Result,
CacheReadTokens, CacheWriteTokens, and TestResultMarshalJSON_TableDriven).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0ec6bae3-1e56-4838-a0bb-afcd76c12307
⛔ Files ignored due to path filters (2)
docs/plans/049_prompt_caching_implement.mdis excluded by!docs/**docs/plans/049_prompt_caching_spec.mdis excluded by!docs/**
📒 Files selected for processing (16)
internal/provider/anthropic.gointernal/provider/anthropic_stream_types.gointernal/provider/anthropic_test.gointernal/provider/bedrock.gointernal/provider/bedrock_test.gointernal/provider/openai.gointernal/provider/openai_test.gointernal/provider/provider.gointernal/provider/provider_test.gointernal/provider/stream.gointernal/testutil/mockserver.gopkg/runner/result.gopkg/runner/result_test.gopkg/runner/run.gopkg/runner/run_test.gopkg/runner/stream_test.go
…closure vars in message_delta anthropic: add TestAnthropic_SendStream_CacheUsageParsing bedrock: consolidate cache-related tests into table-driven TestBedrock_Send_Cache provider: convert cache token and cache config tests to table-driven
There was a problem hiding this comment.
🧹 Nitpick comments (5)
internal/provider/bedrock_test.go (2)
287-295: ⚡ Quick win
setupRequest func() *Requestadds unnecessary indirection — consider a plainrequest *Requestfield.All four test cases construct a static
*Requestwith no dynamic logic, so wrapping it in a factory function doesn't buy anything. Replacing it with a direct field also better matches the prescribed[]struct{name, input, want}table-driven pattern.♻️ Suggested simplification
cases := []struct { name string - setupRequest func() *Request + request *Request serverResponse bedrockResponse wantSystemCachePoint bool wantToolCachePoint bool wantCacheReadTokens int wantCacheWriteTokens int }{ { name: "system block gets cache point", - setupRequest: func() *Request { - return &Request{ - Model: "test-model", - System: "You are helpful.", - CacheConfig: true, - Messages: []Message{{Role: "user", Content: "hi"}}, - } + request: &Request{ + Model: "test-model", + System: "You are helpful.", + CacheConfig: true, + Messages: []Message{{Role: "user", Content: "hi"}}, },…and in the loop body:
-resp, err := b.Send(context.Background(), tc.setupRequest()) +resp, err := b.Send(context.Background(), tc.request)As per coding guidelines, table-driven tests should use the
[]struct{name, input, want}pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/bedrock_test.go` around lines 287 - 295, Replace the indirect test-field setupRequest func() *Request with a concrete request *Request in the test table: change the cases slice struct to include request *Request instead of setupRequest, update each test case to assign the static *Request directly, and in the test loop call tc.request (not tc.setupRequest()). Update any references in the loop (and subtests) to use tc.request when invoking the code under test so the table follows the standard []struct{name, input, want} pattern; keep the other want fields (wantSystemCachePoint, wantToolCachePoint, wantCacheReadTokens, wantCacheWriteTokens) unchanged.
296-374: ⚡ Quick winConsider adding a negative case:
CacheConfig: falseshould produce no cache points.The four existing cases verify that cache points are added when
CacheConfig: true, and that token counts parse correctly — but there's no case that confirms the system block and tool config are not decorated with a cache point whenCacheConfigisfalse(or omitted). That's a pretty easy regression to miss, and a cheap case to add.{ name: "no cache point when CacheConfig false", request: &Request{ Model: "test-model", System: "You are helpful.", Messages: []Message{{Role: "user", Content: "hi"}}, Tools: []Tool{{ Name: "read_file", Description: "Read a file", Parameters: map[string]ToolParameter{"path": {Type: "string", Description: "File path", Required: true}}, }}, // CacheConfig intentionally omitted / false }, serverResponse: bedrockResponse{ Output: bedrockOutput{Message: &bedrockMessage{Role: "assistant", Content: []bedrockBlock{{Text: "ok"}}}}, StopReason: "end_turn", }, wantSystemCachePoint: false, wantToolCachePoint: false, },You'd also want the assertion block to verify absence of a cache point when the
want*flag isfalse:if !tc.wantSystemCachePoint && len(receivedReq.System) > 0 && receivedReq.System[0].CachePoint != nil { t.Error("unexpected cachePoint on system block") } if !tc.wantToolCachePoint && receivedReq.ToolConfig != nil && receivedReq.ToolConfig.CachePoint != nil { t.Error("unexpected cachePoint on ToolConfig") }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/bedrock_test.go` around lines 296 - 374, Add a negative test case to the existing test table that mirrors the "system block gets cache point" / "tool config gets cache point" cases but with Request.CacheConfig set to false or omitted (use Request with System, Messages and Tools populated) and expect wantSystemCachePoint=false and wantToolCachePoint=false; update the assertion section that inspects receivedReq.System and receivedReq.ToolConfig (the variables used in the test) to assert absence of CachePoint when tc.wantSystemCachePoint or tc.wantToolCachePoint is false (i.e., fail if receivedReq.System[0].CachePoint != nil when wantSystemCachePoint==false, and fail if receivedReq.ToolConfig != nil && receivedReq.ToolConfig.CachePoint != nil when wantToolCachePoint==false).internal/provider/openai_test.go (1)
1669-1743: ⚡ Quick winAlign table cases to the repo’s
name/input/wanttest shape.This is table-driven already, but the case struct should follow the project’s standard
name, input, wantpattern for consistency across test files.♻️ Suggested refactor
- cases := []struct { - name string - isStream bool - responsePayload []byte - wantCacheRead int - }{ + cases := []struct { + name string + input struct { + isStream bool + responsePayload []byte + } + want struct { + cacheRead int + } + }{ { name: "send with cached_tokens", - responsePayload: []byte(`{ + input: struct { + isStream bool + responsePayload []byte + }{ + responsePayload: []byte(`{ "model": "gpt-4o", "choices": [{"message": {"content": "Hello from cached prompt"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 100, "completion_tokens": 5, "prompt_tokens_details": {"cached_tokens": 80}} - }`), - wantCacheRead: 80, + }`), + }, + want: struct{ cacheRead int }{cacheRead: 80}, }, ... - if tc.isStream { + if tc.input.isStream { w.Header().Set("Content-Type", "text/event-stream") } - _, _ = w.Write(tc.responsePayload) + _, _ = w.Write(tc.input.responsePayload) })) ... - if tc.isStream { + if tc.input.isStream { ... } - if got != tc.wantCacheRead { - t.Errorf("CacheReadTokens = %d, want %d", got, tc.wantCacheRead) + if got != tc.want.cacheRead { + t.Errorf("CacheReadTokens = %d, want %d", got, tc.want.cacheRead) }As per coding guidelines,
**/*_test.go: Use table-driven tests with[]struct{name, input, want}pattern for all test files.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/openai_test.go` around lines 1669 - 1743, The table-driven test currently uses fields isStream/responsePayload/wantCacheRead; rename the case struct to follow the repo pattern name, input, want by creating an input struct (e.g., Input{IsStream bool, Payload []byte}) and a want struct (e.g., Want{CacheRead int}), then update all usages in the loop and subtests (references: the cases slice, the per-case tc variable, NewOpenAI, Send, SendStream, ev.CacheReadTokens, resp.CacheReadTokens) to read tc.input.IsStream / tc.input.Payload and assert against tc.want.CacheRead to keep the same behavior while matching the name/input/want shape.internal/provider/anthropic.go (1)
217-236: ⚡ Quick winConsider using
anthropicContentBlockinstead of a raw map literal.
anthropicContentBlock(lines 84–94) already definesCacheControl map[string]interface{}— using it here keeps the serialization logic consistent with the rest of the file and gets compile-time field-name checking for free.♻️ Proposed refactor
- return []map[string]interface{}{ - { - "type": "text", - "text": system, - "cache_control": map[string]interface{}{ - "type": "ephemeral", - }, - }, - } + return []anthropicContentBlock{ + { + Type: "text", + Text: system, + CacheControl: map[string]interface{}{ + "type": "ephemeral", + }, + }, + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/anthropic.go` around lines 217 - 236, The buildAnthropicSystem function uses a raw map literal for content blocks; replace that with the existing anthropicContentBlock struct to ensure consistent serialization and compile-time field checks — when cacheEnabled is true, construct an anthropicContentBlock with Type="text", Text=system and CacheControl set to a map (or a dedicated struct) with Type="ephemeral", and return a slice of anthropicContentBlock (or []interface{} containing the struct) instead of []map[string]interface{} so callers and JSON tags remain consistent with the types defined around anthropicContentBlock.internal/provider/anthropic_stream_test.go (1)
214-263: ⚡ Quick winAdd a table-driven structure and a guard on
gotDone.Type.Two small things:
- Per coding guidelines, tests should use the
[]struct{name, input, want}table-driven pattern. This test currently only covers one scenario, but structuring it as a table makes it easy to add edge cases later (e.g., only write tokens present, only read tokens, neither).- If the stream never emits a
StreamEventDone(e.g., a regression drops themessage_deltahandler),gotDonestays zero-valued and the token assertions will produce confusing output. A quick guard surfaces the root cause immediately:if gotDone.Type != StreamEventDone { t.Fatalf("no StreamEventDone received; got type %q", gotDone.Type) }As per coding guidelines, "Use table-driven tests with
[]struct{name, input, want}pattern for all test files."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/provider/anthropic_stream_test.go` around lines 214 - 263, The test TestAnthropic_SendStream_CacheUsageParsing should be converted to a table-driven structure (use []struct{name, input, want} so additional cases like only read/write/none can be added) and must include a guard that verifies a StreamEventDone was actually received before asserting token fields (check gotDone.Type == StreamEventDone and call t.Fatalf if not); update the test harness around stream.Next() and the final assertions to use the table entries' expected token values (InputTokens, OutputTokens, CacheReadTokens, CacheWriteTokens) and reference the same variables (stream, gotDone, StreamEventDone) when implementing the guard and per-case checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/provider/anthropic_stream_test.go`:
- Around line 214-263: The test TestAnthropic_SendStream_CacheUsageParsing
should be converted to a table-driven structure (use []struct{name, input, want}
so additional cases like only read/write/none can be added) and must include a
guard that verifies a StreamEventDone was actually received before asserting
token fields (check gotDone.Type == StreamEventDone and call t.Fatalf if not);
update the test harness around stream.Next() and the final assertions to use the
table entries' expected token values (InputTokens, OutputTokens,
CacheReadTokens, CacheWriteTokens) and reference the same variables (stream,
gotDone, StreamEventDone) when implementing the guard and per-case checks.
In `@internal/provider/anthropic.go`:
- Around line 217-236: The buildAnthropicSystem function uses a raw map literal
for content blocks; replace that with the existing anthropicContentBlock struct
to ensure consistent serialization and compile-time field checks — when
cacheEnabled is true, construct an anthropicContentBlock with Type="text",
Text=system and CacheControl set to a map (or a dedicated struct) with
Type="ephemeral", and return a slice of anthropicContentBlock (or []interface{}
containing the struct) instead of []map[string]interface{} so callers and JSON
tags remain consistent with the types defined around anthropicContentBlock.
In `@internal/provider/bedrock_test.go`:
- Around line 287-295: Replace the indirect test-field setupRequest func()
*Request with a concrete request *Request in the test table: change the cases
slice struct to include request *Request instead of setupRequest, update each
test case to assign the static *Request directly, and in the test loop call
tc.request (not tc.setupRequest()). Update any references in the loop (and
subtests) to use tc.request when invoking the code under test so the table
follows the standard []struct{name, input, want} pattern; keep the other want
fields (wantSystemCachePoint, wantToolCachePoint, wantCacheReadTokens,
wantCacheWriteTokens) unchanged.
- Around line 296-374: Add a negative test case to the existing test table that
mirrors the "system block gets cache point" / "tool config gets cache point"
cases but with Request.CacheConfig set to false or omitted (use Request with
System, Messages and Tools populated) and expect wantSystemCachePoint=false and
wantToolCachePoint=false; update the assertion section that inspects
receivedReq.System and receivedReq.ToolConfig (the variables used in the test)
to assert absence of CachePoint when tc.wantSystemCachePoint or
tc.wantToolCachePoint is false (i.e., fail if receivedReq.System[0].CachePoint
!= nil when wantSystemCachePoint==false, and fail if receivedReq.ToolConfig !=
nil && receivedReq.ToolConfig.CachePoint != nil when wantToolCachePoint==false).
In `@internal/provider/openai_test.go`:
- Around line 1669-1743: The table-driven test currently uses fields
isStream/responsePayload/wantCacheRead; rename the case struct to follow the
repo pattern name, input, want by creating an input struct (e.g., Input{IsStream
bool, Payload []byte}) and a want struct (e.g., Want{CacheRead int}), then
update all usages in the loop and subtests (references: the cases slice, the
per-case tc variable, NewOpenAI, Send, SendStream, ev.CacheReadTokens,
resp.CacheReadTokens) to read tc.input.IsStream / tc.input.Payload and assert
against tc.want.CacheRead to keep the same behavior while matching the
name/input/want shape.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 02d2e4d1-516d-4707-8572-c02531494d1f
📒 Files selected for processing (7)
internal/agent/agent.gointernal/provider/anthropic.gointernal/provider/anthropic_stream_test.gointernal/provider/bedrock_test.gointernal/provider/openai_test.gointernal/provider/provider_test.gopkg/runner/run.go
✅ Files skipped from review due to trivial changes (1)
- internal/provider/provider_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/runner/run.go
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
1 participant
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
This pull request implements prompt caching across Anthropic, Bedrock, and OpenAI providers and adds end-to-end cache token accounting. It introduces a request-level cache toggle, propagates cache read/write token counts through provider responses and stream events, ensures these counts flow through run execution into final Results (with conditional JSON emission), and adds provider-specific wire-format support and tests to exercise caching behavior.
Changelog
Added
CacheConfigboolean onRequestto enable prompt caching.CacheReadTokensandCacheWriteTokensonResponse,StreamEvent, andResult(with conditional JSON serialization).CacheConfigis enabled and streaming parsing to capture cache token usage.CacheConfigis enabled.prompt_tokens_details.cached_tokensinto cache read tokens (non-stream and streaming).AnthropicResponseWithCacheTokensandAnthropicToolUseResponseWithCacheTokens.Changed
cache_controlwhen caching is enabled; request/stream handling maps cache usage into response fields.CacheConfigis enabled; response parsing maps cache token usage into response fields.CacheReadTokens.Send/SendStreamand helpers updated to build cache-aware system/tool request fields whenreq.CacheConfigis set.Result, and includes cache counts in verbose output paths; streaming drain logic propagates cache tokens fromStreamEventDone.Fixed
satisfies #83