619e5ae619
Keep Anthropic request body rewrites attempt-local and synchronize the accepted wire body only after upstream success so failover and retry paths do not reuse stale parsed state.
668 lines
34 KiB
Go
668 lines
34 KiB
Go
//go:build unit
|
||
|
||
package service
|
||
|
||
import (
|
||
"context"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"regexp"
|
||
"strings"
|
||
"testing"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/stretchr/testify/require"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
// ============================================================================
|
||
// 背景
|
||
// ============================================================================
|
||
//
|
||
// Anthropic 上游对 body.context_management 字段实施 Pydantic schema 校验：
|
||
// 当且仅当 anthropic-beta header 含 context-management-2025-06-27 时接受。
|
||
// 否则报：
|
||
// "context_management: Extra inputs are not permitted"
|
||
//
|
||
// 本仓采用能力维度对称约束（与 Bedrock 路径的 sanitizeBedrockFieldsForBetaTokens
|
||
// 对称）：在所有 Anthropic 直连出口，按最终 anthropic-beta header 是否含上述 token
|
||
// 决定 body 是否保留同名字段。
|
||
//
|
||
// 本文件覆盖：
|
||
// 1) sanitizeAnthropicBodyForBetaTokens 纯函数
|
||
// 2) anthropicBetaTokensContains 解析辅助函数
|
||
// 3) computeFinalAnthropicBeta / computeFinalCountTokensAnthropicBeta 各路径
|
||
// 4) normalizeClaudeOAuthRequestBody 的 context_management 补齐行为（不再按 model 短路）
|
||
|
||
// ============================================================================
|
||
// anthropicBetaTokensContains
|
||
// ============================================================================
|
||
|
||
func TestAnthropicBetaTokensContains_EmptyInputs(t *testing.T) {
|
||
require.False(t, anthropicBetaTokensContains("", "context-management-2025-06-27"))
|
||
require.False(t, anthropicBetaTokensContains("oauth-2025-04-20", ""))
|
||
}
|
||
|
||
func TestAnthropicBetaTokensContains_SingleToken(t *testing.T) {
|
||
require.True(t, anthropicBetaTokensContains("context-management-2025-06-27", "context-management-2025-06-27"))
|
||
}
|
||
|
||
func TestAnthropicBetaTokensContains_MultiTokenComma(t *testing.T) {
|
||
header := "oauth-2025-04-20,context-management-2025-06-27,interleaved-thinking-2025-05-14"
|
||
require.True(t, anthropicBetaTokensContains(header, "context-management-2025-06-27"))
|
||
require.True(t, anthropicBetaTokensContains(header, "oauth-2025-04-20"))
|
||
require.False(t, anthropicBetaTokensContains(header, "fast-mode-2026-02-01"))
|
||
}
|
||
|
||
func TestAnthropicBetaTokensContains_ToleratesWhitespace(t *testing.T) {
|
||
header := "oauth-2025-04-20 , context-management-2025-06-27 , interleaved-thinking-2025-05-14"
|
||
require.True(t, anthropicBetaTokensContains(header, "context-management-2025-06-27"))
|
||
}
|
||
|
||
func TestAnthropicBetaTokensContains_SubstringNotMatched(t *testing.T) {
|
||
// 严格 token 比较，不应被子串误匹配
|
||
require.False(t, anthropicBetaTokensContains("context-management-2025-06-27-rev2", "context-management-2025-06-27"),
|
||
"必须按 token 边界匹配，不允许 prefix 子串误命中")
|
||
}
|
||
|
||
// ============================================================================
|
||
// sanitizeAnthropicBodyForBetaTokens
|
||
// ============================================================================
|
||
|
||
func TestSanitizeAnthropicBodyForBetaTokens_NoFieldNoChange(t *testing.T) {
|
||
body := []byte(`{"model":"claude-haiku-4-5","messages":[]}`)
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens(body, "oauth-2025-04-20")
|
||
require.False(t, changed)
|
||
require.Equal(t, string(body), string(out))
|
||
}
|
||
|
||
func TestSanitizeAnthropicBodyForBetaTokens_FieldKeptWhenBetaPresent(t *testing.T) {
|
||
body := []byte(`{"model":"claude-opus-4-7","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens(body,
|
||
"oauth-2025-04-20,context-management-2025-06-27,interleaved-thinking-2025-05-14")
|
||
require.False(t, changed)
|
||
require.True(t, gjson.GetBytes(out, "context_management").Exists())
|
||
require.Equal(t, "clear_thinking_20251015",
|
||
gjson.GetBytes(out, "context_management.edits.0.type").String())
|
||
}
|
||
|
||
func TestSanitizeAnthropicBodyForBetaTokens_FieldStrippedWhenBetaMissing(t *testing.T) {
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens(body, "oauth-2025-04-20,interleaved-thinking-2025-05-14")
|
||
require.True(t, changed)
|
||
require.False(t, gjson.GetBytes(out, "context_management").Exists(),
|
||
"header 不含 context-management beta 时必须 strip 同名字段")
|
||
}
|
||
|
||
func TestSanitizeAnthropicBodyForBetaTokens_FieldStrippedWhenBetaEmpty(t *testing.T) {
|
||
body := []byte(`{"context_management":{"edits":[]},"messages":[]}`)
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens(body, "")
|
||
require.True(t, changed)
|
||
require.False(t, gjson.GetBytes(out, "context_management").Exists())
|
||
}
|
||
|
||
func TestSanitizeAnthropicBodyForBetaTokens_EmptyBody(t *testing.T) {
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens([]byte{}, "")
|
||
require.False(t, changed)
|
||
require.Empty(t, out)
|
||
|
||
out, changed = sanitizeAnthropicBodyForBetaTokens(nil, "")
|
||
require.False(t, changed)
|
||
require.Empty(t, out)
|
||
}
|
||
|
||
// ★ 关键回归断言：能力维度 sanitize 解决了 "真 CC + haiku" 路径的过度删除问题。
|
||
// 真实 Claude Code CLI 2.1.87+ 客户端 header 含 context-management beta；
|
||
// 即使 model 是 haiku，sanitize 也不应剥离功能字段。
|
||
func TestSanitizeAnthropicBodyForBetaTokens_HaikuRealCCClientPreservesField(t *testing.T) {
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[]}`)
|
||
// 真 Claude Code CLI 2.1.87+ 客户端 header 含 context-management beta
|
||
clientBeta := "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27"
|
||
out, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta)
|
||
require.False(t, changed,
|
||
"真 CC 客户端 header 含 context-management beta 时，haiku body 字段必须保留（功能不丢）")
|
||
require.True(t, gjson.GetBytes(out, "context_management").Exists())
|
||
}
|
||
|
||
// ============================================================================
|
||
// computeFinalAnthropicBeta — 关键路径
|
||
// ============================================================================
|
||
|
||
func newTestGatewayServiceForBeta(injectBetaForAPIKey bool) *GatewayService {
|
||
cfg := &config.Config{}
|
||
cfg.Gateway.InjectBetaForAPIKey = injectBetaForAPIKey
|
||
return &GatewayService{cfg: cfg}
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_OAuthMimic_NonHaiku_IncludesContextManagement(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", true, "claude-sonnet-4-6", http.Header{}, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"OAuth mimic non-haiku 必须注入完整 CC mimicry beta，含 context-management-2025-06-27")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaOAuth))
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaClaudeCode))
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_OAuthMimic_Haiku_ExcludesContextManagement(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", true, "claude-haiku-4-5", http.Header{}, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.False(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"OAuth mimic haiku 仅注入 oauth + interleaved-thinking，不含 context-management")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaOAuth))
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaInterleavedThinking))
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_OAuthMimic_IgnoresClientBeta(t *testing.T) {
|
||
// mimic 路径下原代码白名单透传被跳过，client beta 应被忽略
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "custom-experimental-beta")
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", true, "claude-sonnet-4-6", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.False(t, strings.Contains(final, "custom-experimental-beta"),
|
||
"mimic 路径必须忽略客户端 anthropic-beta header")
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_OAuthTransparent_NonHaiku_PreservesClientContextManagement(t *testing.T) {
|
||
// 真 CC 客户端透传：客户端 header 中的 context-management beta 必须保留
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "claude-code-20250219,oauth-2025-04-20,context-management-2025-06-27")
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", false, "claude-sonnet-4-6", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement))
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_OAuthTransparent_Haiku_RealCCPreservesContextManagement(t *testing.T) {
|
||
// haiku 透传 + 客户端带 context-management beta → 必须保留
|
||
// （能力维度核心场景：避免 model-name 误删客户端透传的功能 beta）
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "claude-code-20250219,oauth-2025-04-20,context-management-2025-06-27,interleaved-thinking-2025-05-14")
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", false, "claude-haiku-4-5", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"真 CC + haiku + 客户端带 context-management beta → 透传必须保留")
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_APIKey_PassesClientBetaThroughDropSet(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "oauth-2025-04-20,custom-beta")
|
||
final, ok := s.computeFinalAnthropicBeta("apikey", false, "claude-sonnet-4-6", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, "oauth-2025-04-20"))
|
||
require.True(t, anthropicBetaTokensContains(final, "custom-beta"))
|
||
}
|
||
|
||
func TestComputeFinalAnthropicBeta_APIKey_NoClientBetaInjectOff_ShouldNotSet(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
final, ok := s.computeFinalAnthropicBeta("apikey", false, "claude-sonnet-4-6", http.Header{}, []byte(`{}`), nil)
|
||
require.False(t, ok, "API-key + 客户端未传 + InjectBetaForAPIKey 关 → 不应主动设置 anthropic-beta")
|
||
require.Equal(t, "", final)
|
||
}
|
||
|
||
// ============================================================================
|
||
// computeFinalCountTokensAnthropicBeta
|
||
// ============================================================================
|
||
|
||
func TestComputeFinalCountTokensAnthropicBeta_OAuthMimic_AlwaysIncludesContextManagement(t *testing.T) {
|
||
// count_tokens 路径下 mimic 不按 haiku 排除：始终注入完整 mimicry beta
|
||
s := newTestGatewayServiceForBeta(false)
|
||
final, ok := s.computeFinalCountTokensAnthropicBeta("oauth", true, "claude-haiku-4-5", http.Header{}, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"count_tokens + mimic 即使 haiku 也注入 context-management beta（与 messages 不同）")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaTokenCounting),
|
||
"count_tokens 路径必须含 token-counting beta")
|
||
}
|
||
|
||
// 重构等价性回归：
|
||
// 原 main buildCountTokensRequest 在 count_tokens mimic 分支上不跳过白名单透传
|
||
// （与 messages mimic 不同），incomingBeta 取自客户端透传。重构后必须从 clientHeaders
|
||
// 拿同一个值并 merge，否则会丢失客户端 beta。
|
||
func TestComputeFinalCountTokensAnthropicBeta_OAuthMimic_PreservesClientBeta(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "custom-experimental-beta,context-1m-2025-08-07")
|
||
final, ok := s.computeFinalCountTokensAnthropicBeta("oauth", true, "claude-haiku-4-5", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, "custom-experimental-beta"),
|
||
"count_tokens mimic 不同于 messages mimic：原代码会保留客户端透传的 beta")
|
||
require.True(t, anthropicBetaTokensContains(final, "context-1m-2025-08-07"),
|
||
"客户端透传的其他 beta token 同样需要保留")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"同时 FullClaudeCodeMimicryBetas 不打折扣")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaTokenCounting),
|
||
"同时补齐 token-counting beta")
|
||
}
|
||
|
||
// messages mimic 路径反向验证：原代码会跳过白名单透传，
|
||
// 客户端 beta 不会进入 mimic 计算。重构后 messages computeFinalAnthropicBeta
|
||
// mimic 分支依然不该使用 clientBeta。
|
||
func TestComputeFinalAnthropicBeta_OAuthMimic_IgnoresClientBetaExplicit(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "custom-experimental-beta")
|
||
final, ok := s.computeFinalAnthropicBeta("oauth", true, "claude-sonnet-4-6", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.False(t, anthropicBetaTokensContains(final, "custom-experimental-beta"),
|
||
"messages mimic 原代码跳过白名单透传 → 客户端 beta 不进入计算。"+
|
||
"与 count_tokens mimic 是不同的设计，不能合并为同一函数。")
|
||
}
|
||
|
||
func TestComputeFinalCountTokensAnthropicBeta_OAuthTransparent_NoClientBetaInjectsDefault(t *testing.T) {
|
||
// 真 CC 客户端透传 + 客户端未传 anthropic-beta → 用 CountTokensBetaHeader 兜底
|
||
s := newTestGatewayServiceForBeta(false)
|
||
final, ok := s.computeFinalCountTokensAnthropicBeta("oauth", false, "claude-haiku-4-5", http.Header{}, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.Equal(t, claude.CountTokensBetaHeader, final)
|
||
// CountTokensBetaHeader 不含 context-management beta
|
||
require.False(t, anthropicBetaTokensContains(final, claude.BetaContextManagement))
|
||
}
|
||
|
||
func TestComputeFinalCountTokensAnthropicBeta_OAuthTransparent_AppendsBetaTokenCounting(t *testing.T) {
|
||
s := newTestGatewayServiceForBeta(false)
|
||
hdr := http.Header{}
|
||
hdr.Set("anthropic-beta", "oauth-2025-04-20,context-management-2025-06-27")
|
||
final, ok := s.computeFinalCountTokensAnthropicBeta("oauth", false, "claude-sonnet-4-6", hdr, []byte(`{}`), nil)
|
||
require.True(t, ok)
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaTokenCounting),
|
||
"客户端未带 token-counting beta 时必须补齐")
|
||
require.True(t, anthropicBetaTokensContains(final, claude.BetaContextManagement),
|
||
"客户端带的 context-management beta 必须保留")
|
||
}
|
||
|
||
// ============================================================================
|
||
// normalizeClaudeOAuthRequestBody — 回归：context_management 补齐恢复原行为
|
||
// ============================================================================
|
||
//
|
||
// 重构后该函数不再按 model 名短路：thinking=enabled/adaptive 时补齐 context_management，
|
||
// 与 model 无关。strip 责任移交 sanitizeAnthropicBodyForBetaTokens（在
|
||
// buildUpstreamRequest 层按最终 beta header 执行）。
|
||
|
||
func TestNormalizeClaudeOAuthRequestBody_InjectsContextManagement_ThinkingEnabled(t *testing.T) {
|
||
body := []byte(`{"model":"claude-sonnet-4-6","thinking":{"type":"enabled","budget_tokens":1000},"messages":[]}`)
|
||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-sonnet-4-6", claudeOAuthNormalizeOptions{})
|
||
require.True(t, gjson.GetBytes(out, "context_management").Exists())
|
||
require.Equal(t, "clear_thinking_20251015",
|
||
gjson.GetBytes(out, "context_management.edits.0.type").String())
|
||
}
|
||
|
||
func TestNormalizeClaudeOAuthRequestBody_InjectsContextManagement_ThinkingAdaptive(t *testing.T) {
|
||
body := []byte(`{"model":"claude-opus-4-7","thinking":{"type":"adaptive"},"messages":[]}`)
|
||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-opus-4-7", claudeOAuthNormalizeOptions{})
|
||
require.True(t, gjson.GetBytes(out, "context_management").Exists())
|
||
}
|
||
|
||
func TestNormalizeClaudeOAuthRequestBody_HaikuStillInjects_StripDeferredToSanitize(t *testing.T) {
|
||
// haiku + thinking=enabled：normalize 阶段仍按 CLI mimicry 行为补齐字段；
|
||
// strip 由 buildUpstreamRequest 层的 sanitize 兜底（如果 final beta 不含 token）。
|
||
body := []byte(`{"model":"claude-haiku-4-5","thinking":{"type":"enabled","budget_tokens":1000},"messages":[]}`)
|
||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-haiku-4-5", claudeOAuthNormalizeOptions{})
|
||
require.True(t, gjson.GetBytes(out, "context_management").Exists(),
|
||
"normalize 不再按 model 名短路；strip 责任移交 sanitize 层")
|
||
}
|
||
|
||
func TestNormalizeClaudeOAuthRequestBody_PreservesClientContextManagement(t *testing.T) {
|
||
body := []byte(`{"model":"claude-opus-4-7","context_management":{"edits":[{"type":"custom_strategy"}]},"thinking":{"type":"enabled","budget_tokens":1000},"messages":[]}`)
|
||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-opus-4-7", claudeOAuthNormalizeOptions{})
|
||
require.Equal(t, "custom_strategy",
|
||
gjson.GetBytes(out, "context_management.edits.0.type").String(),
|
||
"客户端透传的 context_management 内容必须原样保留")
|
||
}
|
||
|
||
func TestNormalizeClaudeOAuthRequestBody_NoThinking_NoInject(t *testing.T) {
|
||
body := []byte(`{"model":"claude-sonnet-4-6","messages":[]}`)
|
||
out, _ := normalizeClaudeOAuthRequestBody(body, "claude-sonnet-4-6", claudeOAuthNormalizeOptions{})
|
||
require.False(t, gjson.GetBytes(out, "context_management").Exists())
|
||
}
|
||
|
||
// ============================================================================
|
||
// passthrough 集成测试：buildUpstreamRequest-
|
||
// AnthropicAPIKeyPassthrough 与 buildCountTokensRequestAnthropicAPIKeyPassthrough
|
||
// 路径上 sanitize 是否生效。
|
||
// ============================================================================
|
||
|
||
// passthrough 集成测试不设 base_url，避开 validateUpstreamBaseURL 对 cfg.Security 的依赖。
|
||
// targetURL 会走默认 claudeAPIURL，sanitize 逻辑与 baseURL 是否存在无关。
|
||
func newAnthropicAPIKeyPassthroughAccountForBetaTest() *Account {
|
||
return &Account{
|
||
ID: 501,
|
||
Name: "anthropic-apikey-passthrough-ctxmgmt-test",
|
||
Platform: PlatformAnthropic,
|
||
Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{
|
||
"api_key": "upstream-key",
|
||
},
|
||
Extra: map[string]any{"anthropic_passthrough": true},
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
}
|
||
}
|
||
|
||
func readUpstreamBodyForTest(t *testing.T, req *http.Request) []byte {
|
||
t.Helper()
|
||
require.NotNil(t, req.Body)
|
||
b, err := io.ReadAll(req.Body)
|
||
require.NoError(t, err)
|
||
return b
|
||
}
|
||
|
||
func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_StripsContextManagementWhenClientHeaderMissingBeta(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
// 客户端仅带 oauth beta，不带 context-management-2025-06-27
|
||
c.Request.Header.Set("Anthropic-Beta", "oauth-2025-04-20")
|
||
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
|
||
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
|
||
)
|
||
require.NoError(t, err)
|
||
require.False(t, gjson.GetBytes(readUpstreamBodyForTest(t, req), "context_management").Exists(),
|
||
"API-key passthrough + 客户端未带 context-management beta → strip body 字段")
|
||
}
|
||
|
||
func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_PreservesContextManagementWhenClientHeaderHasBeta(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
c.Request.Header.Set("Anthropic-Beta", "oauth-2025-04-20,context-management-2025-06-27")
|
||
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
|
||
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
|
||
)
|
||
require.NoError(t, err)
|
||
require.True(t, gjson.GetBytes(readUpstreamBodyForTest(t, req), "context_management").Exists(),
|
||
"API-key passthrough + 客户端带 context-management beta → 字段保留（不过度删除）")
|
||
}
|
||
|
||
func TestBuildCountTokensRequestAnthropicAPIKeyPassthrough_StripsContextManagementWhenClientHeaderMissingBeta(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
|
||
c.Request.Header.Set("Anthropic-Beta", "oauth-2025-04-20,token-counting-2024-11-01")
|
||
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, err := svc.buildCountTokensRequestAnthropicAPIKeyPassthrough(
|
||
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
|
||
)
|
||
require.NoError(t, err)
|
||
require.False(t, gjson.GetBytes(readUpstreamBodyForTest(t, req), "context_management").Exists(),
|
||
"count_tokens passthrough + 客户端未带 context-management beta → strip")
|
||
}
|
||
|
||
// ============================================================================
|
||
// 集成测试：buildUpstreamRequest
|
||
// 全路径验证上游 outgoing body 与 anthropic-beta header 严格对称。
|
||
// 这个测试能挡住未来某人忘调 sanitize / 将 sanitize 挪到 CCH 之后 等 regression。
|
||
// ============================================================================
|
||
|
||
func TestBuildUpstreamRequest_OAuthMimicHaiku_StripsContextManagementEndToEnd(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
account := &Account{ID: 401, Platform: PlatformAnthropic, Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"access_token": "oauth-tok"},
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
}
|
||
// haiku + mimic CC → final beta = HaikuBetaHeader（不含 context-management）→
|
||
// body 必须 strip。
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequest(
|
||
context.Background(), c, account, body,
|
||
"oauth-tok", "oauth", "claude-haiku-4-5", false, true, // mimicClaudeCode=true
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
outBeta := getHeaderRaw(req.Header, "anthropic-beta")
|
||
|
||
require.False(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"OAuth mimic + haiku 端到端：outgoing body 不应含 context_management")
|
||
require.False(t, anthropicBetaTokensContains(outBeta, claude.BetaContextManagement),
|
||
"对称约束：outgoing anthropic-beta header 也不带 context-management beta")
|
||
}
|
||
|
||
func TestBuildUpstreamRequest_OAuthMimicNonHaiku_PreservesContextManagementEndToEnd(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
account := &Account{ID: 402, Platform: PlatformAnthropic, Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"access_token": "oauth-tok"},
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
}
|
||
// sonnet + mimic CC → final beta = FullClaudeCodeMimicryBetas（含 context-management）→
|
||
// body 保留。
|
||
body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequest(
|
||
context.Background(), c, account, body,
|
||
"oauth-tok", "oauth", "claude-sonnet-4-6", false, true,
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
outBeta := getHeaderRaw(req.Header, "anthropic-beta")
|
||
|
||
require.True(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"OAuth mimic + non-haiku：outgoing body 必须保留 context_management。")
|
||
require.True(t, anthropicBetaTokensContains(outBeta, claude.BetaContextManagement),
|
||
"对称约束：outgoing anthropic-beta header 同时含 context-management beta")
|
||
}
|
||
|
||
func TestBuildUpstreamRequest_OAuthTransparentHaikuWithRealCCBeta_PreservesField(t *testing.T) {
|
||
// 端到端验证：真 CC 客户端 + haiku + 客户端 header 带 context-management beta
|
||
// → final beta 透传 → 不应该过度删除 body 字段
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
c.Request.Header.Set("Anthropic-Beta",
|
||
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27")
|
||
|
||
account := &Account{ID: 403, Platform: PlatformAnthropic, Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"access_token": "oauth-tok"},
|
||
Status: StatusActive, Schedulable: true,
|
||
}
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequest(
|
||
context.Background(), c, account, body,
|
||
"oauth-tok", "oauth", "claude-haiku-4-5", false, false, // mimicClaudeCode=false（真 CC）
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
outBeta := getHeaderRaw(req.Header, "anthropic-beta")
|
||
|
||
require.True(t, anthropicBetaTokensContains(outBeta, claude.BetaContextManagement),
|
||
"真 CC 透传路径：客户端 header 中的 context-management beta 必须保留")
|
||
require.True(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"回归保护：真 CC + haiku + 客户端带 beta token 时，clear_thinking_20251015 功能不能静默失效")
|
||
}
|
||
|
||
// CCH 顺序语义测试：sanitize 必须在 signBillingHeaderCCH 之前，
|
||
// 否则签名的 hash 与最终发送的 body 不一致，被 Anthropic 判 third-party。
|
||
//
|
||
// 该测试不走 buildUpstreamRequest 完整路径（需要 mock SettingService 成本高），
|
||
// 而是直接验证两个顺序产生的 cch 不同，证明二者不可交换。
|
||
// 测试名本身是语义约束的文档化 marker。
|
||
func TestSanitizeMustBeBeforeCCHSigning_HashConsistency(t *testing.T) {
|
||
// 构造 body：含 context_management + cch=00000 占位符
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.92; cch=00000;"}],"messages":[]}`)
|
||
|
||
// 最终发送场景：final beta 不含 context-management beta → sanitize 会 strip
|
||
finalBeta := "oauth-2025-04-20,interleaved-thinking-2025-05-14"
|
||
|
||
extractCCH := func(t *testing.T, b []byte) string {
|
||
t.Helper()
|
||
m := regexp.MustCompile(`\bcch=([0-9a-fA-F]{5})\b`).FindSubmatch(b)
|
||
require.NotNil(t, m, "body 里找不到 cch=<5hex> ：%s", string(b))
|
||
return string(m[1])
|
||
}
|
||
|
||
// === 正确顺序：sanitize → signBillingHeaderCCH ===
|
||
// 1. strip context_management
|
||
sanitizedFirst, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBeta)
|
||
require.True(t, changed)
|
||
require.False(t, gjson.GetBytes(sanitizedFirst, "context_management").Exists())
|
||
// 2. 基于“strip 后的 body”算 hash
|
||
correctFinal := signBillingHeaderCCH(sanitizedFirst)
|
||
correctCCH := extractCCH(t, correctFinal)
|
||
require.NotEqual(t, "00000", correctCCH, "placeholder 应被替换")
|
||
|
||
// === 错误顺序：signBillingHeaderCCH → sanitize（未来 regression 场景）===
|
||
// 1. 先基于“含 context_management 的 body”算 hash → cch=H_with
|
||
signedFirst := signBillingHeaderCCH(body)
|
||
wrongCCH := extractCCH(t, signedFirst)
|
||
require.NotEqual(t, "00000", wrongCCH)
|
||
// 2. 后 strip context_management → body 变化但 cch 仍是 H_with
|
||
wrongFinal, _ := sanitizeAnthropicBodyForBetaTokens(signedFirst, finalBeta)
|
||
wrongFinalCCH := extractCCH(t, wrongFinal)
|
||
|
||
// === 关键断言 ===
|
||
// 上游验证逻辑：将 outgoing body 的 cch 还原为 00000、重算 hash、与 cch 字段比较。
|
||
// 模拟上游验证：用发送 body 算出“期望的 cch”，与发送 body 里的 cch 字段比。
|
||
recomputeExpected := func(b []byte, currentCCH string) string {
|
||
t.Helper()
|
||
// 把 cch=<currentCCH> 还原为 cch=00000
|
||
re := regexp.MustCompile(`(\bcch=)` + currentCCH + `(\b)`)
|
||
restored := re.ReplaceAll(b, []byte("${1}00000${2}"))
|
||
return extractCCH(t, signBillingHeaderCCH(restored))
|
||
}
|
||
|
||
// 正确顺序：发送 body 的 cch == 重算 hash → 上游验证过
|
||
require.Equal(t, correctCCH, recomputeExpected(correctFinal, correctCCH),
|
||
"正确顺序：final body 里的 cch 与重算 hash 一致 → 上游验证通过")
|
||
|
||
// 错误顺序：发送 body 的 cch 是“含 ctx 算的”，但最终 body 不含 ctx → 重算 hash 不同
|
||
require.NotEqual(t, wrongFinalCCH, recomputeExpected(wrongFinal, wrongFinalCCH),
|
||
"错误顺序：final body 里的 cch 是基于含 ctx 的 body 算的，"+
|
||
"但发送 body 已 strip ctx → 上游重算 hash 与 cch 不一致 → 被判 third-party。"+
|
||
"这是 buildUpstreamRequest / buildCountTokensRequest 里 sanitize 必须在 "+
|
||
"signBillingHeaderCCH 之前的原因。")
|
||
}
|
||
|
||
// count_tokens 主路径 E2E 集成测试
|
||
func TestBuildCountTokensRequest_OAuthMimicHaiku_PreservesContextManagementEndToEnd(t *testing.T) {
|
||
// count_tokens 路径下 mimic 不按 haiku 排除，始终注入 BetaContextManagement
|
||
// → sanitize 看到最终 beta header 含 context-management beta → 字段保留。
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
|
||
|
||
account := &Account{ID: 411, Platform: PlatformAnthropic, Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"access_token": "oauth-tok"},
|
||
Status: StatusActive, Schedulable: true,
|
||
}
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildCountTokensRequest(
|
||
context.Background(), c, account, body,
|
||
"oauth-tok", "oauth", "claude-haiku-4-5", true, // mimicClaudeCode=true
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
outBeta := getHeaderRaw(req.Header, "anthropic-beta")
|
||
|
||
require.True(t, anthropicBetaTokensContains(outBeta, claude.BetaContextManagement),
|
||
"count_tokens mimic 始终注入 context-management beta")
|
||
require.True(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"对称约束：final beta 含 token 时 body 字段保留")
|
||
require.True(t, anthropicBetaTokensContains(outBeta, claude.BetaTokenCounting),
|
||
"count_tokens 路径必须含 token-counting beta")
|
||
}
|
||
|
||
func TestBuildCountTokensRequest_APIKeyHaiku_StripsContextManagementEndToEnd(t *testing.T) {
|
||
// API-key + haiku + 客户端 header 不带 context-management beta → final beta 不含 → strip
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
|
||
c.Request.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
|
||
|
||
account := &Account{ID: 412, Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{"api_key": "sk-ant-xxx"},
|
||
Status: StatusActive, Schedulable: true,
|
||
}
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildCountTokensRequest(
|
||
context.Background(), c, account, body,
|
||
"sk-ant-xxx", "apikey", "claude-haiku-4-5", false,
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
require.False(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"count_tokens API-key + 客户端未带 beta token → body strip")
|
||
}
|
||
|
||
// count_tokens passthrough preserve 测试
|
||
func TestBuildCountTokensRequestAnthropicAPIKeyPassthrough_PreservesContextManagementWhenClientHeaderHasBeta(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
|
||
c.Request.Header.Set("Anthropic-Beta", "oauth-2025-04-20,context-management-2025-06-27,token-counting-2024-11-01")
|
||
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, err := svc.buildCountTokensRequestAnthropicAPIKeyPassthrough(
|
||
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
|
||
)
|
||
require.NoError(t, err)
|
||
require.True(t, gjson.GetBytes(readUpstreamBodyForTest(t, req), "context_management").Exists(),
|
||
"count_tokens passthrough + 客户端带 context-management beta → 字段保留")
|
||
}
|
||
|
||
func TestBuildUpstreamRequest_APIKeyHaikuWithContextManagement_StripsField(t *testing.T) {
|
||
// API-key + haiku + body 带 context_management + 客户端 header 未带 context-management beta
|
||
// → final beta 不含 → body 字段被 strip
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
c.Request.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
|
||
|
||
account := &Account{ID: 404, Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{"api_key": "sk-ant-xxx"},
|
||
Status: StatusActive, Schedulable: true,
|
||
}
|
||
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`)
|
||
svc := &GatewayService{cfg: &config.Config{}}
|
||
req, _, err := svc.buildUpstreamRequest(
|
||
context.Background(), c, account, body,
|
||
"sk-ant-xxx", "apikey", "claude-haiku-4-5", false, false,
|
||
)
|
||
require.NoError(t, err)
|
||
|
||
outBody := readUpstreamBodyForTest(t, req)
|
||
require.False(t, gjson.GetBytes(outBody, "context_management").Exists(),
|
||
"API-key + haiku + 客户端未带 beta token → body 字段必须被 strip")
|
||
}
|