fix: preserve system cache breakpoint in OAuth mimic

This commit is contained in:
jy.liu
2026-07-27 19:28:45 +08:00
parent 59ce11c780
commit 1631b19f84
3 changed files with 86 additions and 27 deletions
@@ -798,11 +798,12 @@ func TestGatewayService_AnthropicOAuthMimic_RewritesSystemWithBillingBlock(t *te
gin.SetMode(gin.TestMode)
tests := []struct {
name string
body string
wantModel string
wantOriginalSystem string
wantMetadataUserID string
name string
body string
wantModel string
wantOriginalSystem string
wantOriginalSystemCacheTTL string
wantMetadataUserID string
}{
{
name: "sonnet system array",
@@ -817,11 +818,12 @@ func TestGatewayService_AnthropicOAuthMimic_RewritesSystemWithBillingBlock(t *te
wantOriginalSystem: "x-anthropic-billing-header keep",
},
{
name: "haiku full mimicry",
body: `{"model":"claude-haiku-4-5","metadata":{"user_id":"pi-session-metadata"},"system":[{"type":"text","text":"Pi project instructions","cache_control":{"type":"ephemeral"}}],"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`,
wantModel: "claude-haiku-4-5-20251001",
wantOriginalSystem: "Pi project instructions",
wantMetadataUserID: "pi-session-metadata",
name: "haiku full mimicry",
body: `{"model":"claude-haiku-4-5","metadata":{"user_id":"pi-session-metadata"},"system":[{"type":"text","text":"Pi project instructions","cache_control":{"type":"ephemeral","ttl":"1h"}}],"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`,
wantModel: "claude-haiku-4-5-20251001",
wantOriginalSystem: "Pi project instructions",
wantOriginalSystemCacheTTL: "1h",
wantMetadataUserID: "pi-session-metadata",
},
}
@@ -912,6 +914,12 @@ func TestGatewayService_AnthropicOAuthMimic_RewritesSystemWithBillingBlock(t *te
firstMsg := messages.Array()[0]
require.Equal(t, "user", firstMsg.Get("role").String())
require.Contains(t, firstMsg.Get("content.0.text").String(), tt.wantOriginalSystem)
if tt.wantOriginalSystemCacheTTL != "" {
require.Equal(t, "ephemeral", firstMsg.Get("content.0.cache_control.type").String())
require.Equal(t, tt.wantOriginalSystemCacheTTL, firstMsg.Get("content.0.cache_control.ttl").String())
} else {
require.False(t, firstMsg.Get("content.0.cache_control").Exists())
}
if tt.wantMetadataUserID != "" {
require.Equal(t, tt.wantMetadataUserID, gjson.GetBytes(upstream.lastBody, "metadata.user_id").String())
@@ -854,26 +854,42 @@ func ValidateClaudeOAuthSystemPromptBlocksConfig(raw string) error {
return nil
}
func extractSystemTextAndCacheControl(system any) (string, any) {
switch v := system.(type) {
case string:
return strings.TrimSpace(v), nil
case []any:
var parts []string
var cacheControl any
for _, item := range v {
m, ok := item.(map[string]any)
if !ok {
continue
}
text, ok := m["text"].(string)
if !ok || strings.TrimSpace(text) == "" {
continue
}
parts = append(parts, text)
// system blocks are collapsed into one messages text block below.
// Preserve the last original breakpoint as the closest equivalent
// boundary, including its client-selected TTL.
if cc, exists := m["cache_control"]; exists && cc != nil {
cacheControl = cc
}
}
return strings.Join(parts, "\n\n"), cacheControl
default:
return "", nil
}
}
func rewriteSystemForNonClaudeCodeWithPromptBlocks(body []byte, system any, expansionPrompt string, blocksConfig string) []byte {
system = normalizeSystemParam(system)
expansionPrompt = defaultClaudeOAuthExpansionPrompt(expansionPrompt)
// 1. 提取原始 system prompt 文本
var originalSystemText string
switch v := system.(type) {
case string:
originalSystemText = strings.TrimSpace(v)
case []any:
var parts []string
for _, item := range v {
if m, ok := item.(map[string]any); ok {
if text, ok := m["text"].(string); ok && strings.TrimSpace(text) != "" {
parts = append(parts, text)
}
}
}
originalSystemText = strings.Join(parts, "\n\n")
}
// 1. 提取原始 system prompt 文本及其缓存断点
originalSystemText, originalSystemCacheControl := extractSystemTextAndCacheControl(system)
// 2. 构造 system 数组,对齐真实 Claude Code CLI 的 3-block 形态:
// [0] billing attribution block(cc_version={cliVer}.{fp}; cc_entrypoint=cli;)
@@ -906,10 +922,17 @@ func rewriteSystemForNonClaudeCodeWithPromptBlocks(body []byte, system any, expa
// 模型仍通过 messages 接收完整指令,保留客户端功能
ccPromptTrimmed := strings.TrimSpace(claudeCodeSystemPrompt)
if originalSystemText != "" && originalSystemText != ccPromptTrimmed && !hasClaudeCodePrefix(originalSystemText) {
instructionBlock := map[string]any{
"type": "text",
"text": "[System Instructions]\n" + originalSystemText,
}
if originalSystemCacheControl != nil {
instructionBlock["cache_control"] = originalSystemCacheControl
}
instrMsg, err1 := json.Marshal(map[string]any{
"role": "user",
"content": []map[string]any{
{"type": "text", "text": "[System Instructions]\n" + originalSystemText},
instructionBlock,
},
})
ackMsg, err2 := json.Marshal(map[string]any{
@@ -482,6 +482,34 @@ func TestRewriteSystemForNonClaudeCodeWithPrompt_UsesCustomExpansionPrompt(t *te
require.Equal(t, "ephemeral", system.Array()[2].Get("cache_control.type").String())
}
func TestRewriteSystemForNonClaudeCode_PreservesSystemCacheControlOnMigratedMessage(t *testing.T) {
body := []byte(`{"model":"claude-3","system":[{"type":"text","text":"Stable project instructions","cache_control":{"type":"ephemeral","ttl":"1h"}}],"messages":[{"role":"user","content":"hello"}]}`)
system := []any{
map[string]any{
"type": "text",
"text": "Stable project instructions",
"cache_control": map[string]any{"type": "ephemeral", "ttl": "1h"},
},
}
result := rewriteSystemForNonClaudeCode(body, system)
require.Equal(t, "[System Instructions]\nStable project instructions", gjson.GetBytes(result, "messages.0.content.0.text").String())
require.Equal(t, "ephemeral", gjson.GetBytes(result, "messages.0.content.0.cache_control.type").String())
require.Equal(t, "1h", gjson.GetBytes(result, "messages.0.content.0.cache_control.ttl").String())
}
func TestRewriteSystemForNonClaudeCode_LeavesMigratedMessageUncachedWithoutSystemBreakpoint(t *testing.T) {
body := []byte(`{"model":"claude-3","system":[{"type":"text","text":"Project instructions"}],"messages":[{"role":"user","content":"hello"}]}`)
system := []any{
map[string]any{"type": "text", "text": "Project instructions"},
}
result := rewriteSystemForNonClaudeCode(body, system)
require.False(t, gjson.GetBytes(result, "messages.0.content.0.cache_control").Exists())
}
func TestRewriteSystemForNonClaudeCodeWithPromptBlocks_UsesConfiguredBlocks(t *testing.T) {
body := []byte(`{"model":"claude-3","system":"Project instructions","messages":[{"role":"user","content":"hello"}]}`)
blocks := `{