diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 46482209ca8..1e1334dc94c 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -464,11 +464,7 @@ func (p *Server) pinnedWorkspaceMCPTools( return nil, xerrors.Errorf("list chat context resources: %w", err) } infos := workspaceMCPToolInfosFromResources(resources) - tools := make([]fantasy.AgentTool, 0, len(infos)) - for _, info := range infos { - tools = append(tools, chattool.NewWorkspaceMCPTool(info, getConn, nil)) - } - return tools, nil + return chattool.NewWorkspaceMCPTools(infos, getConn, nil), nil } type turnWorkspaceContext struct { diff --git a/coderd/x/chatd/chattool/mcpworkspace.go b/coderd/x/chatd/chattool/mcpworkspace.go index 1d2affc6d53..8f851268826 100644 --- a/coderd/x/chatd/chattool/mcpworkspace.go +++ b/coderd/x/chatd/chattool/mcpworkspace.go @@ -6,6 +6,9 @@ import ( "encoding/json" "errors" "net/http" + "regexp" + "slices" + "strconv" "strings" "charm.land/fantasy" @@ -14,27 +17,90 @@ import ( "github.com/coder/coder/v2/codersdk/workspacesdk" ) +// modelToolNameSanitizer matches characters that LLM providers reject in tool +// names. Anthropic and Bedrock require ^[a-zA-Z0-9_-]{1,128}$, and OpenAI +// enforces a 64-character cap over a similar set. A single invalid name would +// otherwise 400 the entire inference request, failing the whole turn. +var modelToolNameSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_-]`) + +// maxModelToolNameLen is the strictest provider tool-name length limit +// (OpenAI allows 64, Bedrock 128); we cap at the lower bound so names are safe +// for every provider. +const maxModelToolNameLen = 64 + // WorkspaceMCPTool wraps a single MCP tool discovered in a // workspace, proxying calls through the workspace agent // connection. It implements fantasy.AgentTool so it can be // registered alongside built-in chat tools. type WorkspaceMCPTool struct { - info fantasy.ToolInfo + info fantasy.ToolInfo + // routingName is the unsanitized "serverName__toolName" form the + // workspace agent expects: it splits on "__" to locate the server and + // calls the original tool name. info.Name is the sanitized, provider-safe + // name shown to the model, so the two can differ when the server or tool + // name contains characters outside the provider's allowed set. + routingName string getConn func(context.Context) (workspacesdk.AgentConn, error) providerOpts fantasy.ProviderOptions invalidateCache func() } -// NewWorkspaceMCPTool creates a tool wrapper from an MCPToolInfo -// discovered on a workspace agent. Each tool proxies calls back -// through the agent connection. The optional invalidateCache -// callback is invoked when CallMCPTool returns a 404 error, -// indicating that the server was removed and the chat's cached -// tool list should be dropped. +// NewWorkspaceMCPTool creates a single tool wrapper from an MCPToolInfo +// discovered on a workspace agent. Each tool proxies calls back through the +// agent connection. The optional invalidateCache callback is invoked when +// CallMCPTool returns a 404 error, indicating that the server was removed and +// the chat's cached tool list should be dropped. +// +// The model-facing name is sanitized to the provider-safe character set and +// length so a server or tool name containing a character such as "@" cannot +// produce an invalid tool name that the provider rejects. The unsanitized name +// is retained as routingName so the workspace agent can still route the call to +// the original server and tool. +// +// Prefer NewWorkspaceMCPTools when building a set of tools, because that path +// also disambiguates names that collide after sanitization. This single-tool +// constructor cannot detect collisions on its own. func NewWorkspaceMCPTool( tool workspacesdk.MCPToolInfo, getConn func(context.Context) (workspacesdk.AgentConn, error), invalidateCache func(), +) *WorkspaceMCPTool { + return buildWorkspaceMCPTool(tool, sanitizeModelToolName(tool.Name), getConn, invalidateCache) +} + +// NewWorkspaceMCPTools builds wrappers for a set of workspace MCP tools. +// Because the model-facing name is sanitized and length-capped, two distinct +// servers or tools can normalize to the same string (for example server keys +// "foo.bar" and "foo_bar" each exposing "echo", or names that share the first +// maxModelToolNameLen bytes). Duplicate names would be sent to the provider, +// which can reject the request, and the model's name-keyed dispatch would make +// one tool unreachable. To keep every tool addressable, colliding model-facing +// names are disambiguated with a numeric suffix while each tool keeps its own +// original routing name. Tools are sorted by routing name first so the suffix +// assignment is stable across turns. +func NewWorkspaceMCPTools( + infos []workspacesdk.MCPToolInfo, + getConn func(context.Context) (workspacesdk.AgentConn, error), + invalidateCache func(), +) []fantasy.AgentTool { + sorted := slices.Clone(infos) + slices.SortFunc(sorted, func(a, b workspacesdk.MCPToolInfo) int { + return strings.Compare(a.Name, b.Name) + }) + tools := make([]fantasy.AgentTool, 0, len(sorted)) + seen := make(map[string]struct{}, len(sorted)) + for _, info := range sorted { + modelName := uniqueModelToolName(sanitizeModelToolName(info.Name), seen) + tools = append(tools, buildWorkspaceMCPTool(info, modelName, getConn, invalidateCache)) + } + return tools +} + +func buildWorkspaceMCPTool( + tool workspacesdk.MCPToolInfo, + modelName string, + getConn func(context.Context) (workspacesdk.AgentConn, error), + invalidateCache func(), ) *WorkspaceMCPTool { required := tool.Required if required == nil { @@ -42,17 +108,57 @@ func NewWorkspaceMCPTool( } return &WorkspaceMCPTool{ info: fantasy.ToolInfo{ - Name: tool.Name, + Name: modelName, Description: tool.Description, Parameters: tool.Schema, Required: required, Parallel: true, }, + routingName: tool.Name, getConn: getConn, invalidateCache: invalidateCache, } } +// sanitizeModelToolName returns the provider-safe form of a workspace MCP tool +// name: characters outside [a-zA-Z0-9_-] become "_" and the result is capped +// at maxModelToolNameLen. The "__" server/tool separator survives because +// underscores are already in the allowed set. +func sanitizeModelToolName(name string) string { + sanitized := modelToolNameSanitizer.ReplaceAllString(name, "_") + if len(sanitized) > maxModelToolNameLen { + sanitized = sanitized[:maxModelToolNameLen] + } + return sanitized +} + +// uniqueModelToolName returns name when it is unused; otherwise it appends an +// incrementing "_N" suffix (starting at 2), truncating the base so the result +// stays within maxModelToolNameLen, until it finds a name absent from seen. +// The returned name is recorded in seen. +func uniqueModelToolName(name string, seen map[string]struct{}) string { + if _, ok := seen[name]; !ok { + seen[name] = struct{}{} + return name + } + for i := 2; ; i++ { + suffix := "_" + strconv.Itoa(i) + base := name + if len(base)+len(suffix) > maxModelToolNameLen { + cut := maxModelToolNameLen - len(suffix) + if cut < 0 { + cut = 0 + } + base = base[:cut] + } + candidate := base + suffix + if _, ok := seen[candidate]; !ok { + seen[candidate] = struct{}{} + return candidate + } + } +} + func (t *WorkspaceMCPTool) Info() fantasy.ToolInfo { return t.info } @@ -80,7 +186,7 @@ func (t *WorkspaceMCPTool) Run( } resp, err := conn.CallMCPTool(ctx, workspacesdk.CallMCPToolRequest{ - ToolName: t.info.Name, + ToolName: t.routingName, Arguments: args, }) if err != nil { diff --git a/coderd/x/chatd/chattool/mcpworkspace_test.go b/coderd/x/chatd/chattool/mcpworkspace_test.go index 4306509abd4..1e32d1a1394 100644 --- a/coderd/x/chatd/chattool/mcpworkspace_test.go +++ b/coderd/x/chatd/chattool/mcpworkspace_test.go @@ -3,6 +3,7 @@ package chattool_test import ( "context" "net/http" + "strings" "sync/atomic" "testing" @@ -153,3 +154,127 @@ func TestWorkspaceMCPTool_InvalidateOn404(t *testing.T) { assert.True(t, resp.IsError) }) } + +func TestWorkspaceMCPTool_SanitizesModelNameKeepsRoutingName(t *testing.T) { + t.Parallel() + + t.Run("InvalidCharsSanitizedForModelOriginalForRouting", func(t *testing.T) { + t.Parallel() + + var gotToolName string + tool := chattool.NewWorkspaceMCPTool( + workspacesdk.MCPToolInfo{ + // "@" is outside the provider's allowed tool-name set; the + // model must never see it or the whole request is rejected. + Name: "weather@home__get_forecast", + Description: "test tool", + }, + func(_ context.Context) (workspacesdk.AgentConn, error) { + return &fakeAgentConn{ + callMCPToolFunc: func(_ context.Context, req workspacesdk.CallMCPToolRequest) (workspacesdk.CallMCPToolResponse, error) { + gotToolName = req.ToolName + return workspacesdk.CallMCPToolResponse{ + Content: []workspacesdk.MCPToolContent{{Type: "text", Text: "ok"}}, + }, nil + }, + }, nil + }, + nil, + ) + + // The model-facing name is sanitized to the provider-safe set. + assert.Equal(t, "weather_home__get_forecast", tool.Info().Name) + + resp, err := tool.Run(context.Background(), fantasy.ToolCall{}) + require.NoError(t, err) + assert.False(t, resp.IsError) + // The agent receives the original name so it can route the call to + // the correct server and original tool. + assert.Equal(t, "weather@home__get_forecast", gotToolName) + }) + + t.Run("ValidNameUnchanged", func(t *testing.T) { + t.Parallel() + + tool := chattool.NewWorkspaceMCPTool( + workspacesdk.MCPToolInfo{ + Name: "github__create_issue", + Description: "test tool", + }, + func(_ context.Context) (workspacesdk.AgentConn, error) { + return &fakeAgentConn{ + callMCPToolFunc: func(_ context.Context, _ workspacesdk.CallMCPToolRequest) (workspacesdk.CallMCPToolResponse, error) { + return workspacesdk.CallMCPToolResponse{}, nil + }, + }, nil + }, + nil, + ) + + // A name already within the allowed set is left untouched. + assert.Equal(t, "github__create_issue", tool.Info().Name) + }) + + t.Run("LongNameTruncatedForModel", func(t *testing.T) { + t.Parallel() + + // A name longer than the provider limit is truncated. "srv__" plus a + // 64-char tool name exceeds the 64-char cap. + longName := "srv__" + strings.Repeat("a", 64) + tool := chattool.NewWorkspaceMCPTool( + workspacesdk.MCPToolInfo{ + Name: longName, + Description: "test tool", + }, + func(_ context.Context) (workspacesdk.AgentConn, error) { + return &fakeAgentConn{ + callMCPToolFunc: func(_ context.Context, _ workspacesdk.CallMCPToolRequest) (workspacesdk.CallMCPToolResponse, error) { + return workspacesdk.CallMCPToolResponse{}, nil + }, + }, nil + }, + nil, + ) + + // The model-facing name is capped at the strictest provider limit. + assert.LessOrEqual(t, len(tool.Info().Name), 64) + }) +} + +func TestNewWorkspaceMCPTools_DisambiguatesCollidingNames(t *testing.T) { + t.Parallel() + + var routed []string + getConn := func(_ context.Context) (workspacesdk.AgentConn, error) { + return &fakeAgentConn{ + callMCPToolFunc: func(_ context.Context, req workspacesdk.CallMCPToolRequest) (workspacesdk.CallMCPToolResponse, error) { + routed = append(routed, req.ToolName) + return workspacesdk.CallMCPToolResponse{}, nil + }, + }, nil + } + + // Both names sanitize to "foo_bar__echo"; the set builder must keep them + // distinct for the model while routing each to its own original name. + infos := []workspacesdk.MCPToolInfo{ + {Name: "foo.bar__echo"}, + {Name: "foo_bar__echo"}, + } + + tools := chattool.NewWorkspaceMCPTools(infos, getConn, nil) + require.Len(t, tools, 2) + + names := []string{tools[0].Info().Name, tools[1].Info().Name} + assert.NotEqual(t, names[0], names[1], + "colliding model-facing names must be disambiguated") + assert.ElementsMatch(t, + []string{"foo_bar__echo", "foo_bar__echo_2"}, names) + + // Each tool routes to its own original (unsanitized) name. + for _, tl := range tools { + _, err := tl.Run(context.Background(), fantasy.ToolCall{}) + require.NoError(t, err) + } + assert.ElementsMatch(t, + []string{"foo.bar__echo", "foo_bar__echo"}, routed) +}