Skip to content

Commit 5ab9afa

Browse files
luispaterluispater
committed
fix(executor): handle OAuth tool name remapping with rename detection and add tests
Closes: router-for-me#2656
1 parent 65ce863 commit 5ab9afa

2 files changed

Lines changed: 96 additions & 46 deletions

File tree

‎internal/runtime/executor/claude_executor.go‎

Lines changed: 54 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,9 @@ var oauthToolRenameMap = map[string]string{
5757
"glob": "Glob",
5858
"grep": "Grep",
5959
"task": "Task",
60-
"webfetch": "WebFetch",
61-
"todowrite": "TodoWrite",
62-
"question": "Question",
60+
"webfetch": "WebFetch",
61+
"todowrite": "TodoWrite",
62+
"question": "Question",
6363
"skill": "Skill",
6464
"ls": "LS",
6565
"todoread": "TodoRead",
@@ -192,14 +192,15 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
192192
bodyForTranslation := body
193193
bodyForUpstream := body
194194
oauthToken := isClaudeOAuthToken(apiKey)
195+
oauthToolNamesRemapped := false
195196
if oauthToken && !auth.ToolPrefixDisabled() {
196197
bodyForUpstream = applyClaudeToolPrefix(body, claudeToolPrefix)
197198
}
198199
// Remap third-party tool names to Claude Code equivalents and remove
199200
// tools without official counterparts. This prevents Anthropic from
200201
// fingerprinting the request as third-party via tool naming patterns.
201202
if oauthToken {
202-
bodyForUpstream = remapOAuthToolNames(bodyForUpstream)
203+
bodyForUpstream, oauthToolNamesRemapped = remapOAuthToolNames(bodyForUpstream)
203204
}
204205
// Enable cch signing by default for OAuth tokens (not just experimental flag).
205206
// Claude Code always computes cch; missing or invalid cch is a detectable fingerprint.
@@ -297,7 +298,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
297298
data = stripClaudeToolPrefixFromResponse(data, claudeToolPrefix)
298299
}
299300
// Reverse the OAuth tool name remap so the downstream client sees original names.
300-
if isClaudeOAuthToken(apiKey) {
301+
if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
301302
data = reverseRemapOAuthToolNames(data)
302303
}
303304
var param any
@@ -373,14 +374,15 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
373374
bodyForTranslation := body
374375
bodyForUpstream := body
375376
oauthToken := isClaudeOAuthToken(apiKey)
377+
oauthToolNamesRemapped := false
376378
if oauthToken && !auth.ToolPrefixDisabled() {
377379
bodyForUpstream = applyClaudeToolPrefix(body, claudeToolPrefix)
378380
}
379381
// Remap third-party tool names to Claude Code equivalents and remove
380382
// tools without official counterparts. This prevents Anthropic from
381383
// fingerprinting the request as third-party via tool naming patterns.
382384
if oauthToken {
383-
bodyForUpstream = remapOAuthToolNames(bodyForUpstream)
385+
bodyForUpstream, oauthToolNamesRemapped = remapOAuthToolNames(bodyForUpstream)
384386
}
385387
// Enable cch signing by default for OAuth tokens (not just experimental flag).
386388
if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) {
@@ -474,7 +476,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
474476
if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {
475477
line = stripClaudeToolPrefixFromStreamLine(line, claudeToolPrefix)
476478
}
477-
if isClaudeOAuthToken(apiKey) {
479+
if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
478480
line = reverseRemapOAuthToolNamesFromStreamLine(line)
479481
}
480482
// Forward the line as-is to preserve SSE format
@@ -504,7 +506,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
504506
if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {
505507
line = stripClaudeToolPrefixFromStreamLine(line, claudeToolPrefix)
506508
}
507-
if isClaudeOAuthToken(apiKey) {
509+
if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
508510
line = reverseRemapOAuthToolNamesFromStreamLine(line)
509511
}
510512
chunks := sdktranslator.TranslateStream(
@@ -561,7 +563,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut
561563
}
562564
// Remap tool names for OAuth token requests to avoid third-party fingerprinting.
563565
if isClaudeOAuthToken(apiKey) {
564-
body = remapOAuthToolNames(body)
566+
body, _ = remapOAuthToolNames(body)
565567
}
566568

567569
url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL)
@@ -1018,7 +1020,8 @@ func isClaudeOAuthToken(apiKey string) bool {
10181020
// It operates on: tools[].name, tool_choice.name, and all tool_use/tool_reference
10191021
// references in messages. Removed tools' corresponding tool_result blocks are preserved
10201022
// (they just become orphaned, which is safe for Claude).
1021-
func remapOAuthToolNames(body []byte) []byte {
1023+
func remapOAuthToolNames(body []byte) ([]byte, bool) {
1024+
renamed := false
10221025
// 1. Rewrite tools array in a single pass (if present).
10231026
// IMPORTANT: do not mutate names first and then rebuild from an older gjson
10241027
// snapshot. gjson results are snapshots of the original bytes; rebuilding from a
@@ -1027,42 +1030,43 @@ func remapOAuthToolNames(body []byte) []byte {
10271030
tools := gjson.GetBytes(body, "tools")
10281031
if tools.Exists() && tools.IsArray() {
10291032

1030-
var toolsJSON strings.Builder
1031-
toolsJSON.WriteByte('[')
1032-
toolCount := 0
1033-
tools.ForEach(func(_, tool gjson.Result) bool {
1034-
// Keep Anthropic built-in tools (web_search, code_execution, etc.) unchanged.
1035-
if tool.Get("type").Exists() && tool.Get("type").String() != "" {
1036-
if toolCount > 0 {
1037-
toolsJSON.WriteByte(',')
1033+
var toolsJSON strings.Builder
1034+
toolsJSON.WriteByte('[')
1035+
toolCount := 0
1036+
tools.ForEach(func(_, tool gjson.Result) bool {
1037+
// Keep Anthropic built-in tools (web_search, code_execution, etc.) unchanged.
1038+
if tool.Get("type").Exists() && tool.Get("type").String() != "" {
1039+
if toolCount > 0 {
1040+
toolsJSON.WriteByte(',')
1041+
}
1042+
toolsJSON.WriteString(tool.Raw)
1043+
toolCount++
1044+
return true
10381045
}
1039-
toolsJSON.WriteString(tool.Raw)
1040-
toolCount++
1041-
return true
1042-
}
10431046

1044-
name := tool.Get("name").String()
1045-
if oauthToolsToRemove[name] {
1046-
return true
1047-
}
1047+
name := tool.Get("name").String()
1048+
if oauthToolsToRemove[name] {
1049+
return true
1050+
}
10481051

1049-
toolJSON := tool.Raw
1050-
if newName, ok := oauthToolRenameMap[name]; ok {
1051-
updatedTool, err := sjson.Set(toolJSON, "name", newName)
1052-
if err == nil {
1053-
toolJSON = updatedTool
1052+
toolJSON := tool.Raw
1053+
if newName, ok := oauthToolRenameMap[name]; ok && newName != name {
1054+
updatedTool, err := sjson.Set(toolJSON, "name", newName)
1055+
if err == nil {
1056+
toolJSON = updatedTool
1057+
renamed = true
1058+
}
10541059
}
1055-
}
10561060

1057-
if toolCount > 0 {
1058-
toolsJSON.WriteByte(',')
1059-
}
1060-
toolsJSON.WriteString(toolJSON)
1061-
toolCount++
1062-
return true
1063-
})
1064-
toolsJSON.WriteByte(']')
1065-
body, _ = sjson.SetRawBytes(body, "tools", []byte(toolsJSON.String()))
1061+
if toolCount > 0 {
1062+
toolsJSON.WriteByte(',')
1063+
}
1064+
toolsJSON.WriteString(toolJSON)
1065+
toolCount++
1066+
return true
1067+
})
1068+
toolsJSON.WriteByte(']')
1069+
body, _ = sjson.SetRawBytes(body, "tools", []byte(toolsJSON.String()))
10661070
}
10671071

10681072
// 2. Rename tool_choice if it references a known tool
@@ -1073,8 +1077,9 @@ func remapOAuthToolNames(body []byte) []byte {
10731077
// The chosen tool was removed from the tools array, so drop tool_choice to
10741078
// keep the payload internally consistent and fall back to normal auto tool use.
10751079
body, _ = sjson.DeleteBytes(body, "tool_choice")
1076-
} else if newName, ok := oauthToolRenameMap[tcName]; ok {
1080+
} else if newName, ok := oauthToolRenameMap[tcName]; ok && newName != tcName {
10771081
body, _ = sjson.SetBytes(body, "tool_choice.name", newName)
1082+
renamed = true
10781083
}
10791084
}
10801085

@@ -1091,15 +1096,17 @@ func remapOAuthToolNames(body []byte) []byte {
10911096
switch partType {
10921097
case "tool_use":
10931098
name := part.Get("name").String()
1094-
if newName, ok := oauthToolRenameMap[name]; ok {
1099+
if newName, ok := oauthToolRenameMap[name]; ok && newName != name {
10951100
path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int())
10961101
body, _ = sjson.SetBytes(body, path, newName)
1102+
renamed = true
10971103
}
10981104
case "tool_reference":
10991105
toolName := part.Get("tool_name").String()
1100-
if newName, ok := oauthToolRenameMap[toolName]; ok {
1106+
if newName, ok := oauthToolRenameMap[toolName]; ok && newName != toolName {
11011107
path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int())
11021108
body, _ = sjson.SetBytes(body, path, newName)
1109+
renamed = true
11031110
}
11041111
case "tool_result":
11051112
// Handle nested tool_reference blocks inside tool_result.content[]
@@ -1110,9 +1117,10 @@ func remapOAuthToolNames(body []byte) []byte {
11101117
nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool {
11111118
if nestedPart.Get("type").String() == "tool_reference" {
11121119
nestedToolName := nestedPart.Get("tool_name").String()
1113-
if newName, ok := oauthToolRenameMap[nestedToolName]; ok {
1120+
if newName, ok := oauthToolRenameMap[nestedToolName]; ok && newName != nestedToolName {
11141121
nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int())
11151122
body, _ = sjson.SetBytes(body, nestedPath, newName)
1123+
renamed = true
11161124
}
11171125
}
11181126
return true
@@ -1125,7 +1133,7 @@ func remapOAuthToolNames(body []byte) []byte {
11251133
})
11261134
}
11271135

1128-
return body
1136+
return body, renamed
11291137
}
11301138

11311139
// reverseRemapOAuthToolNames reverses the tool name mapping for non-stream responses.

‎internal/runtime/executor/claude_executor_test.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1949,3 +1949,45 @@ func TestNormalizeClaudeTemperatureForThinking_AfterForcedToolChoiceKeepsOrigina
19491949
t.Fatalf("temperature = %v, want 0", got)
19501950
}
19511951
}
1952+
1953+
func TestRemapOAuthToolNames_TitleCase_NoReverseNeeded(t *testing.T) {
1954+
body := []byte(`{"tools":[{"name":"Bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
1955+
1956+
out, renamed := remapOAuthToolNames(body)
1957+
if renamed {
1958+
t.Fatalf("renamed = true, want false")
1959+
}
1960+
if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {
1961+
t.Fatalf("tools.0.name = %q, want %q", got, "Bash")
1962+
}
1963+
1964+
resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
1965+
reversed := resp
1966+
if renamed {
1967+
reversed = reverseRemapOAuthToolNames(resp)
1968+
}
1969+
if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "Bash" {
1970+
t.Fatalf("content.0.name = %q, want %q", got, "Bash")
1971+
}
1972+
}
1973+
1974+
func TestRemapOAuthToolNames_Lowercase_ReverseApplied(t *testing.T) {
1975+
body := []byte(`{"tools":[{"name":"bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
1976+
1977+
out, renamed := remapOAuthToolNames(body)
1978+
if !renamed {
1979+
t.Fatalf("renamed = false, want true")
1980+
}
1981+
if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {
1982+
t.Fatalf("tools.0.name = %q, want %q", got, "Bash")
1983+
}
1984+
1985+
resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
1986+
reversed := resp
1987+
if renamed {
1988+
reversed = reverseRemapOAuthToolNames(resp)
1989+
}
1990+
if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "bash" {
1991+
t.Fatalf("content.0.name = %q, want %q", got, "bash")
1992+
}
1993+
}

0 commit comments

Comments
 (0)