mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
+53
-83
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user