diff --git a/agent/x/agentmcp/manager.go b/agent/x/agentmcp/manager.go index 363f61a3dfb..80197cff577 100644 --- a/agent/x/agentmcp/manager.go +++ b/agent/x/agentmcp/manager.go @@ -2,6 +2,7 @@ package agentmcp import ( "context" + "encoding/base64" "errors" "fmt" "io/fs" @@ -14,9 +15,7 @@ import ( "sync" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" + "github.com/modelcontextprotocol/go-sdk/mcp" "golang.org/x/sync/errgroup" "golang.org/x/xerrors" tailscalesingleflight "tailscale.com/util/singleflight" @@ -126,10 +125,9 @@ type Manager struct { connectStartedHook func() } -// serverEntry pairs a server config with its connected client. type serverEntry struct { config ServerConfig - client *client.Client + client *mcp.ClientSession } // NewManager creates a new MCP client manager. The ctx bounds @@ -418,7 +416,7 @@ type serverDiff struct { type connectedServer struct { name string config ServerConfig - client *client.Client + client *mcp.ClientSession } // doReload reads MCP config files and performs a differential @@ -697,11 +695,9 @@ func (m *Manager) CallTool(ctx context.Context, req workspacesdk.CallMCPToolRequ callCtx, cancel := context.WithTimeout(ctx, toolCallTimeout) defer cancel() - result, err := entry.client.CallTool(callCtx, mcp.CallToolRequest{ - Params: mcp.CallToolParams{ - Name: originalName, - Arguments: req.Arguments, - }, + result, err := entry.client.CallTool(callCtx, &mcp.CallToolParams{ + Name: originalName, + Arguments: req.Arguments, }) if err != nil { return workspacesdk.CallMCPToolResponse{}, xerrors.Errorf("call tool %q on %q: %w", originalName, serverName, err) @@ -743,7 +739,7 @@ func (m *Manager) refreshCatalog(ctx context.Context, wanted map[string]ServerCo for name, entry := range servers { eg.Go(func() error { listCtx, cancel := context.WithTimeout(ctx, connectTimeout) - result, err := entry.client.ListTools(listCtx, mcp.ListToolsRequest{}) + result, err := entry.client.ListTools(listCtx, nil) cancel() if err != nil { logger.Warn(ctx, "failed to list tools from MCP server", @@ -858,75 +854,49 @@ func (m *Manager) Close() error { return errors.Join(errs...) } -// connectServer establishes a connection to a single MCP server -// and returns the connected client. It does not modify any Manager -// state. -func (m *Manager) connectServer(ctx context.Context, cfg ServerConfig) (*client.Client, error) { +// connectServer does not modify Manager state. +func (m *Manager) connectServer(ctx context.Context, cfg ServerConfig) (*mcp.ClientSession, error) { + // Use ctx for the transport so a stdio subprocess outlives the + // connect handshake. connectCtx bounds only Connect; closing the + // session or canceling ctx stops the subprocess. tr, err := m.createTransport(ctx, cfg) if err != nil { return nil, xerrors.Errorf("create transport for %q: %w", cfg.Name, err) } - c := client.NewClient(tr) + c := mcp.NewClient(&mcp.Implementation{ + Name: "coder-agent", + Version: buildinfo.Version(), + }, nil) connectCtx, cancel := context.WithTimeout(ctx, connectTimeout) defer cancel() - // Use the parent ctx (not connectCtx) so the subprocess outlives - // the connect/initialize handshake. connectCtx bounds only the - // Initialize call below. The subprocess is cleaned up when the - // Manager is closed or ctx is canceled. - if err := c.Start(ctx); err != nil { - _ = c.Close() - return nil, xerrors.Errorf("start %q: %w", cfg.Name, err) - } - - _, err = c.Initialize(connectCtx, mcp.InitializeRequest{ - Params: mcp.InitializeParams{ - ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - ClientInfo: mcp.Implementation{ - Name: "coder-agent", - Version: buildinfo.Version(), - }, - }, - }) + session, err := c.Connect(connectCtx, tr, nil) if err != nil { - _ = c.Close() - return nil, xerrors.Errorf("initialize %q: %w", cfg.Name, err) + return nil, xerrors.Errorf("connect %q: %w", cfg.Name, err) } - return c, nil + return session, nil } -// createTransport builds the mcp-go transport for a server config. -func (m *Manager) createTransport(ctx context.Context, cfg ServerConfig) (transport.Interface, error) { +func (m *Manager) createTransport(ctx context.Context, cfg ServerConfig) (mcp.Transport, error) { switch cfg.Transport { case "stdio": env := m.buildEnv(ctx, cfg.Env) - return transport.NewStdioWithOptions( - cfg.Command, - env, - cfg.Args, - transport.WithCommandFunc(func(ctx context.Context, command string, cmdEnv []string, args []string) (*exec.Cmd, error) { - cmd := m.execer.CommandContext(ctx, command, args...) - cmd.Env = cmdEnv - return cmd, nil - }), - ), nil + cmd := m.execer.CommandContext(ctx, cfg.Command, cfg.Args...) + cmd.Env = env + return &mcp.CommandTransport{Command: cmd}, nil case "http", "": - var opts []transport.StreamableHTTPCOption - opts = append(opts, transport.WithHTTPHeaders(cfg.Headers)) - if c := mcpHTTPClient(); c != nil { - opts = append(opts, transport.WithHTTPBasicClient(c)) - } - return transport.NewStreamableHTTP(cfg.URL, opts...) + return &mcp.StreamableClientTransport{ + Endpoint: cfg.URL, + HTTPClient: httpClientWithHeaders(cfg.Headers), + }, nil case "sse": - var sseOpts []transport.ClientOption - sseOpts = append(sseOpts, transport.WithHeaders(cfg.Headers)) - if c := mcpHTTPClient(); c != nil { - sseOpts = append(sseOpts, transport.WithHTTPClient(c)) - } - return transport.NewSSE(cfg.URL, sseOpts...) + return &mcp.SSEClientTransport{ + Endpoint: cfg.URL, + HTTPClient: httpClientWithHeaders(cfg.Headers), + }, nil default: return nil, xerrors.Errorf("unsupported transport %q", cfg.Transport) } @@ -993,29 +963,31 @@ func convertResult(result *mcp.CallToolResult) workspacesdk.CallMCPToolResponse var content []workspacesdk.MCPToolContent for _, item := range result.Content { switch c := item.(type) { - case mcp.TextContent: + case *mcp.TextContent: content = append(content, workspacesdk.MCPToolContent{ Type: "text", Text: c.Text, }) - case mcp.ImageContent: + case *mcp.ImageContent: + // The SDK decodes base64 during unmarshal; re-encode to + // keep the agent API's base64 wire format. content = append(content, workspacesdk.MCPToolContent{ Type: "image", - Data: c.Data, + Data: base64.StdEncoding.EncodeToString(c.Data), MediaType: c.MIMEType, }) - case mcp.AudioContent: + case *mcp.AudioContent: content = append(content, workspacesdk.MCPToolContent{ Type: "audio", - Data: c.Data, + Data: base64.StdEncoding.EncodeToString(c.Data), MediaType: c.MIMEType, }) - case mcp.EmbeddedResource: + case *mcp.EmbeddedResource: content = append(content, workspacesdk.MCPToolContent{ Type: "resource", Text: fmt.Sprintf("[embedded resource: %T]", c.Resource), }) - case mcp.ResourceLink: + case *mcp.ResourceLink: content = append(content, workspacesdk.MCPToolContent{ Type: "resource", Text: fmt.Sprintf("[resource link: %s]", c.URI), @@ -1055,23 +1027,21 @@ type ToolInfo struct { InputSchema map[string]any } -// toolInputSchemaMap converts an mcp-go tool input schema into the -// JSON-Schema-shaped map ToolInfo carries. Required is converted to -// []any so the downstream protobuf/structpb encoding accepts it. An -// empty schema yields nil so the tool ships with InputSchema unset. -func toolInputSchemaMap(s mcp.ToolInputSchema) map[string]any { +// Only type, properties, and required are exposed through ToolInfo; +// empty schemas leave InputSchema unset. +func toolInputSchemaMap(schema any) map[string]any { + m, ok := schema.(map[string]any) + if !ok { + return nil + } out := map[string]any{} - if s.Type != "" { - out["type"] = s.Type + if typ, ok := m["type"].(string); ok && typ != "" { + out["type"] = typ } - if len(s.Properties) > 0 { - out["properties"] = s.Properties + if properties, ok := m["properties"].(map[string]any); ok && len(properties) > 0 { + out["properties"] = properties } - if len(s.Required) > 0 { - required := make([]any, len(s.Required)) - for i, req := range s.Required { - required[i] = req - } + if required, ok := m["required"].([]any); ok && len(required) > 0 { out["required"] = required } if len(out) == 0 { diff --git a/agent/x/agentmcp/manager_internal_test.go b/agent/x/agentmcp/manager_internal_test.go index 8ec8bc77e83..75472a926a9 100644 --- a/agent/x/agentmcp/manager_internal_test.go +++ b/agent/x/agentmcp/manager_internal_test.go @@ -3,13 +3,14 @@ package agentmcp import ( "bufio" "context" + "encoding/base64" "encoding/json" "fmt" "os" "testing" "time" - "github.com/mark3labs/mcp-go/mcp" + "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -100,7 +101,7 @@ func TestConvertResult(t *testing.T) { name: "TextContent", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.TextContent{Type: "text", Text: "hello"}, + &mcp.TextContent{Text: "hello"}, }, }, want: workspacesdk.CallMCPToolResponse{ @@ -113,16 +114,15 @@ func TestConvertResult(t *testing.T) { name: "ImageContent", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.ImageContent{ - Type: "image", - Data: "base64data", + &mcp.ImageContent{ + Data: []byte("rawdata"), MIMEType: "image/png", }, }, }, want: workspacesdk.CallMCPToolResponse{ Content: []workspacesdk.MCPToolContent{ - {Type: "image", Data: "base64data", MediaType: "image/png"}, + {Type: "image", Data: base64.StdEncoding.EncodeToString([]byte("rawdata")), MediaType: "image/png"}, }, }, }, @@ -130,16 +130,15 @@ func TestConvertResult(t *testing.T) { name: "AudioContent", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.AudioContent{ - Type: "audio", - Data: "base64audio", + &mcp.AudioContent{ + Data: []byte("rawaudio"), MIMEType: "audio/mp3", }, }, }, want: workspacesdk.CallMCPToolResponse{ Content: []workspacesdk.MCPToolContent{ - {Type: "audio", Data: "base64audio", MediaType: "audio/mp3"}, + {Type: "audio", Data: base64.StdEncoding.EncodeToString([]byte("rawaudio")), MediaType: "audio/mp3"}, }, }, }, @@ -147,7 +146,7 @@ func TestConvertResult(t *testing.T) { name: "IsErrorPropagation", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.TextContent{Type: "text", Text: "fail"}, + &mcp.TextContent{Text: "fail"}, }, IsError: true, }, @@ -162,10 +161,9 @@ func TestConvertResult(t *testing.T) { name: "MultipleContentItems", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.TextContent{Type: "text", Text: "caption"}, - mcp.ImageContent{ - Type: "image", - Data: "imgdata", + &mcp.TextContent{Text: "caption"}, + &mcp.ImageContent{ + Data: []byte("imgdata"), MIMEType: "image/jpeg", }, }, @@ -173,7 +171,7 @@ func TestConvertResult(t *testing.T) { want: workspacesdk.CallMCPToolResponse{ Content: []workspacesdk.MCPToolContent{ {Type: "text", Text: "caption"}, - {Type: "image", Data: "imgdata", MediaType: "image/jpeg"}, + {Type: "image", Data: base64.StdEncoding.EncodeToString([]byte("imgdata")), MediaType: "image/jpeg"}, }, }, }, @@ -181,9 +179,8 @@ func TestConvertResult(t *testing.T) { name: "ResourceLink", input: &mcp.CallToolResult{ Content: []mcp.Content{ - mcp.ResourceLink{ - Type: "resource_link", - URI: "file:///tmp/test.txt", + &mcp.ResourceLink{ + URI: "file:///tmp/test.txt", }, }, }, @@ -242,7 +239,7 @@ func TestConnectServer_StdioProcessSurvivesConnect(t *testing.T) { // alive. Verify by listing tools (requires a live server). listCtx, listCancel := context.WithTimeout(ctx, testutil.WaitShort) defer listCancel() - result, err := client.ListTools(listCtx, mcp.ListToolsRequest{}) + result, err := client.ListTools(listCtx, nil) require.NoError(t, err, "ListTools should succeed, server must be alive after connect") require.Len(t, result.Tools, 1) assert.Equal(t, "echo", result.Tools[0].Name) diff --git a/agent/x/agentmcp/mcphttpclient.go b/agent/x/agentmcp/mcphttpclient.go index 0b4c07ea3c0..e09218046ee 100644 --- a/agent/x/agentmcp/mcphttpclient.go +++ b/agent/x/agentmcp/mcphttpclient.go @@ -5,6 +5,33 @@ import ( "net/http" ) +func httpClientWithHeaders(headers map[string]string) *http.Client { + base := http.DefaultTransport + if isolated := mcpHTTPClient(); isolated != nil { + base = isolated.Transport + } + if len(headers) == 0 { + return &http.Client{Transport: base} + } + return &http.Client{Transport: &headerRoundTripper{ + base: base, + headers: headers, + }} +} + +type headerRoundTripper struct { + base http.RoundTripper + headers map[string]string +} + +func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + for k, v := range h.headers { + clone.Header.Set(k, v) + } + return h.base.RoundTrip(clone) +} + // mcpHTTPClient returns an isolated *http.Client when running // inside tests, or nil for production. During tests, // httptest.Server.Close() calls diff --git a/agent/x/agentmcp/reload_internal_test.go b/agent/x/agentmcp/reload_internal_test.go index 192fef21fe6..d22f7f5d520 100644 --- a/agent/x/agentmcp/reload_internal_test.go +++ b/agent/x/agentmcp/reload_internal_test.go @@ -10,7 +10,6 @@ import ( "sync" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -508,8 +507,9 @@ func TestDifferentialReload(t *testing.T) { origClient := m.servers["srv"].client m.mu.RUnlock() - // Change the server's args to trigger a diff. - entry.Args = append(entry.Args, "-test.v") + // Change the environment because verbose test flags make the + // fake server write non-protocol output, which the SDK rejects. + entry.Env["EXTRA_DIFF_TRIGGER"] = "1" writeMCPConfig(t, dir, map[string]mcpServerEntry{"srv": entry}) err = m.Reload(ctx, []string{configPath}) @@ -561,7 +561,7 @@ func TestDifferentialReload(t *testing.T) { // ListTools on a closed client returns an error. listCtx, cancel := context.WithTimeout(ctx, testutil.WaitShort) defer cancel() - _, listErr := oldClientB.ListTools(listCtx, mcp.ListToolsRequest{}) + _, listErr := oldClientB.ListTools(listCtx, nil) assert.Error(t, listErr, "ListTools on closed client should fail") })