From eb6a768d103da08c039c20dfad28aad2c936fbbc Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 5 Oct 2026 14:16:37 +0800 Subject: [PATCH] feat(claude): support explicit prompt cache options - Honor `prompt_cache_options.mode = "explicit"` to suppress automatic `cache_control` insertion while preserving client-defined cache breakpoints - Strip `prompt_cache_options` from request payloads before forwarding upstream and during token counting Closes: #6380 --- .../runtime/executor/caching_verify_test.go | 3 + .../executor/claude_executor_cloaking.go | 261 +++-- .../executor/claude_executor_execute.go | 15 +- .../executor/claude_executor_stream.go | 15 +- .../executor/claude_executor_tokens.go | 10 +- .../claude_prompt_cache_options_test.go | 950 ++++++++++++++++++ 6 files changed, 1189 insertions(+), 65 deletions(-) create mode 100644 internal/runtime/executor/claude_prompt_cache_options_test.go diff --git a/internal/runtime/executor/caching_verify_test.go b/internal/runtime/executor/caching_verify_test.go index 198bcd9b4..dd9c9f354 100644 --- a/internal/runtime/executor/caching_verify_test.go +++ b/internal/runtime/executor/caching_verify_test.go @@ -441,6 +441,9 @@ func TestShouldEnsureCacheControl(t *testing.T) { {name: "cloaked with marker", payload: withMarker, cloaked: true, want: true}, {name: "unconfirmed markerless", payload: markerless, want: true}, {name: "unconfirmed with marker", payload: withMarker, want: false}, + {name: "explicit mode markerless", payload: []byte(`{"prompt_cache_options":{"mode":"explicit"},"messages":[{"role":"user","content":"x"}]}`), want: false}, + {name: "explicit mode cloaked", payload: []byte(`{"prompt_cache_options":{"mode":"explicit"},"messages":[{"role":"user","content":"x"}]}`), cloaked: true, want: false}, + {name: "explicit mode with marker", payload: []byte(`{"prompt_cache_options":{"mode":"explicit"},"messages":[{"role":"user","content":[{"type":"text","text":"x","cache_control":{"type":"ephemeral"}}]}]}`), cloaked: true, want: false}, } for _, tt := range tests { diff --git a/internal/runtime/executor/claude_executor_cloaking.go b/internal/runtime/executor/claude_executor_cloaking.go index 14bcd5b3d..f43102b6e 100644 --- a/internal/runtime/executor/claude_executor_cloaking.go +++ b/internal/runtime/executor/claude_executor_cloaking.go @@ -320,7 +320,7 @@ func checkSystemInstructionsWithMode(payload []byte, strictMode bool) []byte { // Claude models give it operator-level authority without changing the cached // top-level prefix. func checkSystemInstructionsWithSigningMode(payload []byte, strictMode bool, cchSigning bool, version, entrypoint, workload string) []byte { - return checkSystemInstructionsWithSigningModeAt(payload, strictMode, cchSigning, version, entrypoint, workload, claudeCodeLocalDate(time.Now()), false, "", "") + return checkSystemInstructionsWithSigningModeAt(payload, strictMode, cchSigning, version, entrypoint, workload, claudeCodeLocalDate(time.Now()), false, "", "", "") } // isClaudeFable51Model reports whether the model is specifically Fable 5.1 / Mythos 5.1, @@ -339,6 +339,65 @@ func isClaudeFable51Model(model string) bool { return false } +type forwardedClaudeSystemPromptBlock struct { + Text string + CacheControl *claudeCacheControl + RawCC string +} + +func buildForwardedSystemBlock(text string, block forwardedClaudeSystemPromptBlock, isExplicit bool) string { + if isExplicit { + if block.RawCC != "" { + return `{"type":"text","text":` + marshalJSONStringWithoutHTMLEscape(text) + `,"cache_control":` + block.RawCC + `}` + } + return buildTextBlock(text, nil) + } + return buildTextBlock(text, &claudeCodeCacheControl) +} + +func isValidClaudeCacheControl(cc gjson.Result) bool { + if !cc.Exists() || cc.Type == gjson.Null || !cc.IsObject() { + return false + } + typ := cc.Get("type") + return typ.Type == gjson.String && typ.String() == "ephemeral" +} + +func collectForwardedClaudeSystemPromptBlocksWithCache(system gjson.Result) []forwardedClaudeSystemPromptBlock { + var blocks []forwardedClaudeSystemPromptBlock + appendBlock := func(text string, cc *claudeCacheControl, rawCC string) { + if strings.TrimSpace(text) == "" || util.IsClaudeCodeAttributionSystemText(text) || text == claudeCodeCLIIdentity { + return + } + blocks = append(blocks, forwardedClaudeSystemPromptBlock{Text: text, CacheControl: cc, RawCC: rawCC}) + } + + if system.Type == gjson.String { + appendBlock(system.String(), nil, "") + return blocks + } + if !system.IsArray() { + return blocks + } + system.ForEach(func(_, item gjson.Result) bool { + if item.Get("type").String() == "text" { + var cc *claudeCacheControl + rawCC := "" + if itemCC := item.Get("cache_control"); isValidClaudeCacheControl(itemCC) { + cc = &claudeCacheControl{ + Type: itemCC.Get("type").String(), + TTL: itemCC.Get("ttl").String(), + Scope: itemCC.Get("scope").String(), + } + rawCC = itemCC.Raw + } + appendBlock(item.Get("text").String(), cc, rawCC) + } + return true + }) + return blocks +} + func checkSystemInstructionsWithSigningModeAt( payload []byte, strictMode bool, @@ -354,7 +413,13 @@ func checkSystemInstructionsWithSigningModeAt( billingText := generateBillingHeader(cchSigning, version, messageText, entrypoint, workload, isSubagent, prevReq, promptID, turnOrigin...) billingBlock := buildTextBlock(billingText, nil) - agentBlock := buildTextBlock(claudeCodeCLIIdentity, &claudeCodeCacheControl) + + isExplicit := (len(turnOrigin) > 1 && turnOrigin[1] == "explicit_cache_mode") || isExplicitPromptCacheMode(payload) + var agentCC *claudeCacheControl = &claudeCodeCacheControl + if isExplicit { + agentCC = nil + } + agentBlock := buildTextBlock(claudeCodeCLIIdentity, agentCC) systemBlocks := []string{billingBlock, agentBlock} model := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "model").String())) @@ -363,34 +428,31 @@ func checkSystemInstructionsWithSigningModeAt( } payload, _ = sjson.SetRawBytes(payload, "system", []byte("["+strings.Join(systemBlocks, ",")+"]")) if strictMode { - return injectClaudeCodeCurrentDate(payload, currentDate) + return injectClaudeCodeCurrentDateInternal(payload, currentDate, isExplicit) } - forwardedSystemBlocks := collectForwardedClaudeSystemPromptBlocks(system) + forwardedSystemBlocks := collectForwardedClaudeSystemPromptBlocksWithCache(system) if len(forwardedSystemBlocks) == 0 { - return injectClaudeCodeCurrentDate(payload, currentDate) + return injectClaudeCodeCurrentDateInternal(payload, currentDate, isExplicit) } if claudeHistoryHasAdvisorCallOrResult(payload) { for _, block := range forwardedSystemBlocks { - systemBlocks = append(systemBlocks, buildTextBlock(block, nil)) + systemBlocks = append(systemBlocks, buildForwardedSystemBlock(block.Text, block, isExplicit)) } payload, _ = sjson.SetRawBytes(payload, "system", []byte("["+strings.Join(systemBlocks, ",")+"]")) - return injectClaudeCodeCurrentDate(payload, currentDate) + return injectClaudeCodeCurrentDateInternal(payload, currentDate, isExplicit) } if claudeUsesLegacySystemReminder(payload) { - payload = prependClaudeSystemRemindersToFirstUserMessage(payload, forwardedSystemBlocks) + payload = prependClaudeSystemReminderBlocksToFirstUserMessage(payload, forwardedSystemBlocks, isExplicit) } else if claudeMidConversationSystemMessagesAtEnd(payload) { for _, block := range forwardedSystemBlocks { - systemBlocks = append(systemBlocks, buildTextBlock(block, &claudeCodeCacheControl)) + systemBlocks = append(systemBlocks, buildForwardedSystemBlock(block.Text, block, isExplicit)) } payload, _ = sjson.SetRawBytes(payload, "system", []byte("["+strings.Join(systemBlocks, ",")+"]")) } else { - // Unknown and future model IDs optimistically use the authoritative - // mid-conversation system role. Only empirically unsupported legacy IDs - // stay on the user-reminder compatibility path. - payload = insertClaudeMidConversationSystemMessages(payload, forwardedSystemBlocks) + payload = insertClaudeMidConversationSystemBlocks(payload, forwardedSystemBlocks, isExplicit) } - return injectClaudeCodeCurrentDate(payload, currentDate) + return injectClaudeCodeCurrentDateInternal(payload, currentDate, isExplicit) } // relocateClaudeSystemPromptForCountTokens keeps a cloaked count_tokens request @@ -400,16 +462,17 @@ func checkSystemInstructionsWithSigningModeAt( // using the same positional mapping as the Messages path. That keeps the counted // tokens aligned with the request the caller is about to send while preventing a // third-party system prompt from reaching Anthropic in the system slot. -func relocateClaudeSystemPromptForCountTokens(payload []byte, strictMode bool) []byte { +func relocateClaudeSystemPromptForCountTokens(payload []byte, strictMode bool, explicitCacheMode ...bool) []byte { system := gjson.GetBytes(payload, "system") if !system.Exists() { return payload } + isExplicit := (len(explicitCacheMode) > 0 && explicitCacheMode[0]) || isExplicitPromptCacheMode(payload) // Strict mode drops caller prompts on the Messages path, so it must not // reintroduce them here either. - var forwardedSystemBlocks []string + var forwardedSystemBlocks []forwardedClaudeSystemPromptBlock if !strictMode { - forwardedSystemBlocks = collectForwardedClaudeSystemPromptBlocks(system) + forwardedSystemBlocks = collectForwardedClaudeSystemPromptBlocksWithCache(system) } if len(forwardedSystemBlocks) == 0 { updated, _ := sjson.DeleteBytes(payload, "system") @@ -422,7 +485,7 @@ func relocateClaudeSystemPromptForCountTokens(payload []byte, strictMode bool) [ // remains aligned with the Messages path without splicing messages[]. blocks := make([]string, 0, len(forwardedSystemBlocks)) for _, block := range forwardedSystemBlocks { - blocks = append(blocks, buildTextBlock(block, nil)) + blocks = append(blocks, buildForwardedSystemBlock(block.Text, block, isExplicit)) } updated, _ := sjson.SetRawBytes(payload, "system", []byte("["+strings.Join(blocks, ",")+"]")) return updated @@ -433,17 +496,17 @@ func relocateClaudeSystemPromptForCountTokens(payload []byte, strictMode bool) [ } payload = updated if claudeUsesLegacySystemReminder(payload) { - return prependClaudeSystemRemindersToFirstUserMessage(payload, forwardedSystemBlocks) + return prependClaudeSystemReminderBlocksToFirstUserMessage(payload, forwardedSystemBlocks, isExplicit) } if claudeMidConversationSystemMessagesAtEnd(payload) { blocks := make([]string, 0, len(forwardedSystemBlocks)) for _, block := range forwardedSystemBlocks { - blocks = append(blocks, buildTextBlock(block, &claudeCodeCacheControl)) + blocks = append(blocks, buildForwardedSystemBlock(block.Text, block, isExplicit)) } updated, _ := sjson.SetRawBytes(payload, "system", []byte("["+strings.Join(blocks, ",")+"]")) return updated } - return insertClaudeMidConversationSystemMessages(payload, forwardedSystemBlocks) + return insertClaudeMidConversationSystemBlocks(payload, forwardedSystemBlocks, isExplicit) } // claudeLegacySystemReminderModels lists the official Anthropic model IDs and @@ -635,6 +698,9 @@ func buildTextBlock(text string, cacheControl *claudeCacheControl) string { if cacheControl.TTL != "" { block += `,"ttl":` + marshalJSONStringWithoutHTMLEscape(cacheControl.TTL) } + if cacheControl.Scope != "" { + block += `,"scope":` + marshalJSONStringWithoutHTMLEscape(cacheControl.Scope) + } block += "}" } return block + "}" @@ -649,58 +715,71 @@ func marshalJSONStringWithoutHTMLEscape(value string) string { } func prependClaudeSystemRemindersToFirstUserMessage(payload []byte, texts []string) []byte { - firstUserIdx := firstClaudeUserMessageIndex(payload) - if firstUserIdx < 0 || len(texts) == 0 { - return payload - } - - reminderTexts := make([]string, 0, len(texts)) + blocks := make([]forwardedClaudeSystemPromptBlock, 0, len(texts)) for _, text := range texts { - reminderTexts = append(reminderTexts, claudeCallerSystemReminder(text)) + blocks = append(blocks, forwardedClaudeSystemPromptBlock{Text: text}) + } + return prependClaudeSystemReminderBlocksToFirstUserMessage(payload, blocks, false) +} + +func prependClaudeSystemReminderBlocksToFirstUserMessage(payload []byte, blocks []forwardedClaudeSystemPromptBlock, isExplicit bool) []byte { + firstUserIdx := firstClaudeUserMessageIndex(payload) + if firstUserIdx < 0 || len(blocks) == 0 { + return payload } contentPath := fmt.Sprintf("messages.%d.content", firstUserIdx) content := gjson.GetBytes(payload, contentPath) if content.IsArray() { - blocks := content.Array() - existing := make(map[string]int, len(blocks)) - for _, block := range blocks { + contentBlocks := content.Array() + existing := make(map[string]int, len(contentBlocks)) + for _, block := range contentBlocks { if block.Get("type").String() == "text" { existing[block.Get("text").String()]++ } } - reminderBlocks := make([]string, 0, len(reminderTexts)) - for _, reminderText := range reminderTexts { + reminderBlocks := make([]string, 0, len(blocks)) + for _, b := range blocks { + reminderText := claudeCallerSystemReminder(b.Text) if existing[reminderText] > 0 { existing[reminderText]-- continue } - reminderBlocks = append(reminderBlocks, buildTextBlock(reminderText, nil)) + if isExplicit && b.RawCC != "" { + reminderBlocks = append(reminderBlocks, `{"type":"text","text":`+marshalJSONStringWithoutHTMLEscape(reminderText)+`,"cache_control":`+b.RawCC+`}`) + } else { + reminderBlocks = append(reminderBlocks, buildTextBlock(reminderText, nil)) + } } if len(reminderBlocks) == 0 { return payload } insertAt := 0 - for insertAt < len(blocks) && blocks[insertAt].Get("type").String() == "tool_result" { + for insertAt < len(contentBlocks) && contentBlocks[insertAt].Get("type").String() == "tool_result" { insertAt++ } - rawBlocks := make([]string, 0, len(blocks)+len(reminderBlocks)) - for idx, block := range blocks { + rawBlocks := make([]string, 0, len(contentBlocks)+len(reminderBlocks)) + for idx, block := range contentBlocks { if idx == insertAt { rawBlocks = append(rawBlocks, reminderBlocks...) } rawBlocks = append(rawBlocks, block.Raw) } - if insertAt == len(blocks) { + if insertAt == len(contentBlocks) { rawBlocks = append(rawBlocks, reminderBlocks...) } payload, _ = sjson.SetRawBytes(payload, contentPath, []byte("["+strings.Join(rawBlocks, ",")+"]")) } else if content.Type == gjson.String { - rawBlocks := make([]string, 0, len(reminderTexts)+1) - for _, reminderText := range reminderTexts { - rawBlocks = append(rawBlocks, buildTextBlock(reminderText, nil)) + rawBlocks := make([]string, 0, len(blocks)+1) + for _, b := range blocks { + reminderText := claudeCallerSystemReminder(b.Text) + if isExplicit && b.RawCC != "" { + rawBlocks = append(rawBlocks, `{"type":"text","text":`+marshalJSONStringWithoutHTMLEscape(reminderText)+`,"cache_control":`+b.RawCC+`}`) + } else { + rawBlocks = append(rawBlocks, buildTextBlock(reminderText, nil)) + } } rawBlocks = append(rawBlocks, buildTextBlock(content.String(), nil)) payload, _ = sjson.SetRawBytes(payload, contentPath, []byte("["+strings.Join(rawBlocks, ",")+"]")) @@ -785,8 +864,16 @@ func claudeMidConversationSystemMessagesAtEnd(payload []byte) bool { } func insertClaudeMidConversationSystemMessages(payload []byte, texts []string) []byte { + blocks := make([]forwardedClaudeSystemPromptBlock, 0, len(texts)) + for _, text := range texts { + blocks = append(blocks, forwardedClaudeSystemPromptBlock{Text: text, CacheControl: &claudeCodeCacheControl}) + } + return insertClaudeMidConversationSystemBlocks(payload, blocks, false) +} + +func insertClaudeMidConversationSystemBlocks(payload []byte, blocks []forwardedClaudeSystemPromptBlock, isExplicit bool) []byte { firstUserIdx := firstClaudeUserMessageIndex(payload) - if firstUserIdx < 0 || len(texts) == 0 { + if firstUserIdx < 0 || len(blocks) == 0 { return payload } @@ -799,11 +886,11 @@ func insertClaudeMidConversationSystemMessages(payload []byte, texts []string) [ for insertAt < len(messageBlocks) && messageBlocks[insertAt].Get("role").String() == "user" { insertAt++ } - if len(messageBlocks)-insertAt >= len(texts) { + if len(messageBlocks)-insertAt >= len(blocks) { matches := true - for idx, text := range texts { + for idx, block := range blocks { message := messageBlocks[insertAt+idx] - if message.Get("role").String() != "system" || claudeMessageContentText(message.Get("content")) != text { + if message.Get("role").String() != "system" || claudeMessageContentText(message.Get("content")) != block.Text { matches = false break } @@ -813,9 +900,9 @@ func insertClaudeMidConversationSystemMessages(payload []byte, texts []string) [ } } - systemMessages := make([]string, 0, len(texts)) - for _, text := range texts { - content := "[" + buildTextBlock(text, &claudeCodeCacheControl) + "]" + systemMessages := make([]string, 0, len(blocks)) + for _, block := range blocks { + content := "[" + buildForwardedSystemBlock(block.Text, block, isExplicit) + "]" systemMessages = append(systemMessages, `{"role":"system","content":`+content+"}") } rawMessages := make([]string, 0, len(messageBlocks)+len(systemMessages)) @@ -1202,6 +1289,10 @@ func isClaudeCodeCurrentDateReminder(text string) bool { // one whenever continuity is available) so the reminder text cannot change // between requests of one session. func injectClaudeCodeCurrentDate(payload []byte, date string) []byte { + return injectClaudeCodeCurrentDateInternal(payload, date, false) +} + +func injectClaudeCodeCurrentDateInternal(payload []byte, date string, isExplicit bool) []byte { firstUserIdx := firstClaudeUserMessageIndex(payload) if firstUserIdx < 0 { return payload @@ -1213,7 +1304,11 @@ func injectClaudeCodeCurrentDate(payload []byte, date string) []byte { dateBlock := buildTextBlock(dateText, nil) if content.Type == gjson.String { - userBlock := buildTextBlock(content.String(), &claudeCodeCacheControl) + var userCC *claudeCacheControl = &claudeCodeCacheControl + if isExplicit { + userCC = nil + } + userBlock := buildTextBlock(content.String(), userCC) newArray := "[" + dateBlock + "," + userBlock + "]" payload, _ = sjson.SetRawBytes(payload, contentPath, []byte(newArray)) return payload @@ -1232,9 +1327,11 @@ func injectClaudeCodeCurrentDate(payload []byte, date string) []byte { continue } if !actualTextCached && !isClaudeCodeContextReminder(text) { - rawBlocks = append(rawBlocks, withEphemeralCacheControl(block.Raw)) - actualTextCached = true - continue + if !isExplicit { + rawBlocks = append(rawBlocks, withEphemeralCacheControl(block.Raw)) + actualTextCached = true + continue + } } } rawBlocks = append(rawBlocks, block.Raw) @@ -1446,11 +1543,13 @@ func applyCloakingInternal( confirmedClaudeCode bool, cchSigning bool, obfuscateSensitiveWords bool, + explicitCacheMode ...bool, ) ([]byte, bool, error) { policy, settings := resolveClaudeWirePolicy(cfg, auth, apiKey, confirmedClaudeCode) if !policy.Cloak { return payload, false, nil } + isExplicit := (len(explicitCacheMode) > 0 && explicitCacheMode[0]) || isExplicitPromptCacheModeContext(ctx) // Strict mode drops caller system prompts entirely, so nothing needs a // destination and an unusable block cannot lose information. if !settings.strictMode { @@ -1493,6 +1592,10 @@ func applyCloakingInternal( if !isProbeOrHelper && !isSubagent { turnOrigin = "human" } + turnOriginArgs := []string{turnOrigin} + if isExplicit { + turnOriginArgs = append(turnOriginArgs, "explicit_cache_mode") + } payload = checkSystemInstructionsWithSigningModeAt( payload, settings.strictMode, @@ -1504,7 +1607,7 @@ func applyCloakingInternal( isSubagent, prevReq, promptID, - turnOrigin, + turnOriginArgs..., ) // In native Claude Code 2.1.258, claude-fable-5-1 requests carry: @@ -1525,7 +1628,7 @@ func applyCloakingInternal( // Probes never use 1h cache in native Claude Code; ensure any caller-supplied // 1h ttl is stripped to match extended-cache-ttl beta suppression. Subagents // preserve caller-requested 1h cache TTL (e.g. subagentPromptCacheTtl: 1h). - if isProbeOrHelper || (isSubagent && !helps.ClaudeSubagentRequests1h(incomingHeaders, payload)) { + if !isExplicit && (isProbeOrHelper || (isSubagent && !helps.ClaudeSubagentRequests1h(incomingHeaders, payload))) { payload = stripClaudeCacheControlTTL(payload) } @@ -1550,8 +1653,9 @@ func applyCloakingInternal( } type claudeCacheControl struct { - Type string `json:"type"` - TTL string `json:"ttl,omitempty"` + Type string `json:"type"` + TTL string `json:"ttl,omitempty"` + Scope string `json:"scope,omitempty"` } // claudeCodeCacheControl is the default Claude Code breakpoint shape. @@ -1723,7 +1827,50 @@ func forEachClaudeCacheControlBlock(payload []byte, visit func(path string, bloc } } -func shouldEnsureCacheControl(payload []byte, cloaked, confirmedClaudeCode bool) bool { +// isExplicitPromptCacheMode reports whether any of the candidate payloads +// specifies prompt_cache_options.mode as "explicit" (case-insensitive). +func isExplicitPromptCacheMode(payloads ...[]byte) bool { + for _, payload := range payloads { + if len(payload) == 0 { + continue + } + mode := gjson.GetBytes(payload, "prompt_cache_options.mode").String() + if strings.EqualFold(strings.TrimSpace(mode), "explicit") { + return true + } + } + return false +} + +type explicitPromptCacheModeContextKey struct{} + +func withExplicitPromptCacheMode(ctx context.Context, explicit bool) context.Context { + return context.WithValue(ctx, explicitPromptCacheModeContextKey{}, explicit) +} + +func isExplicitPromptCacheModeContext(ctx context.Context) bool { + if ctx == nil { + return false + } + v, _ := ctx.Value(explicitPromptCacheModeContextKey{}).(bool) + return v +} + +func stripPromptCacheOptions(payload []byte) []byte { + if !gjson.GetBytes(payload, "prompt_cache_options").Exists() { + return payload + } + updated, err := sjson.DeleteBytes(payload, "prompt_cache_options") + if err != nil { + return payload + } + return updated +} + +func shouldEnsureCacheControl(payload []byte, cloaked, confirmedClaudeCode bool, candidates ...[]byte) bool { + if isExplicitPromptCacheMode(append([][]byte{payload}, candidates...)...) { + return false + } return !confirmedClaudeCode && (cloaked || countCacheControls(payload) == 0) } diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go index 3ee64fa61..db3fc46ba 100644 --- a/internal/runtime/executor/claude_executor_execute.go +++ b/internal/runtime/executor/claude_executor_execute.go @@ -91,6 +91,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r // based on client type and configuration. _, wireSettings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, confirmedClaudeCode) isProbeOrHelper := helps.IsClaudeProbeOrHelperRequest(body) + explicitCacheMode := isExplicitPromptCacheMode(originalPayload, req.Payload, body) + if explicitCacheMode { + ctx = withExplicitPromptCacheMode(ctx, true) + } var cloaked bool body, cloaked, err = applyCloakingInternal( ctx, @@ -101,6 +105,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r confirmedClaudeCode, cchSigning, false, + explicitCacheMode, ) if err != nil { return resp, err @@ -146,7 +151,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r // first-user marker cannot suppress system/latest-user breakpoints. // cloaked and confirmedClaudeCode are mutually exclusive: resolveClaudeWirePolicy // forces Cloak off for a confirmed native client. - cpaOwnsCacheControl := shouldEnsureCacheControl(body, cloaked, confirmedClaudeCode) + cpaOwnsCacheControl := shouldEnsureCacheControl(body, cloaked, confirmedClaudeCode, originalPayload, req.Payload) if cpaOwnsCacheControl { body = ensureCacheControl(body) } @@ -155,6 +160,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r // Cloaking and ensureCacheControl may push the total over 4 when the client // already sends multiple cache_control blocks. body = enforceCacheControlLimit(body, 4) + body = stripPromptCacheOptions(body) // Native selects the 1h cache pool only for OAuth credentials and pairs it with // extended-cache-ttl-2025-04-11, which claudeCodeCLIBetas emits on exactly the @@ -173,13 +179,15 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r subagent1h := isSubagent && helps.ClaudeSubagentRequests1h(incomingHeaders, body) if cpaOwnsCacheControl && fp.ProfileClaudeCodeCLI && (!isSubagent || subagent1h) && !isProbeOrHelper { body = upgradeClaudeCacheControlTTL(body, claudeCacheControlTTL1h) - } else if isProbeOrHelper || (isSubagent && !subagent1h) { + } else if (isProbeOrHelper || (isSubagent && !subagent1h)) && !explicitCacheMode { body = stripClaudeCacheControlTTL(body) } // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. // A 1h-TTL block must not appear after a 5m-TTL block in evaluation order (tools→system→messages). - body = normalizeCacheControlTTL(body) + if !explicitCacheMode { + body = normalizeCacheControlTTL(body) + } // Set the built-in transport default before the explicit payload barrier. // Native non-stream Haiku helper requests omit stream rather than sending // false, so preserve that measured wire shape when the transport agrees. @@ -233,6 +241,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r diagnosticsState = claudeDiagnosticsRequestState{} } extraBetas, bodyForUpstream = extractAndRemoveBetas(bodyForUpstream) + bodyForUpstream = stripPromptCacheOptions(bodyForUpstream) if cchSigning { bodyForUpstream, err = signAnthropicMessagesBody(bodyForUpstream) if err != nil { diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index f8b220390..aa22c4f08 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -94,6 +94,10 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A // based on client type and configuration. _, wireSettings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, confirmedClaudeCode) isProbeOrHelper := helps.IsClaudeProbeOrHelperRequest(body) + explicitCacheMode := isExplicitPromptCacheMode(originalPayload, req.Payload, body) + if explicitCacheMode { + ctx = withExplicitPromptCacheMode(ctx, true) + } var cloaked bool body, cloaked, err = applyCloakingInternal( ctx, @@ -104,6 +108,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A confirmedClaudeCode, cchSigning, false, + explicitCacheMode, ) if err != nil { return nil, err @@ -149,7 +154,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A // first-user marker cannot suppress system/latest-user breakpoints. // cloaked and confirmedClaudeCode are mutually exclusive: resolveClaudeWirePolicy // forces Cloak off for a confirmed native client. - cpaOwnsCacheControl := shouldEnsureCacheControl(body, cloaked, confirmedClaudeCode) + cpaOwnsCacheControl := shouldEnsureCacheControl(body, cloaked, confirmedClaudeCode, originalPayload, req.Payload) if cpaOwnsCacheControl { body = ensureCacheControl(body) } @@ -158,6 +163,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A // Cloaking and ensureCacheControl may push the total over 4 when the client // already sends multiple cache_control blocks. body = enforceCacheControlLimit(body, 4) + body = stripPromptCacheOptions(body) // Native selects the 1h cache pool only for OAuth credentials and pairs it with // extended-cache-ttl-2025-04-11, which claudeCodeCLIBetas emits on exactly the @@ -176,12 +182,14 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A subagent1h := isSubagent && helps.ClaudeSubagentRequests1h(incomingHeaders, body) if cpaOwnsCacheControl && fp.ProfileClaudeCodeCLI && (!isSubagent || subagent1h) && !isProbeOrHelper { body = upgradeClaudeCacheControlTTL(body, claudeCacheControlTTL1h) - } else if isProbeOrHelper || (isSubagent && !subagent1h) { + } else if (isProbeOrHelper || (isSubagent && !subagent1h)) && !explicitCacheMode { body = stripClaudeCacheControlTTL(body) } // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. - body = normalizeCacheControlTTL(body) + if !explicitCacheMode { + body = normalizeCacheControlTTL(body) + } // Extract betas from body and convert to header var extraBetas []string @@ -228,6 +236,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A diagnosticsState = claudeDiagnosticsRequestState{} } extraBetas, bodyForUpstream = extractAndRemoveBetas(bodyForUpstream) + bodyForUpstream = stripPromptCacheOptions(bodyForUpstream) if cchSigning { bodyForUpstream, err = signAnthropicMessagesBody(bodyForUpstream) if err != nil { diff --git a/internal/runtime/executor/claude_executor_tokens.go b/internal/runtime/executor/claude_executor_tokens.go index 9d2315660..6e3ecdae6 100644 --- a/internal/runtime/executor/claude_executor_tokens.go +++ b/internal/runtime/executor/claude_executor_tokens.go @@ -53,6 +53,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut } body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel, helps.APIKeyModelIsCompat(req)) body = helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslatedForPayload, req, opts)(body) + body = stripPromptCacheOptions(body) if errValidate := validateClaudeTokenCountRequest(body); errValidate != nil { return cliproxyexecutor.Response{}, errValidate } @@ -171,13 +172,14 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy // Messages path. Kimi opt-in uses the same contract. policy, settings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, confirmedClaudeCode) cloaked := policy.Cloak + explicitCacheMode := isExplicitPromptCacheMode(originalPayload, req.Payload, body) if cloaked { if !settings.strictMode { if errSystem := validateClaudeCallerSystemBlocks(gjson.GetBytes(body, "system")); errSystem != nil { return cliproxyexecutor.Response{}, errSystem } } - body = relocateClaudeSystemPromptForCountTokens(body, settings.strictMode) + body = relocateClaudeSystemPromptForCountTokens(body, settings.strictMode, explicitCacheMode) if len(settings.sensitiveWords) > 0 { body = helps.ObfuscateSensitiveWords(body, helps.BuildSensitiveWordMatcher(settings.sensitiveWords)) } @@ -185,7 +187,9 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy // Keep count_tokens requests compatible with Anthropic cache-control constraints too. body = enforceCacheControlLimit(body, 4) - body = normalizeCacheControlTTL(body) + if !explicitCacheMode { + body = normalizeCacheControlTTL(body) + } // Extract betas from body and convert to header (for count_tokens too) var extraBetas []string @@ -197,6 +201,7 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy body, _ = prepareClaudeOAuthToolNamesForUpstream(body, mcpAliases) } body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel, helps.APIKeyModelIsCompat(req)) + body = stripPromptCacheOptions(body) // Two different reasons converge on the same deletions, and they must stay // separable. // @@ -227,6 +232,7 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy return cliproxyexecutor.Response{}, errMidSystem } body = helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslatedForPayload, req, opts)(body) + body = stripPromptCacheOptions(body) httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) if err != nil { return cliproxyexecutor.Response{}, err diff --git a/internal/runtime/executor/claude_prompt_cache_options_test.go b/internal/runtime/executor/claude_prompt_cache_options_test.go new file mode 100644 index 000000000..2d92d389f --- /dev/null +++ b/internal/runtime/executor/claude_prompt_cache_options_test.go @@ -0,0 +1,950 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" + "github.com/tidwall/gjson" +) + +func writeClaudeSSEMockResponse(w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-sonnet\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n") + _, _ = fmt.Fprint(w, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n") + _, _ = fmt.Fprint(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n") + _, _ = fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_SuppressesAutoCacheControl(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 0 { + t.Fatalf("expected 0 cache_control blocks with mode: explicit and no client breakpoints, got %d: %s", count, string(seenBody)) + } + if gjson.GetBytes(seenBody, "prompt_cache_options").Exists() { + t.Fatalf("prompt_cache_options should not be forwarded upstream: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_PreservesClientBreakpoints(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": [ + {"type": "text", "text": "Cached part", "cache_control": {"type": "ephemeral"}} + ]}, + {"role": "user", "content": "Uncached part"} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control block (the client-specified one), got %d: %s", count, string(seenBody)) + } + if !gjson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("client cache_control was not preserved: %s", string(seenBody)) + } + if gjson.GetBytes(seenBody, "system.0.cache_control").Exists() { + t.Fatalf("system prompt unexpectedly received automatic cache_control: %s", string(seenBody)) + } + if gjson.GetBytes(seenBody, "messages.1.content.0.cache_control").Exists() { + t.Fatalf("second message unexpectedly received automatic cache_control: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Stream_Explicit_SuppressesAutoCacheControl(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "stream": true, + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "user", "content": "Stream query"} + ] + }`) + + streamRes, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + for range streamRes.Chunks { + } + + if count := countCacheControls(seenBody); count != 0 { + t.Fatalf("expected 0 cache_control blocks in stream with mode: explicit, got %d: %s", count, string(seenBody)) + } + if gjson.GetBytes(seenBody, "prompt_cache_options").Exists() { + t.Fatalf("prompt_cache_options should not be forwarded upstream: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_ImplicitOrMissing_PreservesAutoCacheControl(t *testing.T) { + for _, tc := range []struct { + name string + originalReq []byte + }{ + { + name: "missing prompt_cache_options", + originalReq: []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"} + ] + }`), + }, + { + name: "mode implicit", + originalReq: []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "implicit" + }, + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"} + ] + }`), + }, + } { + t.Run(tc.name, func(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: tc.originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: tc.originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + // In normal mode without client markers, system and message should receive auto-injected cache_control. + if !gjson.GetBytes(seenBody, "system.0.cache_control").Exists() { + t.Fatalf("expected system cache_control in implicit mode: %s", string(seenBody)) + } + if !gjson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("expected message cache_control in implicit mode: %s", string(seenBody)) + } + }) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_ToolsFallbackSuppressed(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "tools": [ + {"type": "function", "function": {"name": "tool1", "description": "t1", "parameters": {"type": "object"}}}, + {"type": "function", "function": {"name": "tool2", "description": "t2", "parameters": {"type": "object"}}} + ], + "messages": [ + {"role": "user", "content": "Hello without system"} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 0 { + t.Fatalf("expected 0 cache_control blocks with mode: explicit when only tools are present, got %d: %s", count, string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_ClientToolBreakpointPreserved(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "tools": [ + {"type": "function", "cache_control": {"type": "ephemeral"}, "function": {"name": "tool1", "description": "t1", "parameters": {"type": "object"}}}, + {"type": "function", "function": {"name": "tool2", "description": "t2", "parameters": {"type": "object"}}} + ], + "messages": [ + {"role": "user", "content": "Hello"} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected 1 cache_control block on client-specified tool, got %d: %s", count, string(seenBody)) + } + if !gjson.GetBytes(seenBody, "tools.0.cache_control").Exists() { + t.Fatalf("tools.0 should have client-specified cache_control: %s", string(seenBody)) + } + if gjson.GetBytes(seenBody, "tools.1.cache_control").Exists() { + t.Fatalf("tools.1 should NOT have cache_control: %s", string(seenBody)) + } + if gjson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("message should NOT receive automatic cache_control: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_Cloaked_SuppressesAllAutoCacheControl(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello in cloaked mode"} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 0 { + t.Fatalf("expected 0 cache_control blocks in cloaked mode with explicit mode and no client breakpoints, got %d: %s", count, string(seenBody)) + } + if gjson.GetBytes(seenBody, "prompt_cache_options").Exists() { + t.Fatalf("prompt_cache_options should not be forwarded upstream: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_Cloaked_PreservesClientBreakpoints(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "First turn"}, + {"role": "assistant", "content": "Understood."}, + {"role": "user", "content": [ + {"type": "text", "text": "Specific cached user query", "cache_control": {"type": "ephemeral"}} + ]} + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control block (the client-specified one), got %d: %s", count, string(seenBody)) + } + if gjson.GetBytes(seenBody, "system.0.cache_control").Exists() || gjson.GetBytes(seenBody, "system.1.cache_control").Exists() { + t.Fatalf("system blocks should not have automatic cache_control: %s", string(seenBody)) + } + if !gjson.GetBytes(seenBody, "messages.3.content.0.cache_control").Exists() { + t.Fatalf("client specified cache_control on user turn should be preserved: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_Cloaked_PreservesClientSystemBreakpoint(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + { + "role": "system", + "content": "Special client system prompt with breakpoint", + "cache_control": {"type": "ephemeral", "scope": "global"} + }, + { + "role": "user", + "content": "Regular user query" + } + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control block (the client system one), got %d: %s", count, string(seenBody)) + } + // Verify CPA synthetic agent block does NOT have cache_control + if gjson.GetBytes(seenBody, "system.1.cache_control").Exists() { + t.Fatalf("CPA synthetic agent block must not have cache_control: %s", string(seenBody)) + } + // Verify user message does NOT have cache_control + if gjson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("first user message must not receive auto cache_control: %s", string(seenBody)) + } + // Verify client system prompt has cache_control preserved + foundClientSystemCC := false + messages := gjson.GetBytes(seenBody, "messages") + if messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + if msg.Get("role").String() == "system" { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("text").String() == "Special client system prompt with breakpoint" && part.Get("cache_control.type").String() == "ephemeral" && part.Get("cache_control.scope").String() == "global" { + foundClientSystemCC = true + } + return true + }) + } + } + return true + }) + } + system := gjson.GetBytes(seenBody, "system") + if system.IsArray() { + system.ForEach(func(_, part gjson.Result) bool { + if part.Get("text").String() == "Special client system prompt with breakpoint" && part.Get("cache_control.type").String() == "ephemeral" && part.Get("cache_control.scope").String() == "global" { + foundClientSystemCC = true + } + return true + }) + } + if !foundClientSystemCC { + t.Fatalf("client specified cache_control on system prompt must be preserved: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_CountTokens_StripsPromptCacheOptions(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"input_tokens": 15}`)) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + {"role": "user", "content": "How many tokens?"} + ] + }`) + + resp, err := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("CountTokens() error = %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("expected non-empty CountTokens payload") + } + + // For custom base_url, CountTokens uses the local estimator after stripping fields. + // Now test countTokensUpstream specifically to verify upstream payload body: + upstreamResp, errUpstream := executor.countTokensUpstream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if errUpstream != nil { + t.Fatalf("countTokensUpstream() error = %v", errUpstream) + } + if len(upstreamResp.Payload) == 0 { + t.Fatal("expected non-empty countTokensUpstream payload") + } + if gjson.GetBytes(seenBody, "prompt_cache_options").Exists() { + t.Fatalf("prompt_cache_options should not be forwarded in count_tokens upstream: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_CountTokensUpstream_StripsPromptCacheOptions_EvenWithPayloadRule(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"input_tokens": 15}`)) + })) + defer server.Close() + + cfg := &config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "claude-3-5-sonnet-20241022"}}, + Params: map[string]any{ + "prompt_cache_options.mode": "explicit", + }, + }, + }, + }, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "How many tokens?"} + ] + }`) + + upstreamResp, errUpstream := executor.countTokensUpstream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if errUpstream != nil { + t.Fatalf("countTokensUpstream() error = %v", errUpstream) + } + if len(upstreamResp.Payload) == 0 { + t.Fatal("expected non-empty countTokensUpstream payload") + } + if gjson.GetBytes(seenBody, "prompt_cache_options").Exists() { + t.Fatalf("prompt_cache_options should not be forwarded in count_tokens upstream even if injected by payload rule: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_ProbePreservesClient1hTTL(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + // "Hi" is a classic probe text that would trigger isProbeOrHelper in native cloaking. + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ] + } + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if gotTTL := gjson.GetBytes(seenBody, "messages.0.content.0.cache_control.ttl").String(); gotTTL != "1h" { + t.Fatalf("client specified 1h TTL must be preserved in explicit mode even on probe request, got %q: %s", gotTTL, string(seenBody)) + } +} + +func TestClaudeExecutor_CountTokensUpstream_Cloaked_ExplicitPreservesClientSystemBreakpoint(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"input_tokens": 20}`)) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + { + "role": "system", + "content": "Count tokens system instructions", + "cache_control": {"type": "ephemeral", "scope": "global"} + }, + { + "role": "user", + "content": "Count tokens user message" + } + ] + }`) + + upstreamResp, errUpstream := executor.countTokensUpstream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if errUpstream != nil { + t.Fatalf("countTokensUpstream() error = %v", errUpstream) + } + if len(upstreamResp.Payload) == 0 { + t.Fatal("expected non-empty countTokensUpstream payload") + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control in count_tokens explicit mode, got %d: %s", count, string(seenBody)) + } + // Verify user message did not get auto-cached + if gjson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("user message must not receive auto cache_control: %s", string(seenBody)) + } + // Verify client system breakpoint was preserved with scope: global + foundClientSystemCC := false + messages := gjson.GetBytes(seenBody, "messages") + if messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + if msg.Get("role").String() == "system" { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("text").String() == "Count tokens system instructions" && + part.Get("cache_control.type").String() == "ephemeral" && + part.Get("cache_control.scope").String() == "global" { + foundClientSystemCC = true + } + return true + }) + } + } + return true + }) + } + if !foundClientSystemCC { + t.Fatalf("client specified cache_control with scope: global must be preserved in count_tokens: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_PromptCacheOptionsMode_Explicit_Cloaked_LegacyModelPreservesClientSystemBreakpoint(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + writeClaudeSSEMockResponse(w) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + // claude-3-5-haiku-20241022 is in claudeLegacySystemReminderModels, so it uses legacy system reminders + originalReq := []byte(`{ + "model": "claude-3-5-haiku-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + { + "role": "system", + "content": "Legacy system instructions with scope", + "cache_control": {"type": "ephemeral", "scope": "global"} + }, + { + "role": "user", + "content": "User prompt" + } + ] + }`) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-haiku-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control block in legacy model explicit mode, got %d: %s", count, string(seenBody)) + } + // Verify system.1 agent block does not have cache_control + if gjson.GetBytes(seenBody, "system.1.cache_control").Exists() { + t.Fatalf("agent block must not have cache_control: %s", string(seenBody)) + } + // Verify the prepended reminder on user message preserves scope: global + foundReminderCC := false + messages := gjson.GetBytes(seenBody, "messages") + if messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, part gjson.Result) bool { + if strings.Contains(part.Get("text").String(), "Legacy system instructions with scope") && + part.Get("cache_control.type").String() == "ephemeral" && + part.Get("cache_control.scope").String() == "global" { + foundReminderCC = true + } + return true + }) + } + return true + }) + } + if !foundReminderCC { + t.Fatalf("legacy system reminder must preserve client cache_control with scope: global: %s", string(seenBody)) + } +} + +func TestClaudeExecutor_CountTokensUpstream_Cloaked_LegacyModel_ExplicitPreservesClientSystemBreakpoint(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"input_tokens": 20}`)) + })) + defer server.Close() + + cfg := &config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + Cloak: &config.CloakConfig{Mode: "always"}, + }}, + } + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + originalReq := []byte(`{ + "model": "claude-3-5-haiku-20241022", + "prompt_cache_options": { + "mode": "explicit" + }, + "messages": [ + { + "role": "system", + "content": "Count tokens legacy system instructions", + "cache_control": {"type": "ephemeral", "scope": "global"} + }, + { + "role": "user", + "content": "User prompt" + } + ] + }`) + + upstreamResp, errUpstream := executor.countTokensUpstream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "claude-3-5-haiku-20241022", + Payload: originalReq, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai"), + OriginalRequest: originalReq, + }) + if errUpstream != nil { + t.Fatalf("countTokensUpstream() error = %v", errUpstream) + } + if len(upstreamResp.Payload) == 0 { + t.Fatal("expected non-empty countTokensUpstream payload") + } + + if count := countCacheControls(seenBody); count != 1 { + t.Fatalf("expected exactly 1 cache_control in legacy count_tokens explicit mode, got %d: %s", count, string(seenBody)) + } + foundReminderCC := false + messages := gjson.GetBytes(seenBody, "messages") + if messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, part gjson.Result) bool { + if strings.Contains(part.Get("text").String(), "Count tokens legacy system instructions") && + part.Get("cache_control.type").String() == "ephemeral" && + part.Get("cache_control.scope").String() == "global" { + foundReminderCC = true + } + return true + }) + } + return true + }) + } + if !foundReminderCC { + t.Fatalf("legacy count_tokens system reminder must preserve client cache_control with scope: global: %s", string(seenBody)) + } +}