diff --git a/internal/runtime/executor/claude_executor_request.go b/internal/runtime/executor/claude_executor_request.go index a15e3204..25bd571d 100644 --- a/internal/runtime/executor/claude_executor_request.go +++ b/internal/runtime/executor/claude_executor_request.go @@ -10,6 +10,7 @@ import ( "io" "net/http" "net/url" + "sort" "strings" "github.com/andybalholm/brotli" @@ -849,7 +850,283 @@ func remapOAuthToolNames(body []byte) ([]byte, map[string]string) { return remapOAuthToolNamesWithOptions(body, claudeMCPAliasOptions{secret: "cpa-claude-mcp-default-caller"}) } +type claudeRawJSONEdit struct { + start int + end int + replacement string +} + func remapOAuthToolNamesWithOptions(body []byte, mcpAliases claudeMCPAliasOptions) ([]byte, map[string]string) { + remapped, reverseMap, ok := remapOAuthToolNamesWithBatchedEdits(body, mcpAliases) + if ok { + return remapped, reverseMap + } + return remapOAuthToolNamesWithOptionsLegacy(body, mcpAliases) +} + +// remapOAuthToolNamesWithBatchedEdits records offsets from the original JSON +// and applies every rename in one copy. Repeated sjson.SetBytes calls copy most +// of the request for every historical tool reference, turning this path into +// O(body size * reference count) allocation growth. +func remapOAuthToolNamesWithBatchedEdits(body []byte, mcpAliases claudeMCPAliasOptions) ([]byte, map[string]string, bool) { + if !gjson.ValidBytes(body) { + return nil, nil, false + } + + reverseMap := make(map[string]string) + recordRename := func(original, renamed string) { + // Preserve the first-seen original name if the same upstream name is + // produced from multiple call sites; they all map back identically. + if _, exists := reverseMap[renamed]; !exists { + reverseMap[renamed] = original + } + } + + // Build one request-specific forward map from declarations. Every client + // tool, including typed custom declarations and names resembling Claude + // built-ins, gets an MCP alias. Historical references use this same map. + tools := gjson.GetBytes(body, "tools") + forwardMap := make(map[string]string) + protectedNames := make(map[string]bool) + reservedNames := helps.AugmentClaudeBuiltinToolRegistry(body, nil) + if tools.Exists() && tools.IsArray() { + tools.ForEach(func(_, tool gjson.Result) bool { + name := tool.Get("name").String() + if name != "" { + reservedNames[name] = true + } + if helps.IsClaudeServerToolType(tool.Get("type").String()) { + protectedNames[name] = true + } + return true + }) + tools.ForEach(func(_, tool gjson.Result) bool { + if helps.IsClaudeServerToolType(tool.Get("type").String()) { + return true + } + name := tool.Get("name").String() + if name == "" || helps.IsClaudeMCPToolName(name) { + return true + } + if _, exists := forwardMap[name]; exists { + return true + } + for attempt := uint32(0); ; attempt++ { + alias := helps.ClaudeMCPToolAlias(mcpAliases.secret, name, attempt) + if reservedNames[alias] { + continue + } + forwardMap[name] = alias + reservedNames[alias] = true + break + } + return true + }) + } + + rewriteName := func(name string) (string, bool) { + if name == "" || protectedNames[name] || helps.IsClaudeMCPToolName(name) { + return name, false + } + if newName, ok := forwardMap[name]; ok && newName != name { + return newName, true + } + return name, false + } + + edits := make([]claudeRawJSONEdit, 0, len(forwardMap)+1) + appendRawEdit := func(result gjson.Result, replacement string) bool { + start := result.Index + end := start + len(result.Raw) + if result.Raw == "" || start < 0 || end < start || end > len(body) || !bytes.Equal(body[start:end], []byte(result.Raw)) { + return false + } + edits = append(edits, claudeRawJSONEdit{start: start, end: end, replacement: replacement}) + return true + } + appendStringEdit := func(result gjson.Result, replacement string) bool { + // ClaudeMCPToolAlias only emits [A-Za-z0-9_-], so adding quotes is + // byte-identical to sjson's encoding without another allocation. + return appendRawEdit(result, `"`+replacement+`"`) + } + + // 1. Rebuild typed custom tools exactly as before, but replace the original + // tools array only after all offsets have been collected. + toolsNeedRewrite := false + if tools.Exists() && tools.IsArray() { + tools.ForEach(func(_, tool gjson.Result) bool { + toolType := tool.Get("type").String() + if helps.IsClaudeServerToolType(toolType) { + return true + } + if strings.TrimSpace(toolType) != "" { + toolsNeedRewrite = true + return false + } + name := tool.Get("name").String() + _, toolsNeedRewrite = rewriteName(name) + return !toolsNeedRewrite + }) + } + if toolsNeedRewrite { + var toolsJSON strings.Builder + toolsJSON.WriteByte('[') + toolCount := 0 + tools.ForEach(func(_, tool gjson.Result) bool { + if helps.IsClaudeServerToolType(tool.Get("type").String()) { + if toolCount > 0 { + toolsJSON.WriteByte(',') + } + toolsJSON.WriteString(tool.Raw) + toolCount++ + return true + } + + name := tool.Get("name").String() + toolJSON := tool.Raw + if strings.TrimSpace(tool.Get("type").String()) != "" { + if updatedTool, errDelete := sjson.Delete(toolJSON, "type"); errDelete == nil { + toolJSON = updatedTool + } + } + if newName, renamed := rewriteName(name); renamed { + updatedTool, err := sjson.Set(toolJSON, "name", newName) + if err == nil { + toolJSON = updatedTool + recordRename(name, newName) + } + } + + if toolCount > 0 { + toolsJSON.WriteByte(',') + } + toolsJSON.WriteString(toolJSON) + toolCount++ + return true + }) + toolsJSON.WriteByte(']') + if !appendRawEdit(tools, toolsJSON.String()) { + return nil, nil, false + } + } + + // 2. Rename tool_choice if it references a declared client tool. + toolChoice := gjson.GetBytes(body, "tool_choice") + if toolChoice.Get("type").String() == "tool" { + nameResult := toolChoice.Get("name") + tcName := nameResult.String() + if newName, renamed := rewriteName(tcName); renamed { + if !appendStringEdit(nameResult, newName) { + return nil, nil, false + } + recordRename(tcName, newName) + } + } + + // 3. Rename tool references in messages while every Result.Index still + // points into the original request bytes. + messages := gjson.GetBytes(body, "messages") + validOffsets := true + if messages.Exists() && messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if !content.Exists() || !content.IsArray() { + return true + } + content.ForEach(func(_, part gjson.Result) bool { + switch part.Get("type").String() { + case "tool_use": + nameResult := part.Get("name") + name := nameResult.String() + if newName, renamed := rewriteName(name); renamed { + if !appendStringEdit(nameResult, newName) { + validOffsets = false + return false + } + recordRename(name, newName) + } + case "tool_reference": + nameResult := part.Get("tool_name") + toolName := nameResult.String() + if newName, renamed := rewriteName(toolName); renamed { + if !appendStringEdit(nameResult, newName) { + validOffsets = false + return false + } + recordRename(toolName, newName) + } + case "tool_result": + nestedContent := part.Get("content") + if nestedContent.Exists() && nestedContent.IsArray() { + nestedContent.ForEach(func(_, nestedPart gjson.Result) bool { + if nestedPart.Get("type").String() != "tool_reference" { + return true + } + nameResult := nestedPart.Get("tool_name") + nestedToolName := nameResult.String() + if newName, renamed := rewriteName(nestedToolName); renamed { + if !appendStringEdit(nameResult, newName) { + validOffsets = false + return false + } + recordRename(nestedToolName, newName) + } + return true + }) + } + } + return validOffsets + }) + return validOffsets + }) + } + if !validOffsets { + return nil, nil, false + } + + remapped, ok := applyClaudeRawJSONEdits(body, edits) + if !ok { + return nil, nil, false + } + return remapped, reverseMap, true +} + +func applyClaudeRawJSONEdits(body []byte, edits []claudeRawJSONEdit) ([]byte, bool) { + if len(edits) == 0 { + return body, true + } + sort.Slice(edits, func(i, j int) bool { + return edits[i].start < edits[j].start + }) + + finalSize := len(body) + cursor := 0 + for _, edit := range edits { + if edit.start < cursor || edit.start < 0 || edit.end < edit.start || edit.end > len(body) { + return nil, false + } + finalSize += len(edit.replacement) - (edit.end - edit.start) + if finalSize < 0 { + return nil, false + } + cursor = edit.end + } + + out := make([]byte, 0, finalSize) + cursor = 0 + for _, edit := range edits { + out = append(out, body[cursor:edit.start]...) + out = append(out, edit.replacement...) + cursor = edit.end + } + out = append(out, body[cursor:]...) + return out, true +} + +// remapOAuthToolNamesWithOptionsLegacy is the byte-for-byte compatibility +// fallback for malformed JSON or an unexpected GJSON offset. Keep it available +// as a differential-test oracle for the batched implementation. +func remapOAuthToolNamesWithOptionsLegacy(body []byte, mcpAliases claudeMCPAliasOptions) ([]byte, map[string]string) { reverseMap := make(map[string]string) recordRename := func(original, renamed string) { // Preserve the first-seen original name if the same upstream name is diff --git a/internal/runtime/executor/claude_executor_request_bench_test.go b/internal/runtime/executor/claude_executor_request_bench_test.go new file mode 100644 index 00000000..15584520 --- /dev/null +++ b/internal/runtime/executor/claude_executor_request_bench_test.go @@ -0,0 +1,105 @@ +package executor + +import ( + "encoding/json" + "fmt" + "strings" + "testing" +) + +type claudeOAuthRemapBenchmarkBody struct { + Model string `json:"model"` + Tools []claudeOAuthRemapBenchmarkTool `json:"tools"` + ToolChoice claudeOAuthRemapBenchmarkChoice `json:"tool_choice"` + Messages []claudeOAuthRemapBenchmarkMessage `json:"messages"` + Padding string `json:"padding"` +} + +type claudeOAuthRemapBenchmarkTool struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema map[string]any `json:"input_schema"` +} + +type claudeOAuthRemapBenchmarkChoice struct { + Type string `json:"type"` + Name string `json:"name"` +} + +type claudeOAuthRemapBenchmarkMessage struct { + Role string `json:"role"` + Content []any `json:"content"` +} + +func BenchmarkRemapOAuthToolNames(b *testing.B) { + benchmarks := []struct { + name string + targetSize int + references int + }{ + {name: "4KiB_8Refs", targetSize: 4 << 10, references: 8}, + {name: "64KiB_100Refs", targetSize: 64 << 10, references: 100}, + {name: "256KiB_500Refs", targetSize: 256 << 10, references: 500}, + } + + for _, benchmark := range benchmarks { + b.Run(benchmark.name, func(b *testing.B) { + body := buildClaudeOAuthRemapBenchmarkBody(b, benchmark.targetSize, benchmark.references) + options := claudeMCPAliasOptions{secret: "benchmark-caller"} + b.ReportAllocs() + b.SetBytes(int64(len(body))) + for b.Loop() { + remapped, reverseMap := remapOAuthToolNamesWithOptions(body, options) + if len(remapped) == 0 || len(reverseMap) == 0 { + b.Fatal("remap returned empty output") + } + } + }) + } +} + +func buildClaudeOAuthRemapBenchmarkBody(tb testing.TB, targetSize, references int) []byte { + tb.Helper() + + const toolCount = 20 + tools := make([]claudeOAuthRemapBenchmarkTool, 0, toolCount) + for i := range toolCount { + tools = append(tools, claudeOAuthRemapBenchmarkTool{ + Name: fmt.Sprintf("benchmark_tool_%02d", i), + Description: "Benchmark tool with a stable representative schema.", + InputSchema: map[string]any{"type": "object", "properties": map[string]any{"value": map[string]any{"type": "string"}}}, + }) + } + + content := make([]any, 0, references) + for i := range references { + name := tools[i%len(tools)].Name + switch i % 3 { + case 0: + content = append(content, map[string]any{"type": "tool_use", "id": fmt.Sprintf("toolu_%04d", i), "name": name, "input": map[string]any{"value": i}}) + case 1: + content = append(content, map[string]any{"type": "tool_reference", "tool_name": name}) + default: + content = append(content, map[string]any{"type": "tool_result", "tool_use_id": fmt.Sprintf("toolu_%04d", i), "content": []any{map[string]any{"type": "tool_reference", "tool_name": name}}}) + } + } + + request := claudeOAuthRemapBenchmarkBody{ + Model: "claude-opus-5", + Tools: tools, + ToolChoice: claudeOAuthRemapBenchmarkChoice{Type: "tool", Name: tools[0].Name}, + Messages: []claudeOAuthRemapBenchmarkMessage{{Role: "assistant", Content: content}}, + } + body, errMarshal := json.Marshal(request) + if errMarshal != nil { + tb.Fatalf("marshal benchmark request: %v", errMarshal) + } + if remaining := targetSize - len(body); remaining > 0 { + request.Padding = strings.Repeat("x", remaining) + body, errMarshal = json.Marshal(request) + if errMarshal != nil { + tb.Fatalf("marshal padded benchmark request: %v", errMarshal) + } + } + return body +} diff --git a/internal/runtime/executor/claude_executor_request_remap_test.go b/internal/runtime/executor/claude_executor_request_remap_test.go new file mode 100644 index 00000000..0488ce9d --- /dev/null +++ b/internal/runtime/executor/claude_executor_request_remap_test.go @@ -0,0 +1,142 @@ +package executor + +import ( + "bytes" + "fmt" + "maps" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/tidwall/gjson" +) + +func TestRemapOAuthToolNamesWithBatchedEditsMatchesLegacyBytes(t *testing.T) { + secret := "differential-caller" + collision := helps.ClaudeMCPToolAlias(secret, "fetch_url", 0) + longName := "读取_" + strings.Repeat("very_long_tool_name_", 8) + tests := []struct { + name string + body []byte + }{ + { + name: "all reference shapes and undeclared history", + body: []byte(`{"model":"claude-opus-5","tools":[{"name":"search_web","input_schema":{"type":"object"}},{"name":"Search_Web","input_schema":{"type":"object"}}],"tool_choice":{"type":"tool","name":"search_web"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"search_web","input":{}},{"type":"tool_reference","tool_name":"Search_Web"},{"type":"tool_result","tool_use_id":"toolu_1","content":[{"type":"tool_reference","tool_name":"search_web"}]},{"type":"tool_use","id":"toolu_unknown","name":"not_declared","input":{}}]}]}`), + }, + { + name: "typed custom server existing MCP and duplicate declaration", + body: []byte(`{"tools":[{"type":"custom","name":"client_custom","input_schema":{"type":"object"}},{"type":"web_search_20250305","name":"web_search"},{"name":"mcp__context7__query-docs"},{"name":"client_custom"}],"messages":[{"role":"assistant","content":[{"type":"tool_use","name":"client_custom","id":"toolu_1","input":{}},{"type":"tool_reference","tool_name":"web_search"},{"type":"tool_reference","tool_name":"mcp__context7__query-docs"}]}]}`), + }, + { + name: "alias collision", + body: []byte(fmt.Sprintf(`{"tools":[{"name":%q},{"name":"fetch_url"}],"tool_choice":{"type":"tool","name":"fetch_url"}}`, collision)), + }, + { + name: "unicode long and case distinct names", + body: []byte(fmt.Sprintf(`{"messages":[{"content":[{"name":%q,"type":"tool_use"},{"tool_name":"read_file","type":"tool_reference"}]}],"tools":[{"name":%q},{"name":"read_file"}]}`, longName, longName)), + }, + { + name: "whitespace key order and escaped original", + body: []byte("{\n \"messages\" : [ { \"content\" : [ { \"name\" : \"fetch\\u005furl\", \"input\":{}, \"type\" : \"tool_use\" } ], \"role\" : \"assistant\" } ],\n \"unknown\" : {\"number\":1.2300,\"escaped\":\"a\\/b\\n<>&\"},\n \"tool_choice\" : { \"name\" : \"fetch\\u005furl\", \"type\" : \"tool\" },\n \"tools\" : [ { \"description\" : \"keep \\\"bytes\\\"\", \"name\" : \"fetch\\u005furl\", \"input_schema\" : { \"type\" : \"object\" } } ]\n}"), + }, + { + name: "non-string names follow legacy coercion", + body: []byte(`{"tools":[{"name":42}],"tool_choice":{"type":"tool","name":42},"messages":[{"content":[{"type":"tool_reference","tool_name":42}]}]}`), + }, + { + name: "no edits", + body: []byte(`{"tools":[{"type":"web_search_20250305","name":"web_search"},{"name":"mcp__server__existing"}],"messages":[{"content":[{"type":"tool_reference","tool_name":"unknown"}]}]}`), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + options := claudeMCPAliasOptions{secret: secret} + wantBody, wantReverseMap := remapOAuthToolNamesWithOptionsLegacy(test.body, options) + gotBody, gotReverseMap, ok := remapOAuthToolNamesWithBatchedEdits(test.body, options) + if !ok { + t.Fatal("batched remap unexpectedly rejected valid JSON offsets") + } + if !bytes.Equal(gotBody, wantBody) { + t.Fatalf("batched body differs from legacy bytes\n got: %s\nwant: %s", gotBody, wantBody) + } + if !maps.Equal(gotReverseMap, wantReverseMap) { + t.Fatalf("batched reverseMap = %v, want %v", gotReverseMap, wantReverseMap) + } + }) + } +} + +func TestRemapOAuthToolNamesWithBatchedEditsReturnsOriginalSliceWithoutEdits(t *testing.T) { + body := []byte(`{"tools":[{"type":"web_search_20250305","name":"web_search"}]}`) + out, reverseMap, ok := remapOAuthToolNamesWithBatchedEdits(body, claudeMCPAliasOptions{secret: "no-edits"}) + if !ok { + t.Fatal("batched remap rejected valid JSON") + } + if len(reverseMap) != 0 { + t.Fatalf("reverseMap = %v, want empty", reverseMap) + } + if len(out) == 0 || &out[0] != &body[0] { + t.Fatal("no-edit remap did not return the original slice") + } +} + +func TestRemapOAuthToolNamesWithOptionsFallsBackForMalformedJSON(t *testing.T) { + body := []byte(`{"tools":[{"name":"search_web"}],"messages":[`) + options := claudeMCPAliasOptions{secret: "malformed"} + if _, _, ok := remapOAuthToolNamesWithBatchedEdits(body, options); ok { + t.Fatal("batched remap accepted malformed JSON") + } + wantBody, wantReverseMap := remapOAuthToolNamesWithOptionsLegacy(body, options) + gotBody, gotReverseMap := remapOAuthToolNamesWithOptions(body, options) + if !bytes.Equal(gotBody, wantBody) || !maps.Equal(gotReverseMap, wantReverseMap) { + t.Fatalf("fallback differs from legacy: body=%q map=%v, want body=%q map=%v", gotBody, gotReverseMap, wantBody, wantReverseMap) + } +} + +func TestApplyClaudeRawJSONEditsRejectsInvalidRanges(t *testing.T) { + body := []byte(`{"a":"one","b":"two"}`) + tests := []struct { + name string + edits []claudeRawJSONEdit + }{ + {name: "overlap", edits: []claudeRawJSONEdit{{start: 5, end: 10}, {start: 8, end: 12}}}, + {name: "negative", edits: []claudeRawJSONEdit{{start: -1, end: 1}}}, + {name: "reversed", edits: []claudeRawJSONEdit{{start: 5, end: 4}}}, + {name: "past end", edits: []claudeRawJSONEdit{{start: 5, end: len(body) + 1}}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if _, ok := applyClaudeRawJSONEdits(body, test.edits); ok { + t.Fatal("invalid edits unexpectedly succeeded") + } + }) + } +} + +func FuzzRemapOAuthToolNamesWithBatchedEditsMatchesLegacy(f *testing.F) { + seeds := [][]byte{ + []byte(`{}`), + []byte(`{"tools":[{"name":"search_web"}]}`), + []byte(`{"tools":[{"type":"custom","name":"读取文件"}],"tool_choice":{"type":"tool","name":"读取文件"},"messages":[{"content":[{"type":"tool_use","name":"读取文件"}]}]}`), + []byte("{\n\"messages\":[{\"content\":[{\"type\":\"tool_reference\",\"tool_name\":\"a\\u005fb\"}]}],\"tools\":[{\"name\":\"a\\u005fb\"}]}"), + } + for _, seed := range seeds { + f.Add(seed, "fuzz-caller") + } + + f.Fuzz(func(t *testing.T, body []byte, secret string) { + if len(body) > 1<<20 || !gjson.ValidBytes(body) { + return + } + options := claudeMCPAliasOptions{secret: secret} + wantBody, wantReverseMap := remapOAuthToolNamesWithOptionsLegacy(body, options) + gotBody, gotReverseMap, ok := remapOAuthToolNamesWithBatchedEdits(body, options) + if !ok { + t.Fatal("batched remap rejected valid JSON offsets") + } + if !bytes.Equal(gotBody, wantBody) || !maps.Equal(gotReverseMap, wantReverseMap) { + t.Fatalf("batched result differs from legacy\nbody: %q\n got: %q %v\nwant: %q %v", body, gotBody, gotReverseMap, wantBody, wantReverseMap) + } + }) +}