From a2b680ab2979064727adcf8a3b77d0ddf21d849b Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Fri, 19 Jun 2026 11:48:30 +0200 Subject: [PATCH] fix: sanitize MCP tool names to satisfy LLM provider constraints (#26539) Fixes #26325 --- aibridge/mcp/proxy_streamable_http.go | 9 ++ aibridge/mcp/tool.go | 30 ++++- aibridge/mcp/tool_test.go | 128 +++++++++++++++++++++ coderd/x/chatd/mcpclient/mcpclient.go | 37 +++++- coderd/x/chatd/mcpclient/mcpclient_test.go | 64 +++++++++++ 5 files changed, 264 insertions(+), 4 deletions(-) create mode 100644 aibridge/mcp/tool_test.go diff --git a/aibridge/mcp/proxy_streamable_http.go b/aibridge/mcp/proxy_streamable_http.go index 108a710d19..8d9e3583c1 100644 --- a/aibridge/mcp/proxy_streamable_http.go +++ b/aibridge/mcp/proxy_streamable_http.go @@ -156,6 +156,15 @@ func (p *StreamableHTTPServerProxy) fetchTools(ctx context.Context) (_ map[strin out := make(map[string]*Tool, len(tools.Tools)) for _, tool := range tools.Tools { encodedID := EncodeToolID(p.serverName, tool.Name) + if existing, ok := out[encodedID]; ok { + p.logger.Warn(ctx, + "duplicate tool ID after sanitization; previous tool will be unreachable", + slog.F("tool_id", encodedID), + slog.F("new_tool", tool.Name), + slog.F("replaced_tool", existing.Name), + slog.F("server", p.serverName), + ) + } out[encodedID] = &Tool{ Client: p.client, ID: encodedID, diff --git a/aibridge/mcp/tool.go b/aibridge/mcp/tool.go index 8fbca9d224..bb13d626ef 100644 --- a/aibridge/mcp/tool.go +++ b/aibridge/mcp/tool.go @@ -20,8 +20,28 @@ const ( maxSpanInputAttrLen = 100 // truncates tool.Call span input attribute to first `maxSpanInputAttrLen` letters injectedToolPrefix = "bmcp" // "bridged MCP" injectedToolDelimiter = "_" + + // MaxToolNameLen is the strictest provider limit for tool names. + // OpenAI allows 64 characters; Bedrock allows 128. We use the + // lower bound so names are safe for every provider. + MaxToolNameLen = 64 ) +// toolNameSanitizer replaces characters that violate LLM provider tool +// name constraints. Bedrock requires ^[a-zA-Z0-9_-]{1,128}$ and OpenAI +// enforces a 64-character limit with a similar character set. Characters +// outside [a-zA-Z0-9_-] are replaced with "_" so a single invalid +// server or tool name cannot 400 the entire inference request. +var toolNameSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_-]`) + +// SanitizeToolName replaces characters outside [a-zA-Z0-9_-] with +// underscores so the resulting name is accepted by LLM providers. +// Callers that assemble a full tool name from multiple components +// should truncate the final result to MaxToolNameLen. +func SanitizeToolName(name string) string { + return toolNameSanitizer.ReplaceAllString(name, "_") +} + // ToolCaller is the narrowest interface which describes the behavior required from [mcp.Client], // which will normally be passed into [Tool] for interaction with an MCP server. // TODO: don't expose github.com/mark3labs/mcp-go outside this package. @@ -110,10 +130,14 @@ func EncodeToolID(server, tool string) string { var sb strings.Builder _, _ = sb.WriteString(injectedToolPrefix) _, _ = sb.WriteString(injectedToolDelimiter) - _, _ = sb.WriteString(server) + _, _ = sb.WriteString(SanitizeToolName(server)) _, _ = sb.WriteString(injectedToolDelimiter) - _, _ = sb.WriteString(tool) - return sb.String() + _, _ = sb.WriteString(SanitizeToolName(tool)) + id := sb.String() + if len(id) > MaxToolNameLen { + id = id[:MaxToolNameLen] + } + return id } // FilterAllowedTools filters tools based on the given allow/denylists. diff --git a/aibridge/mcp/tool_test.go b/aibridge/mcp/tool_test.go new file mode 100644 index 0000000000..4f41cec6a6 --- /dev/null +++ b/aibridge/mcp/tool_test.go @@ -0,0 +1,128 @@ +package mcp_test + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/aibridge/mcp" +) + +func TestSanitizeToolName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + expected string + }{ + { + name: "AlreadyValid", + input: "my_tool-name123", + expected: "my_tool-name123", + }, + { + name: "DotsReplaced", + input: "awslabs.aws-documentation-mcp-server", + expected: "awslabs_aws-documentation-mcp-server", + }, + { + name: "MultipleDots", + input: "com.example.tool.v2", + expected: "com_example_tool_v2", + }, + { + name: "Spaces", + input: "my tool name", + expected: "my_tool_name", + }, + { + name: "SpecialCharacters", + input: "tool@v2#special!", + expected: "tool_v2_special_", + }, + { + name: "Empty", + input: "", + expected: "", + }, + { + name: "AllInvalid", + input: "...", + expected: "___", + }, + { + name: "Slashes", + input: "org/repo/tool", + expected: "org_repo_tool", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := mcp.SanitizeToolName(tt.input) + assert.Equal(t, tt.expected, got) + }) + } +} + +func TestEncodeToolID_SanitizesComponents(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + server string + tool string + expected string + }{ + { + name: "ValidNames", + server: "my-server", + tool: "my_tool", + expected: "bmcp_my-server_my_tool", + }, + { + name: "DottedServerName", + server: "awslabs.aws-documentation-mcp-server", + tool: "read_documentation", + expected: "bmcp_awslabs_aws-documentation-mcp-server_read_documentation", + }, + { + name: "DottedToolName", + server: "server", + tool: "com.example.action", + expected: "bmcp_server_com_example_action", + }, + { + name: "BothDotted", + server: "org.server", + tool: "ns.action", + expected: "bmcp_org_server_ns_action", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := mcp.EncodeToolID(tt.server, tt.tool) + assert.Equal(t, tt.expected, got) + }) + } +} + +func TestEncodeToolID_TruncatesLongNames(t *testing.T) { + t.Parallel() + + // "bmcp_" prefix = 5 chars, "_" delimiter = 1 char, so + // server + tool budget is MaxToolNameLen - 6. + longServer := strings.Repeat("a", 40) + longTool := strings.Repeat("b", 40) + + id := mcp.EncodeToolID(longServer, longTool) + require.LessOrEqual(t, len(id), mcp.MaxToolNameLen, + "encoded ID must not exceed MaxToolNameLen") + assert.True(t, strings.HasPrefix(id, "bmcp_")) +} diff --git a/coderd/x/chatd/mcpclient/mcpclient.go b/coderd/x/chatd/mcpclient/mcpclient.go index cb7e0322c2..0ff65db0aa 100644 --- a/coderd/x/chatd/mcpclient/mcpclient.go +++ b/coderd/x/chatd/mcpclient/mcpclient.go @@ -24,6 +24,7 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" + aidmcp "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/coderd/database" ) @@ -39,6 +40,15 @@ import ( // directly when calling the remote server. const toolNameSep = "__" +// truncateToolName caps the assembled tool name at MaxToolNameLen so +// it fits within provider limits (e.g. OpenAI 64, Bedrock 128). +func truncateToolName(name string) string { + if len(name) > aidmcp.MaxToolNameLen { + return name[:aidmcp.MaxToolNameLen] + } + return name +} + // connectTimeout bounds how long we wait for a single MCP server // to start its transport and complete initialization. Servers that // take longer are skipped so one slow server cannot block the @@ -163,6 +173,31 @@ func ConnectAll( ) }) + // Warn about name collisions that may result from sanitization + // and truncation. When two tools resolve to the same name, the + // LLM tool-call dispatch map keeps only one, so the other + // becomes silently unreachable. + for i := 1; i < len(tools); i++ { + if tools[i-1].Info().Name == tools[i].Info().Name { + prevTool, ok := tools[i-1].(MCPToolIdentifier) + if !ok { + continue + } + currTool, ok := tools[i].(MCPToolIdentifier) + if !ok { + continue + } + if prevTool.MCPServerConfigID() != currTool.MCPServerConfigID() { + logger.Warn(ctx, + "duplicate tool name after sanitization; one tool will be unreachable", + slog.F("tool_name", tools[i].Info().Name), + slog.F("prev_config_id", prevTool.MCPServerConfigID()), + slog.F("curr_config_id", currTool.MCPServerConfigID()), + ) + } + } + } + return tools, cleanup } @@ -520,7 +555,7 @@ func newMCPTool( ) *mcpToolWrapper { return &mcpToolWrapper{ configID: configID, - prefixedName: serverSlug + toolNameSep + tool.Name, + prefixedName: truncateToolName(aidmcp.SanitizeToolName(serverSlug) + toolNameSep + aidmcp.SanitizeToolName(tool.Name)), originalName: tool.Name, description: tool.Description, parameters: tool.InputSchema.Properties, diff --git a/coderd/x/chatd/mcpclient/mcpclient_test.go b/coderd/x/chatd/mcpclient/mcpclient_test.go index d91788fd2f..d6451e30a7 100644 --- a/coderd/x/chatd/mcpclient/mcpclient_test.go +++ b/coderd/x/chatd/mcpclient/mcpclient_test.go @@ -6,6 +6,7 @@ import ( "encoding/base64" "encoding/json" "net/http/httptest" + "strings" "sync" "testing" "time" @@ -113,6 +114,69 @@ func TestConnectAll_DiscoverTools(t *testing.T) { assert.Equal(t, "Echoes the input", echoInfo.Description) } +func TestConnectAll_SanitizesDottedSlug(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + ts := newTestMCPServer(t, echoTool()) + + // Use a dotted slug like awslabs.* MCP servers ship with. + // Dots violate Bedrock's tool name pattern ^[a-zA-Z0-9_-]{1,128}$. + cfg := makeConfig("awslabs.aws-documentation-mcp-server", ts.URL) + tools, cleanup := mcpclient.ConnectAll(ctx, logger, []database.MCPServerConfig{cfg}, nil, uuid.Nil, nil, nil) + t.Cleanup(cleanup) + + require.Len(t, tools, 1) + + // Dots in the slug must be replaced with underscores. + names := toolNames(tools) + assert.Equal(t, []string{"awslabs_aws-documentation-mcp-server__echo"}, names) + + // The tool should still be callable; the original name is + // used when contacting the remote MCP server. + resp, err := tools[0].Run(ctx, fantasy.ToolCall{ + ID: "call-1", + Name: "awslabs_aws-documentation-mcp-server__echo", + Input: `{"input":"hello"}`, + }) + require.NoError(t, err) + assert.Equal(t, "echo: hello", resp.Content) +} + +func TestConnectAll_TruncationCollisionWarning(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + // Two servers whose slugs differ only in a trailing suffix. + // After sanitization + truncation to 64 chars, both produce + // the same prefixed tool name, triggering a collision warning. + // slug (65) + "__" (2) + "echo" (4) = 71 chars; truncated to + // 64 chops the suffix and tool name entirely. + base := strings.Repeat("a", 64) + slug1 := base + "x" + slug2 := base + "y" + + ts := newTestMCPServer(t, echoTool()) + + cfg1 := makeConfig(slug1, ts.URL) + cfg2 := makeConfig(slug2, ts.URL) + + tools, cleanup := mcpclient.ConnectAll( + ctx, logger, + []database.MCPServerConfig{cfg1, cfg2}, + nil, uuid.Nil, nil, nil, + ) + t.Cleanup(cleanup) + + // Both tools should be present (the caller decides policy), + // but their names collide after truncation. + require.Len(t, tools, 2) + assert.Equal(t, tools[0].Info().Name, tools[1].Info().Name, + "truncated names should collide") +} + func TestConnectAll_CallTool(t *testing.T) { t.Parallel() ctx := context.Background()