From df5eb4d84aab55b74d95124d7621d52e4e7ee95b Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 12 Aug 2026 08:36:53 +0000 Subject: [PATCH] feat(agent/x/agentmcp): migrate workspace agent MCP client to official Go SDK Replace the mark3labs client with the official SDK client for .mcp.json-configured servers. Stdio servers run via CommandTransport with the same execer-built command and enriched environment; the subprocess still outlives the connect handshake because the command context is the manager's, not the bounded connect context. HTTP and SSE headers move to an http.RoundTripper on the transport's client. Config parsing, transport inference, tool-name prefixing, catalog snapshots, and the agent call-tool API keep their external shapes. The SDK re-encodes binary content as base64 for the agent API and, unlike mark3labs, closes stdio connections on non-protocol stdout output, which is what the spec requires of servers. --- agent/x/agentmcp/manager.go | 136 +++++++++------------- agent/x/agentmcp/manager_internal_test.go | 37 +++--- agent/x/agentmcp/mcphttpclient.go | 27 +++++ agent/x/agentmcp/reload_internal_test.go | 8 +- 4 files changed, 101 insertions(+), 107 deletions(-) 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") })