zizi-al/sub2api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(gateway): allow Claude Code count_tokens

Browse Source
This commit is contained in:
wucm667 2026-05-31 08:43:20 +08:00
parent f18451e56f
commit bf3787de1f
2 changed files with 102 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
backend/internal/service/claude_code_validator.go
View File

@ -56,7 +56,7 @@ func NewClaudeCodeValidator() *ClaudeCodeValidator {
// 采用与 claude-relay-service 完全一致的验证策略：
//
// Step 1: User-Agent 检查 (必需) - 必须是 claude-cli/x.x.x
// Step 2: 对于非 messages 路径，只要 UA 匹配就通过
// Step 2: 对于非 messages 路径和 /messages/count_tokens，只要 UA 匹配就通过
// Step 3: 检查 max_tokens=1 + haiku 探测请求绕过（UA 已验证）
// Step 4: 对于 messages 路径，进行严格验证：
// - System prompt 相似度检查
@ -71,12 +71,17 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo
return false
}
// Step 2: 非 messages 路径，只要 UA 匹配就通过
// Step 2: 非 messages 路径只要 UA 匹配就通过
path := r.URL.Path
if !strings.Contains(path, "messages") {
return true
}
// count_tokens 是 Claude Code 官方辅助请求，通常不携带完整 messages system prompt。
if isMessagesCountTokensPath(path) {
return true
}
// Step 3: 检查 max_tokens=1 + haiku 探测请求绕过
// 这类请求用于 Claude Code 验证 API 连通性，不携带 system prompt
if isMaxTokensOneHaiku, ok := IsMaxTokensOneHaikuRequestFromContext(r.Context()); ok && isMaxTokensOneHaiku {
@ -128,6 +133,10 @@ func (v *ClaudeCodeValidator) Validate(r *http.Request, body map[string]any) boo
return true
}
func isMessagesCountTokensPath(path string) bool {
return strings.HasSuffix(path, "/messages/count_tokens")
}
// hasClaudeCodeSystemPrompt 检查请求是否包含 Claude Code 系统提示词
// 使用字符串相似度匹配（Dice coefficient）
func (v *ClaudeCodeValidator) hasClaudeCodeSystemPrompt(body map[string]any) bool {

91
backend/internal/service/claude_code_validator_test.go
View File

@ -48,6 +48,97 @@ func TestClaudeCodeValidator_MessagesWithoutProbeStillNeedStrictValidation(t *te
require.False(t, ok)
}
func TestClaudeCodeValidator_CountTokensPathUAOnly(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
})
require.True(t, ok)
}
func TestClaudeCodeValidator_CountTokensPathRequiresUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesPathFullValid(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"system": []any{
map[string]any{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.True(t, ok)
}
func TestClaudeCodeValidator_MessagesPathRejectsNonClaudeCodeUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"system": []any{
map[string]any{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesPathWithoutSystemPromptStillRejected(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"messages": []any{
map[string]any{"role": "user", "content": "hello"},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_NonMessagesPathUAOnly(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/models", nil)