feat(agent/x/agentmcp): migrate workspace agent MCP client to official Go SDK (#28059)

## Stack Context

PR 4 of 6 in a stack that migrates every Coder MCP surface from the
archived `github.com/mark3labs/mcp-go` library to the official
`github.com/modelcontextprotocol/go-sdk` v1.7.0.

Stack: #28056 -> #28057 -> #28058 -> #28059 -> #28060 -> #28061

## Why

The workspace agent MCP manager now stores `*mcp.ClientSession` per
configured server.

- stdio servers use `mcp.CommandTransport` with an `exec.Cmd` built from
Coder's `agentexec.Execer`, preserving environment enrichment; the
command uses the manager's parent context so a stdio subprocess outlives
the connect handshake and stops when the session closes.
- HTTP and SSE servers use header-injecting HTTP clients.
- Binary tool content is re-encoded to base64 for the agent API because
the official SDK decodes it into raw bytes.
- The reload test now triggers config diffs via an environment variable
because the official SDK drops connections on non-protocol stdout output
(flags like `-test.v` made the fake server chatty).

> Mux created this PR on Mike's behalf.
This commit is contained in:
Michael Suchacz
2026-08-13 10:29:13 +00:00
committed by GitHub
parent 1e546ea8a3
commit 7720e283f5
4 changed files with 101 additions and 107 deletions
+53 -83
View File
@@ -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 {
+17 -20
View File
@@ -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)
+27
View File
@@ -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
+4 -4
View File
@@ -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")
})