History

fix(claude): recover OAuth tool aliases for repeated prefixes and malformed IDs

- Normalize repeated `mcp__<server>__` alias prefixes during reverse remapping to resolve stacked aliases. - Add a semantic-suffix fallback when parsing fails, allowing unambiguous recovery from malformed tool IDs.

Luis Pater committed 3 months ago

Commit f2d272d1 parent 7ea9c672 files · +106 −19Browse files

Unified · Split

2 files
  1. internal/runtime/executor/claude_executor_request.go +27 −10
  2. internal/runtime/executor/claude_executor_request_remap_test.go +79 −9
internal/runtime/executor/claude_executor_request.go+27−10
@@ -1586,9 +1586,17 @@ func (resolver claudeMCPAliasResolver) resolve(name string) (string, bool, error · 1585 unmodified lines
15861586 return "", false, nil
15871587 }
15881588
1589− repeatedServerPrefix := "mcp__" + server + "__" + server + "__"
1590− if suffix, repeatedServer := strings.CutPrefix(name, repeatedServerPrefix); repeatedServer {
1591− if original, exact := resolver.exact["mcp__"+server+"__"+suffix]; exact {
1589+ canonicalServerPrefix := "mcp__" + server + "__"
1590+ normalizedName := name
1591+ suffix := strings.TrimPrefix(name, canonicalServerPrefix)
1592+ for {
1593+ strippedSuffix, repeatedServer := strings.CutPrefix(suffix, server+"__")
1594+ if !repeatedServer {
1595+ break
1596+ }
1597+ suffix = strippedSuffix
1598+ normalizedName = canonicalServerPrefix + suffix
1599+ if original, exact := resolver.exact[normalizedName]; exact {
15921600 return original, true, nil
15931601 }
15941602 }
@@ -1608,21 +1616,30 @@ func (resolver claudeMCPAliasResolver) resolve(name string) (string, bool, error · 13 unmodified lines
16081616 return "", false, fmt.Errorf("cannot restore Claude OAuth MCP tool alias %q: matched multiple declared aliases", name)
16091617 }
16101618
1611− parts, ok := parseClaudeMCPAlias(name)
1612− if ok {
1619+ parts, validAlias := parseClaudeMCPAlias(normalizedName)
1620+ if validAlias {
16131621 for _, entry := range resolver.aliases {
16141622 if entry.parts.server == parts.server && entry.parts.semantic == parts.semantic {
16151623 matchedOriginal = entry.original
16161624 matchCount++
16171625 }
16181626 }
1619− if matchCount == 1 {
1620− return matchedOriginal, true, nil
1621− }
1622− if matchCount > 1 {
1623− return "", false, fmt.Errorf("cannot restore Claude OAuth MCP tool alias %q: semantic suffix matches multiple declared tools", name)
1627+ } else {
1628+ // Keep generated aliases strict while allowing a malformed response tool ID
1629+ // to recover only when its request-local semantic suffix is unambiguous.
1630+ for _, entry := range resolver.aliases {
1631+ if entry.parts.server == server && strings.HasSuffix(normalizedName, "_"+entry.parts.semantic) {
1632+ matchedOriginal = entry.original
1633+ matchCount++
1634+ }
16241635 }
16251636 }
1637+ if matchCount == 1 {
1638+ return matchedOriginal, true, nil
1639+ }
1640+ if matchCount > 1 {
1641+ return "", false, fmt.Errorf("cannot restore Claude OAuth MCP tool alias %q: semantic suffix matches multiple declared tools", name)
1642+ }
16261643
16271644 return "", false, fmt.Errorf("cannot restore Claude OAuth MCP tool alias %q: no unique request-local match", name)
16281645 }
internal/runtime/executor/claude_executor_request_remap_test.go+79−9
@@ -163,18 +163,83 @@ func TestReverseRemapOAuthToolNamesRecoversMangledAliases(t *testing.T) { · 162 unmodified lines
163163 }
164164 }
165165
166−func TestReverseRemapOAuthToolNamesRecoversRepeatedServerAlias(t *testing.T) {
166+func TestReverseRemapOAuthToolNamesRecoversRepeatedServerAliases(t *testing.T) {
167167 const alias = "mcp__hmzqrngkulqv__xuo7jlxlpzee_Bash"
168− reverseMap := map[string]string{alias: "Bash"}
169− responseAlias := "mcp__hmzqrngkulqv__hmzqrngkulqv__xuo7jlxlpzee_Bash"
170− response := []byte(fmt.Sprintf(`{"content":[{"type":"tool_use","id":"toolu_1","name":%q,"input":{}}]}`, responseAlias))
168+ reverseMap := map[string]string{
169+ alias: "Bash",
170+ "mcp__hmzqrngkulqv__aaaaaaaaaaaa_Bash": "OtherBash",
171+ }
172+ tests := []struct {
173+ name string
174+ responseAlias string
175+ }{
176+ {
177+ name: "single repetition",
178+ responseAlias: "mcp__hmzqrngkulqv__hmzqrngkulqv__xuo7jlxlpzee_Bash",
179+ },
180+ {
181+ name: "multiple repetitions",
182+ responseAlias: "mcp__hmzqrngkulqv__hmzqrngkulqv__hmzqrngkulqv__xuo7jlxlpzee_Bash",
183+ },
184+ }
171185
172− restored, errReverse := reverseRemapOAuthToolNames(response, reverseMap)
173− if errReverse != nil {
174− t.Fatalf("reverseRemapOAuthToolNames() error = %v", errReverse)
186+ for _, test := range tests {
187+ t.Run(test.name, func(t *testing.T) {
188+ response := []byte(fmt.Sprintf(`{"content":[{"type":"tool_use","id":"toolu_1","name":%q,"input":{}}]}`, test.responseAlias))
189+ restored, errReverse := reverseRemapOAuthToolNames(response, reverseMap)
190+ if errReverse != nil {
191+ t.Fatalf("reverseRemapOAuthToolNames() error = %v", errReverse)
192+ }
193+ if got := gjson.GetBytes(restored, "content.0.name").String(); got != "Bash" {
194+ t.Fatalf("repeated server alias restored to %q, want Bash", got)
195+ }
196+
197+ line := []byte(fmt.Sprintf(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":%q,"input":{}}}`, test.responseAlias))
198+ restoredLine, errStream := reverseRemapOAuthToolNamesFromStreamLine(line, reverseMap)
199+ if errStream != nil {
200+ t.Fatalf("reverseRemapOAuthToolNamesFromStreamLine() error = %v", errStream)
201+ }
202+ if got := gjson.GetBytes(helps.JSONPayload(restoredLine), "content_block.name").String(); got != "Bash" {
203+ t.Fatalf("stream repeated server alias restored to %q, want Bash", got)
204+ }
205+ })
175206 }
176− if got := gjson.GetBytes(restored, "content.0.name").String(); got != "Bash" {
177− t.Fatalf("repeated server alias restored to %q, want Bash", got)
207+}
208+
209+func TestReverseRemapOAuthToolNamesRecoversMalformedToolIDBySemanticSuffix(t *testing.T) {
210+ const alias = "mcp__hmzqrngkulqv__xuo7jlxlpzee_Bash"
211+ reverseMap := map[string]string{alias: "Bash"}
212+ tests := []struct {
213+ name string
214+ responseAlias string
215+ }{
216+ {name: "short tool ID", responseAlias: "mcp__hmzqrngkulqv__xuo7jlxlpze_Bash"},
217+ {name: "long tool ID", responseAlias: "mcp__hmzqrngkulqv__xuo7jlxlpzeea_Bash"},
218+ {name: "invalid base32 tool ID", responseAlias: "mcp__hmzqrngkulqv__xuo7jlxlpze0_Bash"},
219+ {name: "substituted base32 tool ID", responseAlias: "mcp__hmzqrngkulqv__auo7jlxlpzee_Bash"},
220+ {name: "repeated server and short tool ID", responseAlias: "mcp__hmzqrngkulqv__hmzqrngkulqv__xuo7jlxlpze_Bash"},
221+ }
222+
223+ for _, test := range tests {
224+ t.Run(test.name, func(t *testing.T) {
225+ response := []byte(fmt.Sprintf(`{"content":[{"type":"tool_use","id":"toolu_1","name":%q,"input":{}}]}`, test.responseAlias))
226+ restored, errReverse := reverseRemapOAuthToolNames(response, reverseMap)
227+ if errReverse != nil {
228+ t.Fatalf("reverseRemapOAuthToolNames() error = %v", errReverse)
229+ }
230+ if got := gjson.GetBytes(restored, "content.0.name").String(); got != "Bash" {
231+ t.Fatalf("malformed tool ID alias restored to %q, want Bash", got)
232+ }
233+
234+ line := []byte(fmt.Sprintf(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":%q,"input":{}}}`, test.responseAlias))
235+ restoredLine, errStream := reverseRemapOAuthToolNamesFromStreamLine(line, reverseMap)
236+ if errStream != nil {
237+ t.Fatalf("reverseRemapOAuthToolNamesFromStreamLine() error = %v", errStream)
238+ }
239+ if got := gjson.GetBytes(helps.JSONPayload(restoredLine), "content_block.name").String(); got != "Bash" {
240+ t.Fatalf("stream malformed tool ID alias restored to %q, want Bash", got)
241+ }
242+ })
178243 }
179244 }
180245
@@ -209,6 +274,11 @@ func TestReverseRemapOAuthToolNamesRejectsUnsafeMangledAliases(t *testing.T) { · 28 unmodified lines
209274 alias: "mcp__" + firstParts.server + "__" + unknownToolID + "_" + firstParts.semantic,
210275 wantError: "semantic suffix matches multiple declared tools",
211276 },
277+ {
278+ name: "ambiguous semantic suffix with malformed tool ID",
279+ alias: "mcp__" + firstParts.server + "__" + unknownToolID[:len(unknownToolID)-1] + "_" + firstParts.semantic,
280+ wantError: "semantic suffix matches multiple declared tools",
281+ },
212282 {
213283 name: "unrecoverable semantic suffix",
214284 alias: "mcp__" + firstParts.server + "__" + unknownToolID + "_missing_tool",