From 95d5b2485f7d7feedfe9e50efc1ee480b0e7a384 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 25 Jul 2026 23:19:51 +0800 Subject: [PATCH] fix(runtime): skip deferred tools when injecting cache_control - Updated `ensureCacheControl` to inject cache_control into the last non-deferred tool instead of the last tool. - Adjusted `injectToolsCacheControl` to ignore tools marked with `defer_loading` for prompt caching. - Added unit tests to validate behavior for deferred tools, existing cache_control, and mixed tool arrays. Closes: #4561 --- .../runtime/executor/caching_verify_test.go | 87 +++++++++++++++++++ internal/runtime/executor/claude_executor.go | 26 +++--- 2 files changed, 99 insertions(+), 14 deletions(-) diff --git a/internal/runtime/executor/caching_verify_test.go b/internal/runtime/executor/caching_verify_test.go index 6088d304..0ac96d37 100644 --- a/internal/runtime/executor/caching_verify_test.go +++ b/internal/runtime/executor/caching_verify_test.go @@ -215,6 +215,93 @@ func TestEnsureCacheControl(t *testing.T) { }) } +func TestInjectToolsCacheControlSkipsDeferredTools(t *testing.T) { + tests := []struct { + name string + input string + wantCacheIndex int + wantCacheTTL string + }{ + { + name: "trailing deferred tool", + input: `{"tools":[ + {"name":"resident","defer_loading":false}, + {"name":"deferred","defer_loading":true} + ]}`, + wantCacheIndex: 0, + }, + { + name: "multiple trailing deferred tools", + input: `{"tools":[ + {"name":"resident"}, + {"name":"deferred_1","defer_loading":true}, + {"name":"deferred_2","defer_loading":true} + ]}`, + wantCacheIndex: 0, + }, + { + name: "middle deferred tool", + input: `{"tools":[ + {"name":"resident_1"}, + {"name":"deferred","defer_loading":true}, + {"name":"resident_2"} + ]}`, + wantCacheIndex: 2, + }, + { + name: "all tools deferred", + input: `{"tools":[ + {"name":"deferred_1","defer_loading":true}, + {"name":"deferred_2","defer_loading":true} + ]}`, + wantCacheIndex: -1, + }, + { + name: "existing cache control", + input: `{"tools":[ + {"name":"resident_1","cache_control":{"type":"ephemeral","ttl":"1h"}}, + {"name":"resident_2"} + ]}`, + wantCacheIndex: 0, + wantCacheTTL: "1h", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + output := injectToolsCacheControl([]byte(tt.input)) + tools := gjson.GetBytes(output, "tools").Array() + cacheCount := 0 + for index, tool := range tools { + cacheControl := tool.Get("cache_control") + if cacheControl.Exists() { + cacheCount++ + if index != tt.wantCacheIndex { + t.Errorf("cache_control added to tool %d, want tool %d: %s", index, tt.wantCacheIndex, string(output)) + } + } + if tool.Get("defer_loading").Bool() && cacheControl.Exists() { + t.Errorf("deferred tool %d must not have cache_control: %s", index, string(output)) + } + } + + wantCacheCount := 1 + if tt.wantCacheIndex < 0 { + wantCacheCount = 0 + } + if cacheCount != wantCacheCount { + t.Errorf("cache_control count = %d, want %d: %s", cacheCount, wantCacheCount, string(output)) + } + if tt.wantCacheTTL != "" { + path := fmt.Sprintf("tools.%d.cache_control.ttl", tt.wantCacheIndex) + if got := gjson.GetBytes(output, path).String(); got != tt.wantCacheTTL { + t.Errorf("cache_control TTL = %q, want %q: %s", got, tt.wantCacheTTL, string(output)) + } + } + }) + } +} + // TestCacheControlOrder verifies the correct order: tools -> system -> messages func TestCacheControlOrder(t *testing.T) { input := []byte(`{ diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index abbe6fc1..8e10f989 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -2206,7 +2206,7 @@ func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.A // ensureCacheControl injects cache_control breakpoints into the payload for optimal prompt caching. // According to Anthropic's documentation, cache prefixes are created in order: tools -> system -> messages. // This function adds cache_control to: -// 1. The LAST tool in the tools array (caches all tool definitions) +// 1. The LAST non-deferred tool in the tools array (caches all preceding tool definitions) // 2. The LAST system prompt element // 3. The SECOND-TO-LAST user turn (caches conversation history for multi-turn) // @@ -2214,7 +2214,7 @@ func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.A // This enables up to 90% cost reduction on cached tokens (cache read = 0.1x base price). // See: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching func ensureCacheControl(payload []byte) []byte { - // 1. Inject cache_control into the LAST tool (caches all tool definitions) + // 1. Inject cache_control into the LAST non-deferred tool // Tools are cached first in the hierarchy, so this is the most important breakpoint. payload = injectToolsCacheControl(payload) @@ -2623,8 +2623,8 @@ func injectMessagesCacheControl(payload []byte) []byte { return payload } -// injectToolsCacheControl adds cache_control to the last tool in the tools array. -// Per Anthropic docs: "The cache_control parameter on the last tool definition caches all tool definitions." +// injectToolsCacheControl adds cache_control to the last non-deferred tool in the tools array. +// Deferred tools cannot use prompt caching, so trailing deferred tools are skipped. // This only adds cache_control if NO tool in the array already has it. func injectToolsCacheControl(payload []byte) []byte { tools := gjson.GetBytes(payload, "tools") @@ -2632,26 +2632,24 @@ func injectToolsCacheControl(payload []byte) []byte { return payload } - toolCount := int(tools.Get("#").Int()) - if toolCount == 0 { - return payload - } - - // Check if ANY tool already has cache_control - if so, don't modify tools + // Check if ANY tool already has cache_control and find the last eligible tool. hasCacheControlInTools := false - tools.ForEach(func(_, tool gjson.Result) bool { + lastEligibleToolIndex := -1 + tools.ForEach(func(index, tool gjson.Result) bool { if tool.Get("cache_control").Exists() { hasCacheControlInTools = true return false } + if !tool.Get("defer_loading").Bool() { + lastEligibleToolIndex = int(index.Int()) + } return true }) - if hasCacheControlInTools { + if hasCacheControlInTools || lastEligibleToolIndex < 0 { return payload } - // Add cache_control to the last tool - lastToolPath := fmt.Sprintf("tools.%d.cache_control", toolCount-1) + lastToolPath := fmt.Sprintf("tools.%d.cache_control", lastEligibleToolIndex) result, err := sjson.SetBytes(payload, lastToolPath, map[string]string{"type": "ephemeral"}) if err != nil { log.Warnf("failed to inject cache_control into tools array: %v", err) -- 2.54.0