From b0bd93b15cb3fe38c66f530cf0cf5c65e97b6fdd Mon Sep 17 00:00:00 2001 From: Aska Date: Sun, 6 Sep 2026 17:27:14 +0800 Subject: [PATCH] =?UTF-8?q?feat(mcp):=20Top5#3=20=E5=AF=B9=E9=BD=90=20MCP?= =?UTF-8?q?=20Go=20SDK=E2=80=94=E2=80=94=E5=8D=8F=E8=AE=AE=E7=89=88?= =?UTF-8?q?=E6=9C=AC=E5=8D=8F=E5=95=86=20+=20list=5Fchanged=20=E5=B7=AE?= =?UTF-8?q?=E9=87=8F=E5=88=B7=E6=96=B0=20+=20keepalive=20+=20Streamable=20?= =?UTF-8?q?HTTP?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 协议版本协商：initialize 先发最高版本(2025-11-25)并声明能力集合，遇 UnsupportedProtocolVersionError(-32022)/服务器不支持版本 → negotiateMutuallySupportedVersion 求出双方最高共同版本降级重试(最多2次)；无共同版本→ErrUnsupportedProtocolVersion。 协商结果记 negotiatedVersion，HTTP 传输回填 Mcp-Protocol-Version 头。 - tools/list_changed 订阅：capabilities 声明 tools.listChanged:true；传输读无 id 帧路由到 handleServerNotification；refreshPending 原子旗标合并突发 + 新 goroutine 调 onToolListChanged (避免 stdio lineRPC 重入自锁)；Registry 真差量刷新(重新 DiscoverTools→移除消失/注册新增， 跳过同名冲突保外来工具)。 - keepalive：WithKeepAlive(interval,failures) 起 goroutine 定时 Call(ping)，连续失败≥阈值关 session(默认3,可配)；服务器不支持 ping(-32601)停 keepalive 不误关；ctx+Close 取消不泄漏。 - Streamable HTTP 传输：newTransport 支持 TransportType=http|streamable(保留 stdio/SSE 推断， 空值走旧逻辑不破坏)；Call POST+Accept json/event-stream 响应可单 JSON 或 SSE 流找匹配 id； Mcp-Session-Id 会话 + 独立 GET SSE 流收 server 通知；config 读 MCP_TRANSPORT。 - schema 缓存判定不做：客户端透传 map[string]any 无重复编译，SDK SchemaCache 是服务端 按 Go 类型反射场景，客户端无对应。 go build/vet/test 全绿(count=2 无 flaky) + windows/darwin 交叉编译。 --- internal/capability/mcp/client.go | 209 +++++++++- internal/capability/mcp/config.go | 12 +- internal/capability/mcp/config_test.go | 19 + internal/capability/mcp/enhancements_test.go | 194 +++++++++ internal/capability/mcp/helpers_test.go | 105 ++++- internal/capability/mcp/http.go | 412 +++++++++++++++++++ internal/capability/mcp/http_test.go | 197 +++++++++ internal/capability/mcp/mcp.go | 215 +++++++++- internal/capability/mcp/registry.go | 72 +++- internal/capability/mcp/registry_test.go | 53 +++ internal/capability/mcp/rpc.go | 5 +- internal/capability/mcp/sse.go | 46 ++- internal/capability/mcp/stdio.go | 44 +- 13 files changed, 1517 insertions(+), 66 deletions(-) create mode 100644 internal/capability/mcp/enhancements_test.go create mode 100644 internal/capability/mcp/http.go create mode 100644 internal/capability/mcp/http_test.go diff --git a/internal/capability/mcp/client.go b/internal/capability/mcp/client.go index a04782e..9840121 100644 --- a/internal/capability/mcp/client.go +++ b/internal/capability/mcp/client.go @@ -2,8 +2,11 @@ package mcp import ( "context" + "encoding/json" + "errors" "fmt" "log/slog" + "time" "git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk" ) @@ -61,11 +64,17 @@ func (c *Client) Connect(ctx context.Context) error { tr.Close() return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err) } + // Wire server→client notification routing (list_changed) before the server + // can start sending them after notifications/initialized. + attachNotificationSink(tr, c.handleServerNotification) + info, err := c.initialize(ctx, tr) if err != nil { tr.Close() return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err) } + // Let the transport echo the negotiated protocol version (Streamable HTTP). + attachProtocolProvider(tr, c.protocolVersion) c.sendInitialized(ctx, tr) c.mu.Lock() @@ -73,7 +82,9 @@ func (c *Client) Connect(ctx context.Context) error { c.info = info c.state = stateConnected c.mu.Unlock() - c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name) + + c.startKeepalive() + c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name, "protocol", c.protocolVersion()) return nil } @@ -87,19 +98,94 @@ func (c *Client) Initialize(ctx context.Context) (ServerInfo, error) { return c.initialize(ctx, tr) } -// initialize sends the initialize request over tr and records the server -// identity. +// initialize performs the MCP initialize handshake with protocol version +// negotiation (DevDocs/design/09 §1.1). It does not hard-code a single version: +// it first requests latestProtocolVersion, and if the server rejects that version +// (a JSON-RPC error with code codeUnsupportedProtocolVersion carrying the +// supported set) or replies with a version this client cannot speak, it downgrades +// once to the highest mutually supported version and retries. It records the +// agreed version so the HTTP transport can echo it back as Mcp-Protocol-Version. func (c *Client) initialize(ctx context.Context, tr Transport) (ServerInfo, error) { - params := initializeParams{ - ProtocolVersion: protocolVersion, - Capabilities: map[string]any{}, - ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"}, + sent := c.preferredProtocolVersion + if sent == "" { + sent = latestProtocolVersion } - var res initializeResult - if err := tr.Call(ctx, "initialize", params, &res); err != nil { - return ServerInfo{}, err + var lastErr error + for attempt := 0; attempt < maxInitializeAttempts; attempt++ { + params := initializeParams{ + ProtocolVersion: sent, + Capabilities: c.capabilities(sent), + ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"}, + } + var res initializeResult + err := tr.Call(ctx, "initialize", params, &res) + if err != nil { + // Server rejected the requested version and may advertise the versions + // it supports; downgrade to the highest mutually supported one. + if supported, ok := unsupportedVersions(err); ok { + if v := negotiateMutuallySupportedVersion(supported); v != "" && v != sent { + c.log.Debug("mcp downgrading protocol version", + "name", c.cfg.Name, "requested", sent, "version", v) + sent = v + lastErr = err + continue + } + return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, err) + } + return ServerInfo{}, err + } + negotiated := res.ProtocolVersion + if negotiated == "" { + negotiated = sent + } + if !isSupportedVersion(negotiated) { + // The server replied with a version this client does not speak; + // downgrade once to the highest version we both accept. + if v := highestCommonVersion(negotiated); v != "" && v != sent { + c.log.Debug("mcp renegotiating protocol version", + "name", c.cfg.Name, "server", negotiated, "version", v) + sent = v + lastErr = err + continue + } + return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: server advertised %q", c.cfg.Name, ErrUnsupportedProtocolVersion, negotiated) + } + c.setProtocolVersion(negotiated) + return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil } - return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil + return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, lastErr) +} + +// isMethodNotFound reports whether err wraps a JSON-RPC "method not found" error +// (code -32601), which a peer returns when it does not implement a method such as +// the keepalive ping. +func isMethodNotFound(err error) bool { + var re *rpcError + return errors.As(err, &re) && re.Code == codeMethodNotFound +} + +// unsupportedVersions reports the protocol versions the server supports, extracted +// from the data field of a codeUnsupportedProtocolVersion JSON-RPC error. It +// reports ok=false when err does not carry such an error or the data is unusable. +func unsupportedVersions(err error) ([]string, bool) { + var re *rpcError + if !errors.As(err, &re) || re.Code != codeUnsupportedProtocolVersion { + return nil, false + } + if re.Data == nil { + return nil, false + } + data, err := json.Marshal(re.Data) + if err != nil { + return nil, false + } + var payload struct { + Supported []string `json:"supported"` + } + if err := json.Unmarshal(data, &payload); err != nil || len(payload.Supported) == 0 { + return nil, false + } + return payload.Supported, true } // sendInitialized best-effort sends the notifications/initialized advisory ack. @@ -113,6 +199,87 @@ func (c *Client) sendInitialized(ctx context.Context, tr Transport) { } } +// handleServerNotification is the transport-level sink for server→client +// notifications. It offloads the actual work to a fresh goroutine because it runs +// on the transport read path and must not block or re-enter the same transport +// synchronously; refreshPending coalesces bursts so at most one refresh runs at a +// time. +func (c *Client) handleServerNotification(method string) { + if method != notificationToolsListChanged || c.onToolListChanged == nil { + return + } + if !c.refreshPending.CompareAndSwap(false, true) { + return + } + go func() { + defer c.refreshPending.Store(false) + c.onToolListChanged() + }() +} + +// ping sends a JSON-RPC ping request and reports whether the server responded. +// Pings are used by the keepalive loop to detect a dead connection. +func (c *Client) ping(ctx context.Context) error { + tr := c.transport() + if tr == nil { + return fmt.Errorf("mcp ping: %w", ErrNotConnected) + } + if err := tr.Call(ctx, "ping", map[string]any{}, nil); err != nil { + return fmt.Errorf("mcp ping: %w", err) + } + return nil +} + +// startKeepalive launches the ping keepalive loop when configured. It closes the +// client once keepAliveFailureThreshold consecutive pings fail, so the host can +// rebuild the session. The loop exits when the client is closed. +func (c *Client) startKeepalive() { + interval := c.keepAlive + if interval <= 0 { + return + } + ctx, cancel := context.WithCancel(context.Background()) + c.mu.Lock() + c.keepaliveCancel = cancel + c.mu.Unlock() + threshold := c.keepAliveFailures + if threshold < 1 { + threshold = keepAliveFailureThreshold + } + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + failures := 0 + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2) + err := c.ping(pingCtx) + pingCancel() + if err != nil { + // A peer that does not implement ping is still healthy; stop + // keepalive rather than treating every ping as a failure. + if isMethodNotFound(err) { + c.log.Debug("mcp keepalive stopped; server does not support ping", "name", c.cfg.Name) + return + } + failures++ + c.log.Warn("mcp keepalive ping failed", "name", c.cfg.Name, "err", err, "failures", failures) + if failures < threshold { + continue + } + c.log.Error("mcp keepalive ping failed; closing session", "name", c.cfg.Name, "failures", failures) + _ = c.Close() + return + } + failures = 0 + } + } + }() +} + // ListTools returns the tools exposed by the connected MCP server. func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) { tr := c.transport() @@ -162,7 +329,8 @@ func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) { return out, nil } -// Close closes the transport and marks the client closed. It is idempotent. +// Close closes the transport and marks the client closed. It is idempotent. It +// also cancels the keepalive loop so no ping goroutine outlives the client. func (c *Client) Close() error { c.mu.Lock() if c.state == stateClosed { @@ -171,6 +339,10 @@ func (c *Client) Close() error { } c.state = stateClosed tr := c.tr + if c.keepaliveCancel != nil { + c.keepaliveCancel() + c.keepaliveCancel = nil + } c.mu.Unlock() if tr != nil { return tr.Close() @@ -178,9 +350,18 @@ func (c *Client) Close() error { return nil } -// newTransport selects the transport matching cfg: stdio when Command is set, SSE -// when URL is set. +// newTransport selects the transport matching cfg. An explicit TransportType +// ("stdio", "sse" or "http") wins; otherwise the transport is inferred from the +// fields, preserving the legacy behaviour of stdio-when-Command and SSE-when-URL. func newTransport(cfg MCPConfig, log *slog.Logger) (Transport, error) { + switch cfg.TransportType { + case "http", "streamable": + return newHTTPTransport(cfg, log), nil + case "sse": + return newSSETransport(cfg, log), nil + case "stdio": + return newStdioTransport(cfg, log), nil + } switch { case cfg.Command != "": return newStdioTransport(cfg, log), nil diff --git a/internal/capability/mcp/config.go b/internal/capability/mcp/config.go index 91bd9a3..f9cd784 100644 --- a/internal/capability/mcp/config.go +++ b/internal/capability/mcp/config.go @@ -74,15 +74,17 @@ func readOne(env map[string]string, prefix string) (MCPConfig, bool) { name := env[p+"_NAME"] command := env[p+"_COMMAND"] url := env[p+"_URL"] + transp := env[p+"_TRANSPORT"] if name == "" && command == "" && url == "" { return MCPConfig{}, false } return MCPConfig{ - Name: name, - Command: command, - Args: splitArgs(env[p+"_ARGS"]), - URL: url, - Enabled: parseBool(env[p+"_ENABLED"], false), + Name: name, + Command: command, + Args: splitArgs(env[p+"_ARGS"]), + URL: url, + TransportType: transp, + Enabled: parseBool(env[p+"_ENABLED"], false), }, true } diff --git a/internal/capability/mcp/config_test.go b/internal/capability/mcp/config_test.go index b9606ac..996344c 100644 --- a/internal/capability/mcp/config_test.go +++ b/internal/capability/mcp/config_test.go @@ -52,6 +52,25 @@ func TestFromEnv_MultipleServers(t *testing.T) { } } +func TestFromEnv_TransportType(t *testing.T) { + t.Parallel() + pairs := []string{ + "MCP_NAME=web", + "MCP_URL=http://localhost:9090/mcp", + "MCP_TRANSPORT=http", + } + cfg, err := FromEnv(pairs) + if err != nil { + t.Fatalf("FromEnv: %v", err) + } + if len(cfg.Servers) != 1 { + t.Fatalf("servers = %d, want 1", len(cfg.Servers)) + } + if cfg.Servers[0].TransportType != "http" { + t.Fatalf("TransportType = %q, want http", cfg.Servers[0].TransportType) + } +} + func TestFromEnv_NoServers(t *testing.T) { t.Parallel() cfg, err := FromEnv([]string{"MCP_ENABLED=true"}) diff --git a/internal/capability/mcp/enhancements_test.go b/internal/capability/mcp/enhancements_test.go new file mode 100644 index 0000000..1154db4 --- /dev/null +++ b/internal/capability/mcp/enhancements_test.go @@ -0,0 +1,194 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "reflect" + "testing" + "time" +) + +func TestClient_Connect_NegotiatesDowngradeOnUnsupportedError(t *testing.T) { + t.Parallel() + ft := &fakeTransport{ + initErrors: []error{unsupportedProtocolError("2024-11-05")}, + initResponseVersion: protocolVersion20241105, + } + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + got := ft.initSentVersions() + if !reflect.DeepEqual(got, []string{"2025-11-25", "2024-11-05"}) { + t.Fatalf("init versions = %v, want [2025-11-25 2024-11-05]", got) + } + if c.protocolVersion() != protocolVersion20241105 { + t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105) + } +} + +func TestClient_Connect_AcceptsServerSupportedVersion(t *testing.T) { + t.Parallel() + ft := &fakeTransport{} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + if got := ft.initSentVersions(); !reflect.DeepEqual(got, []string{"2025-11-25"}) { + t.Fatalf("init versions = %v, want [2025-11-25]", got) + } + // The fake echoes the legacy version; the client must accept it as supported. + if c.protocolVersion() != protocolVersion20241105 { + t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105) + } +} + +func TestClient_Connect_NoMutualVersion(t *testing.T) { + t.Parallel() + // The server only supports a version the client cannot speak, so negotiation + // must exhaust and surface ErrUnsupportedProtocolVersion. + ft := &fakeTransport{initErrors: []error{unsupportedProtocolError("2000-01-01")}} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) + err := c.Connect(context.Background()) + if !errors.Is(err, ErrUnsupportedProtocolVersion) { + t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err) + } +} + +func TestClient_Connect_ServerReturnsUnsupportedVersion(t *testing.T) { + t.Parallel() + // A server that echoes a version this client does not support (and cannot + // downgrade to) must fail the handshake with ErrUnsupportedProtocolVersion. + ft := &fakeTransport{initResponseVersion: "2099-01-01"} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) + err := c.Connect(context.Background()) + if !errors.Is(err, ErrUnsupportedProtocolVersion) { + t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err) + } +} + +func TestClient_ToolListChangedInvokesHandler(t *testing.T) { + t.Parallel() + called := make(chan struct{}, 1) + ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}} + c := New(MCPConfig{Name: "s", Command: "x"}, + WithTransport(ft), + WithToolListChangedHandler(func() { called <- struct{}{} })) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + // No handler invoked for unrelated notifications. + ft.notifyServer("notifications/prompts/list_changed") + select { + case <-called: + t.Fatal("handler invoked for a non-tool notification") + case <-time.After(50 * time.Millisecond): + } + // The real notification must invoke the handler. + ft.notifyServer(notificationToolsListChanged) + select { + case <-called: + case <-time.After(time.Second): + t.Fatal("tool list changed handler was not invoked") + } +} + +func TestClient_KeepAlive_SendsPings(t *testing.T) { + t.Parallel() + ft := &fakeTransport{} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 3)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 }) +} + +func TestClient_KeepAlive_FailureClosesSession(t *testing.T) { + t.Parallel() + ft := &fakeTransport{pingErr: errors.New("dead peer")} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + // A single failed ping (threshold 1) must close the client/transport. + waitFor(t, time.Second, func() bool { return c.transport() == nil }) + waitFor(t, time.Second, func() bool { + ft.mu.Lock() + defer ft.mu.Unlock() + return ft.closes >= 1 + }) +} + +func TestClient_KeepAlive_StopsOnUnsupportedPing(t *testing.T) { + t.Parallel() + ft := &fakeTransport{pingErr: fmt.Errorf("mcp ping: %w: %w", ErrRPC, &rpcError{Code: codeMethodNotFound, Message: "method not found"})} + c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + // A method-not-found ping must stop keepalive without closing the session. + waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 }) + time.Sleep(80 * time.Millisecond) + if c.transport() == nil { + t.Fatal("client was closed after an unsupported ping, want it to stay connected") + } + if ft.pingCount() > 1 { + t.Fatalf("keepalive kept pinging after method-not-found; pings = %d", ft.pingCount()) + } +} + +func TestNewTransport_SelectsTransportType(t *testing.T) { + t.Parallel() + tests := []struct { + name string + cfg MCPConfig + want any + }{ + {name: "http", cfg: MCPConfig{URL: "http://x", TransportType: "http"}, want: &streamableHTTPTransport{}}, + {name: "streamable-alias", cfg: MCPConfig{URL: "http://x", TransportType: "streamable"}, want: &streamableHTTPTransport{}}, + {name: "sse", cfg: MCPConfig{URL: "http://x", TransportType: "sse"}, want: &sseTransport{}}, + {name: "stdio", cfg: MCPConfig{Command: "x", TransportType: "stdio"}, want: &stdioTransport{}}, + {name: "infer-stdio", cfg: MCPConfig{Command: "x"}, want: &stdioTransport{}}, + {name: "infer-sse", cfg: MCPConfig{URL: "http://x"}, want: &sseTransport{}}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + tr, err := newTransport(tt.cfg, defaultLogger()) + if err != nil { + t.Fatalf("newTransport: %v", err) + } + if reflect.TypeOf(tr) != reflect.TypeOf(tt.want) { + t.Fatalf("transport type = %T, want %T", tr, tt.want) + } + }) + } +} + +func TestNewTransport_ExplicitTypeSelectsEvenWithoutField(t *testing.T) { + t.Parallel() + // Selecting an explicit transport does not validate the field; Start does. + if tr, err := newTransport(MCPConfig{TransportType: "http"}, defaultLogger()); err != nil { + t.Fatalf("newTransport(http): %v", err) + } else if _, ok := tr.(*streamableHTTPTransport); !ok { + t.Fatalf("transport type = %T, want *streamableHTTPTransport", tr) + } +} + +// waitFor polls cond until it returns true or the deadline passes, failing the +// test on timeout. +func waitFor(t *testing.T, timeout time.Duration, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("condition not met before timeout") +} diff --git a/internal/capability/mcp/helpers_test.go b/internal/capability/mcp/helpers_test.go index df11772..c40fccc 100644 --- a/internal/capability/mcp/helpers_test.go +++ b/internal/capability/mcp/helpers_test.go @@ -2,23 +2,40 @@ package mcp import ( "context" + "fmt" "sync" ) // fakeTransport is a scripted Transport for tests. It records starts, closes and // the methods it was asked to call so assertions can verify the client's usage. +// It also lets tests script the initialize handshake, simulate server→client +// notifications and control the ping result. type fakeTransport struct { startErr error closeErr error initErr error listErr error callErr error + pingErr error tools []ToolInfo - mu sync.Mutex - starts int - closes int - methods []string + // initErrors is a per-call queue of initialize errors. When non-empty one + // entry is consumed per initialize call; once exhausted the handshake + // succeeds. initErr is returned for every call when initErrors is empty. + initErrors []error + // initVersions records the protocolVersion sent in each initialize call. + initVersions []string + // initResponseVersion is the protocolVersion echoed on a successful + // initialize; default 2024-11-05. + initResponseVersion string + + mu sync.Mutex + starts int + closes int + methods []string + pings int + onNotifyFn func(method string) + activeTools []ToolInfo } // Start records a start. @@ -26,29 +43,42 @@ func (f *fakeTransport) Start(_ context.Context) error { f.mu.Lock() defer f.mu.Unlock() f.starts++ + f.activeTools = f.tools return f.startErr } -// Call answers the three MCP methods the client issues. -func (f *fakeTransport) Call(_ context.Context, method string, _ any, result any) error { +// Call answers the MCP methods the client issues. +func (f *fakeTransport) Call(_ context.Context, method string, params any, result any) error { f.mu.Lock() defer f.mu.Unlock() f.methods = append(f.methods, method) switch method { case "initialize": + if p, ok := params.(initializeParams); ok { + f.initVersions = append(f.initVersions, p.ProtocolVersion) + } + if len(f.initErrors) > 0 { + err := f.initErrors[0] + f.initErrors = f.initErrors[1:] + return err + } if f.initErr != nil { return f.initErr } if r, ok := result.(*initializeResult); ok { r.ServerInfo = serverInfo{Name: "fake-server", Version: "1.0"} - r.ProtocolVersion = protocolVersion + r.ProtocolVersion = f.initResponseVersion + if r.ProtocolVersion == "" { + r.ProtocolVersion = protocolVersion20241105 + } + f.activeTools = f.tools } case "tools/list": if f.listErr != nil { return f.listErr } if r, ok := result.(*listToolsResult); ok { - for _, t := range f.tools { + for _, t := range f.activeTools { r.Tools = append(r.Tools, toolSpec{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema}) } } @@ -59,6 +89,11 @@ func (f *fakeTransport) Call(_ context.Context, method string, _ any, result any if r, ok := result.(*callToolResult); ok { r.Content = []contentBlock{{Type: "text", Text: "hello"}} } + case "ping": + f.pings++ + if f.pingErr != nil { + return f.pingErr + } } return nil } @@ -71,6 +106,25 @@ func (f *fakeTransport) notify(_ context.Context, method string, _ any) error { return nil } +// setOnNotification records the client's server→client notification handler so +// tests can deliver one via notifyServer. +func (f *fakeTransport) setOnNotification(fn func(method string)) { + f.mu.Lock() + defer f.mu.Unlock() + f.onNotifyFn = fn +} + +// notifyServer simulates a server→client notification (e.g. +// notifications/tools/list_changed) delivered to the client's handler. +func (f *fakeTransport) notifyServer(method string) { + f.mu.Lock() + fn := f.onNotifyFn + f.mu.Unlock() + if fn != nil { + fn(method) + } +} + // Close records a close. func (f *fakeTransport) Close() error { f.mu.Lock() @@ -97,3 +151,38 @@ func (f *fakeTransport) hasMethod(method string) bool { } return false } + +// initSentVersions returns the protocol versions the client requested across all +// initialize calls. +func (f *fakeTransport) initSentVersions() []string { + f.mu.Lock() + defer f.mu.Unlock() + out := make([]string, len(f.initVersions)) + copy(out, f.initVersions) + return out +} + +// pingCount returns the recorded ping request count. +func (f *fakeTransport) pingCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.pings +} + +// setTools replaces the tool list served on subsequent tools/list calls. +func (f *fakeTransport) setTools(tools []ToolInfo) { + f.mu.Lock() + defer f.mu.Unlock() + f.tools = tools + f.activeTools = tools +} + +// unsupportedProtocolError builds a JSON-RPC error that rejects the requested +// protocol version and advertises the supported set, as a real MCP server does. +func unsupportedProtocolError(supported ...string) error { + return fmt.Errorf("mcp initialize: %w: %w", ErrRPC, &rpcError{ + Code: codeUnsupportedProtocolVersion, + Message: "unsupported protocol version", + Data: map[string]any{"supported": supported}, + }) +} diff --git a/internal/capability/mcp/http.go b/internal/capability/mcp/http.go new file mode 100644 index 0000000..4b3f177 --- /dev/null +++ b/internal/capability/mcp/http.go @@ -0,0 +1,412 @@ +package mcp + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "strconv" + "strings" + "sync" + "time" +) + +// mcpSessionIDHeader carries the Streamable HTTP session identifier returned by a +// server that maintains a session (MCP 2025-11-25 Streamable HTTP). +const mcpSessionIDHeader = "Mcp-Session-Id" + +// maxHTTPResponse bounds how much of a single POST reply body is read before the +// response is considered non-matching. +const maxHTTPResponse = 1 << 20 + +// streamableHTTPTransport speaks JSON-RPC 2.0 to an MCP server over Streamable +// HTTP (MCP 2025-11-25). Client→server messages are POSTed to the endpoint, and +// the server answers either with a single application/json body or an SSE stream. +// Server-initiated notifications (e.g. notifications/tools/list_changed) are +// delivered over a standalone GET SSE stream opened at Start. +type streamableHTTPTransport struct { + cfg MCPConfig + log *slog.Logger + url string + http *http.Client + + mu sync.Mutex + nextID int + started bool + session string + protocol func() string + + onNotification func(method string) + + closeCh chan struct{} + closeOnce sync.Once + getCancel context.CancelFunc +} + +// newHTTPTransport builds a Streamable HTTP transport for cfg. +func newHTTPTransport(cfg MCPConfig, log *slog.Logger) *streamableHTTPTransport { + return &streamableHTTPTransport{ + cfg: cfg, + log: log, + url: cfg.URL, + http: http.DefaultClient, + closeCh: make(chan struct{}), + } +} + +// Start validates the endpoint and opens the standalone GET SSE stream for +// server-initiated messages. The GET stream is best-effort: a server that only +// answers via POST replies 4xx, which we tolerate because request/response still +// works over POST. +func (s *streamableHTTPTransport) Start(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return fmt.Errorf("mcp http start: %w", err) + } + if s.cfg.URL == "" { + return fmt.Errorf("mcp http: %w", ErrNoTransport) + } + if !strings.HasPrefix(s.cfg.URL, "http://") && !strings.HasPrefix(s.cfg.URL, "https://") { + return fmt.Errorf("mcp http: invalid url %q", s.cfg.URL) + } + s.mu.Lock() + s.started = true + s.mu.Unlock() + + getCtx, cancel := context.WithCancel(context.Background()) + s.mu.Lock() + s.getCancel = cancel + s.mu.Unlock() + go s.readStandalone(getCtx) + return nil +} + +// setOnNotification records the handler invoked when the server sends a +// server-initiated notification over the standalone SSE stream. +func (s *streamableHTTPTransport) setOnNotification(fn func(method string)) { + s.mu.Lock() + defer s.mu.Unlock() + s.onNotification = fn +} + +// setProtocolProvider records the getter for the negotiated MCP protocol version, +// echoing it back in the Mcp-Protocol-Version header on every request. +func (s *streamableHTTPTransport) setProtocolProvider(fn func() string) { + s.mu.Lock() + defer s.mu.Unlock() + s.protocol = fn +} + +// Call POSTs a request and decodes the matching response, from either a JSON body +// or an SSE stream returned by the server. +func (s *streamableHTTPTransport) Call(ctx context.Context, method string, params any, result any) error { + s.mu.Lock() + if !s.started { + s.mu.Unlock() + return fmt.Errorf("mcp http %s: %w", method, ErrNotConnected) + } + s.nextID++ + id := strconv.Itoa(s.nextID) + s.mu.Unlock() + + reqBody, err := encodeRequest(id, method, params) + if err != nil { + return fmt.Errorf("mcp http %s: %w", method, err) + } + resp, err := s.post(ctx, reqBody) + if err != nil { + return fmt.Errorf("mcp http %s: %w", method, err) + } + defer resp.Body.Close() + s.captureSession(resp) + // Close the response body if the caller cancels, so a blocked streamed read + // unblocks rather than leaking a goroutine. + stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() }) + defer stopClose() + + return s.decodeResponse(resp, id, method, result) +} + +// post sends a JSON-RPC request body to the endpoint with the streamable HTTP +// headers and returns the response. +func (s *streamableHTTPTransport) post(ctx context.Context, body []byte) (*http.Response, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + s.applyHeaders(req) + return s.http.Do(req) +} + +// applyHeaders sets the Streamable HTTP request headers. +func (s *streamableHTTPTransport) applyHeaders(req *http.Request) { + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + s.mu.Lock() + session := s.session + proto := "" + if s.protocol != nil { + proto = s.protocol() + } + s.mu.Unlock() + if session != "" { + req.Header.Set(mcpSessionIDHeader, session) + } + if proto != "" { + req.Header.Set("Mcp-Protocol-Version", proto) + } +} + +// captureSession stores a session id returned by the server. +func (s *streamableHTTPTransport) captureSession(resp *http.Response) { + if sid := resp.Header.Get(mcpSessionIDHeader); sid != "" { + s.mu.Lock() + s.session = sid + s.mu.Unlock() + } +} + +// decodeResponse reads a POST response and decodes the reply for id. +func (s *streamableHTTPTransport) decodeResponse(resp *http.Response, id, method string, result any) error { + if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") { + return s.readEventStream(resp.Body, resp, id, method, result) + } + // A single JSON body. + body, err := io.ReadAll(io.LimitReader(resp.Body, maxHTTPResponse)) + if err != nil { + return fmt.Errorf("mcp http %s read: %w", method, err) + } + var rr rpcResponse + if err := json.Unmarshal(bytes.TrimSpace(body), &rr); err != nil { + return fmt.Errorf("mcp http %s decode: %w", method, err) + } + if !decodeIDString(rr, id) { + return fmt.Errorf("mcp http %s: %w", method, ErrClosed) + } + return s.applyResponse(rr, method, result) +} + +// readEventStream parses an SSE stream and stops once the response matching id is +// received. Notifications (no id) are routed to the handler. +func (s *streamableHTTPTransport) readEventStream(r io.Reader, resp *http.Response, id, method string, result any) error { + br := bufio.NewReader(r) + var event, data string + for { + line, err := br.ReadString('\n') + if err != nil { + if ctxErr := s.cancelErr(resp); ctxErr != nil && ctxErr != context.Canceled { + return fmt.Errorf("mcp http %s: %w", method, ctxErr) + } + return fmt.Errorf("mcp http %s: %w", method, ErrClosed) + } + line = strings.TrimRight(line, "\r\n") + switch { + case line == "": + if data != "" { + stop, herr := s.handleEvent(event, data, id, method, result) + if herr != nil { + return herr + } + if stop { + return nil + } + } + event, data = "", "" + case strings.HasPrefix(line, "event:"): + event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + d := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" { + data = d + } else { + data += "\n" + d + } + } + } +} + +// handleEvent dispatches a single SSE data payload. It returns true when the +// event is the response being awaited, so the reader can stop. +func (s *streamableHTTPTransport) handleEvent(event, data, id, method string, result any) (bool, error) { + if data == "" { + return false, nil + } + var resp rpcResponse + if err := json.Unmarshal([]byte(data), &resp); err != nil { + return false, nil + } + if resp.ID == nil { + s.routeNotification(resp.Method) + return false, nil + } + if decodeIDString(resp, id) { + if err := s.applyResponse(resp, method, result); err != nil { + return false, err + } + return true, nil + } + return false, nil +} + +// routeNotification hands a server-initiated notification to the client handler. +func (s *streamableHTTPTransport) routeNotification(method string) { + if method == "" { + return + } + s.mu.Lock() + fn := s.onNotification + s.mu.Unlock() + if fn != nil { + fn(method) + } +} + +// cancelErr reports whether the response's request context was cancelled. +func (s *streamableHTTPTransport) cancelErr(resp *http.Response) error { + if resp.Request != nil && resp.Request.Context().Err() != nil { + return resp.Request.Context().Err() + } + return nil +} + +// applyResponse decodes a JSON-RPC response into result, surfacing an error +// object if present. +func (s *streamableHTTPTransport) applyResponse(resp rpcResponse, method string, result any) error { + if resp.Error != nil { + return fmt.Errorf("mcp http %s: %w: %w", method, ErrRPC, resp.Error) + } + if result != nil { + if err := json.Unmarshal(resp.Result, result); err != nil { + return fmt.Errorf("mcp http %s decode: %w", method, err) + } + } + return nil +} + +// notify POSTs a fire-and-forget JSON-RPC notification (no id). The server +// typically answers with 202 Accepted and no body. +func (s *streamableHTTPTransport) notify(ctx context.Context, method string, params any) error { + s.mu.Lock() + if !s.started { + s.mu.Unlock() + return fmt.Errorf("mcp notify %s: %w", method, ErrNotConnected) + } + s.mu.Unlock() + reqBody, err := encodeRequest(nil, method, params) + if err != nil { + return fmt.Errorf("mcp notify %s: %w", method, err) + } + resp, err := s.post(ctx, reqBody) + if err != nil { + return fmt.Errorf("mcp notify %s: %w", method, err) + } + defer resp.Body.Close() + s.captureSession(resp) + _, _ = io.Copy(io.Discard, resp.Body) + return nil +} + +// readStandalone reads the standalone GET SSE stream for server-initiated +// messages until the connection is closed or the stream terminates. +func (s *streamableHTTPTransport) readStandalone(ctx context.Context) { + for { + if ctx.Err() != nil { + return + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil) + if err != nil { + return + } + req.Header.Set("Accept", "text/event-stream") + s.mu.Lock() + session := s.session + proto := "" + if s.protocol != nil { + proto = s.protocol() + } + s.mu.Unlock() + if session != "" { + req.Header.Set(mcpSessionIDHeader, session) + } + if proto != "" { + req.Header.Set("Mcp-Protocol-Version", proto) + } + resp, err := s.http.Do(req) + if err != nil { + if ctx.Err() != nil { + return + } + s.log.Debug("mcp http get stream error", "name", s.cfg.Name, "err", err) + s.sleepBackoff(ctx) + continue + } + if resp.StatusCode != http.StatusOK { + _ = resp.Body.Close() + s.log.Debug("mcp http get stream not accepted", "name", s.cfg.Name, "status", resp.StatusCode) + return + } + s.readStandaloneEvents(ctx, resp) + if ctx.Err() != nil { + return + } + s.sleepBackoff(ctx) + } +} + +// readStandaloneEvents consumes one GET stream, routing notifications until it +// ends or the connection is closed. +func (s *streamableHTTPTransport) readStandaloneEvents(ctx context.Context, resp *http.Response) { + stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() }) + defer stopClose() + br := bufio.NewReader(resp.Body) + var event, data string + for { + line, err := br.ReadString('\n') + if err != nil { + return + } + line = strings.TrimRight(line, "\r\n") + switch { + case line == "": + if data != "" { + s.handleEvent(event, data, "", "", nil) + } + event, data = "", "" + case strings.HasPrefix(line, "event:"): + event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + d := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" { + data = d + } else { + data += "\n" + d + } + } + } +} + +// sleepBackoff waits between standalone stream reconnect attempts, exiting early +// when the connection is closed. +func (s *streamableHTTPTransport) sleepBackoff(ctx context.Context) { + select { + case <-time.After(time.Second): + case <-ctx.Done(): + } +} + +// Close closes the standalone stream and marks the transport closed. It is +// idempotent. +func (s *streamableHTTPTransport) Close() error { + s.closeOnce.Do(func() { + close(s.closeCh) + s.mu.Lock() + if s.getCancel != nil { + s.getCancel() + } + s.mu.Unlock() + }) + return nil +} diff --git a/internal/capability/mcp/http_test.go b/internal/capability/mcp/http_test.go new file mode 100644 index 0000000..ebbb258 --- /dev/null +++ b/internal/capability/mcp/http_test.go @@ -0,0 +1,197 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +// httpTestServer is a Streamable HTTP test server. Its GET handler keeps an SSE +// stream open until the request context is done (so it does not leak), while POST +// handling is supplied by postFn. +func httpTestServer(t *testing.T, postFn func(w http.ResponseWriter, r *http.Request)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + <-r.Context().Done() + return + } + postFn(w, r) + })) +} + +// httpJSONResponse writes a single JSON-RPC response echo back, using the request +// id so the client matches it. +func httpJSONResponse(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req rpcRequest + _ = json.Unmarshal(body, &req) + resp := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"direct"}]}}` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(resp)) +} + +func TestHTTP_Call_DirectJSON(t *testing.T) { + t.Parallel() + srv := httpTestServer(t, httpJSONResponse) + defer srv.Close() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) + defer tr.Close() + if err := tr.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + var res callToolResult + if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil { + t.Fatalf("Call: %v", err) + } + if len(res.Content) != 1 || res.Content[0].Text != "direct" { + t.Fatalf("Content = %+v, want [direct]", res.Content) + } +} + +func TestHTTP_Call_EventStream(t *testing.T) { + t.Parallel() + srv := httpTestServer(t, func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req rpcRequest + _ = json.Unmarshal(body, &req) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f, _ := w.(http.Flusher) + msg := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"stream"}]}}` + _, _ = w.Write([]byte("event: message\ndata: " + msg + "\n\n")) + f.Flush() + }) + defer srv.Close() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) + defer tr.Close() + if err := tr.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + var res callToolResult + if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil { + t.Fatalf("Call: %v", err) + } + if len(res.Content) != 1 || res.Content[0].Text != "stream" { + t.Fatalf("Content = %+v, want [stream]", res.Content) + } +} + +func TestHTTP_ServerNotification(t *testing.T) { + t.Parallel() + notif := make(chan string, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f, _ := w.(http.Flusher) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\n\n")) + f.Flush() + <-r.Context().Done() + return + } + w.WriteHeader(http.StatusAccepted) + })) + defer srv.Close() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) + tr.setOnNotification(func(m string) { notif <- m }) + defer tr.Close() + if err := tr.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + select { + case m := <-notif: + if m != notificationToolsListChanged { + t.Fatalf("notification = %q, want %q", m, notificationToolsListChanged) + } + case <-time.After(time.Second): + t.Fatal("no server notification received") + } +} + +func TestHTTP_SendsProtocolVersionHeader(t *testing.T) { + t.Parallel() + got := make(chan string, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + <-r.Context().Done() + return + } + if v := r.Header.Get("Mcp-Protocol-Version"); v != "" { + select { + case got <- v: + default: + } + } + httpJSONResponse(w, r) + })) + defer srv.Close() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) + tr.setProtocolProvider(func() string { return "2025-11-25" }) + defer tr.Close() + if err := tr.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + if err := tr.Call(context.Background(), "tools/list", map[string]any{}, &listToolsResult{}); err != nil { + t.Fatalf("Call: %v", err) + } + select { + case v := <-got: + if v != "2025-11-25" { + t.Fatalf("protocol header = %q, want 2025-11-25", v) + } + case <-time.After(time.Second): + t.Fatal("Mcp-Protocol-Version header not sent") + } +} + +func TestHTTP_Call_NotStarted(t *testing.T) { + t.Parallel() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger()) + err := tr.Call(context.Background(), "tools/list", nil, nil) + if !errors.Is(err, ErrNotConnected) { + t.Fatalf("Call error = %v, want ErrNotConnected", err) + } +} + +func TestHTTP_Start_NoURL(t *testing.T) { + t.Parallel() + tr := newHTTPTransport(MCPConfig{Name: "s"}, defaultLogger()) + if err := tr.Start(context.Background()); !errors.Is(err, ErrNoTransport) { + t.Fatalf("Start error = %v, want ErrNoTransport", err) + } +} + +func TestHTTP_Start_InvalidURL(t *testing.T) { + t.Parallel() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: "ftp://x"}, defaultLogger()) + if err := tr.Start(context.Background()); err == nil { + t.Fatal("Start accepted a non-http url, want error") + } +} + +func TestHTTP_Close_Idempotent(t *testing.T) { + t.Parallel() + tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger()) + if err := tr.Close(); err != nil { + t.Fatalf("first Close: %v", err) + } + if err := tr.Close(); err != nil { + t.Fatalf("second Close: %v", err) + } +} diff --git a/internal/capability/mcp/mcp.go b/internal/capability/mcp/mcp.go index 4599280..1b3dcfc 100644 --- a/internal/capability/mcp/mcp.go +++ b/internal/capability/mcp/mcp.go @@ -2,11 +2,18 @@ // auto-registration glue that surfaces MCP server tools as capabilities. // // An MCP server is an external tool provider. This package speaks JSON-RPC 2.0 -// to it over stdio (a spawned subprocess) or SSE, discovers the tools it -// exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and registers -// them into the capability registry so any character can invoke them +// to it over stdio (a spawned subprocess), SSE or Streamable HTTP, discovers the +// tools it exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and +// registers them into the capability registry so any character can invoke them // (DevDocs/design/09 插件生态 §1). // +// The client aligns with the MCP Go SDK's mature behaviour: it negotiates the +// protocol version during initialize (requesting the newest version first and +// downgrading to a mutually supported one), subscribes to +// notifications/tools/list_changed via its capabilities so the registry can +// re-sync its tool set, and runs an optional ping keepalive that closes a dead +// session so the host can rebuild it. +// // It is deliberately dependency-light: only the standard library and the // project packages (plugins/sdk, capability/registry, internal/logger) are // used. The wire protocol is implemented directly, so no third-party MCP SDK is @@ -30,14 +37,78 @@ import ( "errors" "io" "log/slog" + "slices" "sync" + "sync/atomic" + "time" "git.yeij.top/AskaEth/Cyrene_v2/internal/logger" ) -// protocolVersion is the MCP protocol version this client advertises during the -// initialize handshake. -const protocolVersion = "2024-11-05" +// MCP protocol versions this client can speak, newest first. The client advertises +// its preferred version (latestProtocolVersion) during initialize and accepts a +// compatible version from the server, downgrading to a mutually supported version +// when the server rejects or cannot match the request (DevDocs/design/09 §1.1). +const ( + // protocolVersion20251125 is the current preferred protocol version (the + // MCP 2025-11-25 revision that introduced the Streamable HTTP migration). + protocolVersion20251125 = "2025-11-25" + // protocolVersion20241105 is the legacy protocol version, retained for + // servers that have not yet migrated. + protocolVersion20241105 = "2024-11-05" + // latestProtocolVersion is what the client requests first during initialize. + latestProtocolVersion = protocolVersion20251125 + // codeUnsupportedProtocolVersion is the JSON-RPC error code a server returns + // during initialize when it does not support the requested protocol version + // and wishes to advertise the versions it does support. + codeUnsupportedProtocolVersion = -32022 + // codeMethodNotFound is the JSON-RPC error code a server returns when it does + // not implement a requested method (e.g. a keepalive ping). + codeMethodNotFound = -32601 + // maxInitializeAttempts bounds the initialize version-negotiation retries + // (the first attempt plus one downgrade). + maxInitializeAttempts = 2 + // notificationToolsListChanged is the server→client notification signalling + // that the tool list has changed. + notificationToolsListChanged = "notifications/tools/list_changed" + // keepAliveFailureThreshold is the number of consecutive ping failures the + // keepalive loop tolerates before closing the connection for a rebuild. + keepAliveFailureThreshold = 3 +) + +// supportedProtocolVersions lists every protocol version this client understands, +// newest first. Version negotiation picks the first entry that the server also +// supports. +var supportedProtocolVersions = []string{protocolVersion20251125, protocolVersion20241105} + +// isSupportedVersion reports whether v is a protocol version this client speaks. +func isSupportedVersion(v string) bool { + return slices.Contains(supportedProtocolVersions, v) +} + +// negotiateMutuallySupportedVersion returns the highest protocol version that is +// supported by both this client and the server (which supports supported). It +// returns "" when there is no overlap. +func negotiateMutuallySupportedVersion(supported []string) string { + for _, v := range supportedProtocolVersions { + if slices.Contains(supported, v) { + return v + } + } + return "" +} + +// highestCommonVersion returns the highest protocol version this client supports +// that is not newer than v (used to downgrade when the server picks a version the +// client cannot speak). It returns "" when no acceptable version exists. +func highestCommonVersion(v string) string { + for _, s := range supportedProtocolVersions { + if s <= v { + return s + } + } + return "" +} // Sentinel errors returned by this package. Test with errors.Is. var ( @@ -54,10 +125,16 @@ var ( // object. The concrete server error is wrapped and accessible via // errors.Is/errors.As on an *rpcError. ErrRPC = errors.New("mcp: rpc error") + // ErrUnsupportedProtocolVersion is returned when the client and server + // cannot agree on a mutually supported protocol version after the + // negotiation retries are exhausted. + ErrUnsupportedProtocolVersion = errors.New("mcp: unsupported protocol version") ) -// MCPConfig describes a single MCP server connection. Exactly one transport is -// selected: stdio when Command is set, SSE when URL is set. +// MCPConfig describes a single MCP server connection. The transport is selected +// by TransportType when set, otherwise inferred from the fields: stdio when +// Command is set, SSE when URL is set. A blank TransportType keeps that legacy +// inference so existing configuration keeps working unchanged. type MCPConfig struct { // Name is the human-readable label of the server, used in logs and as the // prefix for registering its tools uniquely. @@ -66,8 +143,12 @@ type MCPConfig struct { Command string `env:"MCP_COMMAND"` // Args are the arguments passed to Command. Args []string `env:"MCP_ARGS"` - // URL is the SSE endpoint of an MCP server. + // URL is the endpoint of an MCP server (SSE stream endpoint, or the + // Streamable HTTP endpoint when TransportType is "http"). URL string `env:"MCP_URL"` + // TransportType selects the transport explicitly: "stdio", "sse" or + // "http" (Streamable HTTP). Empty means infer from Command/URL. + TransportType string `env:"MCP_TRANSPORT"` // Enabled reports whether this server should be activated at startup. Enabled bool `env:"MCP_ENABLED"` } @@ -122,6 +203,53 @@ type Client struct { mu sync.Mutex state state info ServerInfo + + // negotiatedVersion is the protocol version agreed during initialize; it is + // sent back to the server in the Mcp-Protocol-Version header by the HTTP + // transport. + negotiatedVersion string + + // preferredProtocolVersion, when non-empty, overrides latestProtocolVersion + // as the version requested first during initialize. Tests use it to drive the + // negotiation path deterministically. + preferredProtocolVersion string + + // keepAlive is the ping interval for the connection keepalive; zero disables + // it. keepAliveFailures is the tolerated consecutive-failure count. + keepAlive time.Duration + keepAliveFailures int + keepaliveCancel context.CancelFunc + + // onToolListChanged is invoked (in a fresh goroutine) when the server sends + // notifications/tools/list_changed, so the host can refresh its tool set. + onToolListChanged func() + refreshPending atomic.Bool +} + +// protocolVersion returns the protocol version negotiated by initialize, or "" +// before a successful handshake. +func (c *Client) protocolVersion() string { + c.mu.Lock() + defer c.mu.Unlock() + return c.negotiatedVersion +} + +// setProtocolVersion records the version negotiated by initialize. +func (c *Client) setProtocolVersion(v string) { + c.mu.Lock() + defer c.mu.Unlock() + c.negotiatedVersion = v +} + +// capabilities returns the client capability declaration for the given protocol +// version. It advertises tools.listChanged so the server notifies the client when +// the tool list changes. +func (c *Client) capabilities(version string) map[string]any { + caps := map[string]any{} + if c.onToolListChanged != nil { + caps["tools"] = map[string]any{"listChanged": true} + } + return caps } // state tracks the client lifecycle. @@ -134,8 +262,8 @@ const ( ) // Transport is the JSON-RPC message exchange with an MCP server over a concrete -// transport (stdio subprocess or SSE). Implementations must be safe for -// concurrent use. This is exported so tests and hosts can inject a fake +// transport (stdio subprocess, SSE or Streamable HTTP). Implementations must be +// safe for concurrent use. This is exported so tests and hosts can inject a fake // transport. type Transport interface { // Start makes the transport ready to exchange messages. @@ -154,6 +282,37 @@ type notifier interface { notify(ctx context.Context, method string, params any) error } +// serverNotifier is implemented by transports that can deliver server-initiated +// notifications (messages with no id, e.g. notifications/tools/list_changed) back +// to the client. The handler is invoked from the transport's read path, so it must +// not block or re-enter the same transport synchronously. +type serverNotifier interface { + setOnNotification(fn func(method string)) +} + +// attachNotificationSink wires a server→client notification handler into tr, if +// tr supports one. +func attachNotificationSink(tr Transport, fn func(method string)) { + if on, ok := tr.(serverNotifier); ok { + on.setOnNotification(fn) + } +} + +// protocolProvider is implemented by transports that include the negotiated MCP +// protocol version in requests (the Mcp-Protocol-Version header of the HTTP +// transport). +type protocolProvider interface { + setProtocolProvider(fn func() string) +} + +// attachProtocolProvider wires the negotiated-version getter into tr, if tr +// supports one. +func attachProtocolProvider(tr Transport, fn func() string) { + if p, ok := tr.(protocolProvider); ok { + p.setProtocolProvider(fn) + } +} + // Option customizes a Client built by New. type Option func(*Client) @@ -176,6 +335,40 @@ func WithTransport(t Transport) Option { } } +// WithKeepAlive enables the connection keepalive, sending an MCP ping every +// interval and closing the session after keepAliveFailures consecutive missed +// pings. A zero interval leaves keepalive disabled. When keepAliveFailures is +// omitted, the package default (keepAliveFailureThreshold) is used. +func WithKeepAlive(interval time.Duration, keepAliveFailures ...int) Option { + failures := keepAliveFailureThreshold + if len(keepAliveFailures) > 0 && keepAliveFailures[0] > 0 { + failures = keepAliveFailures[0] + } + return func(c *Client) { + c.keepAlive = interval + c.keepAliveFailures = failures + } +} + +// WithToolListChangedHandler registers a handler that is invoked (in a fresh +// goroutine) whenever the server sends notifications/tools/list_changed, letting +// the host refresh its registered tool set. +func WithToolListChangedHandler(fn func()) Option { + return func(c *Client) { + c.onToolListChanged = fn + } +} + +// WithProtocolVersion overrides the protocol version requested first during +// initialize. It is primarily for tests that drive the negotiation path +// deterministically; callers should normally leave it unset to use +// latestProtocolVersion. +func WithProtocolVersion(v string) Option { + return func(c *Client) { + c.preferredProtocolVersion = v + } +} + // New builds a Client for cfg. By default it uses a discard logger and selects a // stdio or SSE transport from cfg; pass WithLogger and/or WithTransport to // override. diff --git a/internal/capability/mcp/registry.go b/internal/capability/mcp/registry.go index 5e1e957..8b1e92a 100644 --- a/internal/capability/mcp/registry.go +++ b/internal/capability/mcp/registry.go @@ -24,10 +24,15 @@ type Registry struct { // NewRegistry builds a Registry that registers cfg's tools into reg. Options // apply to the underlying Client (logger, injected transport). It is named // NewRegistry rather than New because this package already exports New for -// Client. +// Client. A tool-list-changed handler is wired automatically so the registry +// re-syncs when the server announces tools/list_changed. func NewRegistry(reg *registry.Registry, cfg MCPConfig, opts ...Option) *Registry { + r := &Registry{reg: reg} + opts = append(opts, WithToolListChangedHandler(r.refresh)) c := New(cfg, opts...) - return &Registry{reg: reg, mcp: c, log: c.log} + r.mcp = c + r.log = c.log + return r } // Register connects the MCP server, discovers its tools, and registers each one @@ -59,6 +64,69 @@ func (r *Registry) Register(ctx context.Context) error { return nil } +// refresh re-syncs the registered tool set with the server after a +// notifications/tools/list_changed notification. It re-discovers the current +// tools and applies only the diff: it removes names that are gone and registers +// newly appeared ones, so tools registered by other providers (which collide by +// name) are preserved. It is invoked (off the transport read path) by the +// client's tool-list-changed handler and coalesced by the client. +func (r *Registry) refresh() { + if r.reg == nil { + return + } + tools, err := r.mcp.DiscoverTools(context.Background()) + if err != nil { + r.log.Warn("mcp tool refresh failed", "name", r.mcp.cfg.Name, "err", err) + return + } + + incoming := make(map[string]bool, len(tools)) + for _, t := range tools { + incoming[t.Name()] = true + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return + } + + // Drop names the server no longer exposes. Tools registered by other + // providers were skipped at registration time and are not in r.names, so they + // are never removed here. + kept := r.names[:0] + for _, name := range r.names { + if !incoming[name] { + if err := r.reg.Remove(name); err != nil && !errors.Is(err, registry.ErrNotFound) { + r.log.Warn("mcp refresh remove failed", "name", name, "err", err) + } + continue + } + kept = append(kept, name) + } + r.names = kept + + // Register newly appeared tools. Skips names already present (own or + // foreign) so a name collision with another provider is not doubled. + existing := make(map[string]bool, len(r.names)) + for _, name := range r.names { + existing[name] = true + } + for _, t := range tools { + if existing[t.Name()] { + continue + } + if err := r.reg.RegisterTool(t); err != nil { + if !errors.Is(err, registry.ErrDuplicate) { + r.log.Warn("mcp refresh register failed", "name", t.Name(), "err", err) + } + continue + } + r.names = append(r.names, t.Name()) + } + r.log.Info("mcp tools refreshed", "name", r.mcp.cfg.Name, "count", len(r.names)) +} + // Unregister removes every tool this registry previously registered and closes // the underlying client. It is safe to call multiple times. func (r *Registry) Unregister(ctx context.Context) error { diff --git a/internal/capability/mcp/registry_test.go b/internal/capability/mcp/registry_test.go index fa6cfda..6f4b255 100644 --- a/internal/capability/mcp/registry_test.go +++ b/internal/capability/mcp/registry_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "testing" + "time" "git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry" ) @@ -107,6 +108,58 @@ func TestRegistry_Unregister_Idempotent(t *testing.T) { } } +func TestRegistry_ToolListChangedRefreshes(t *testing.T) { + t.Parallel() + ft := &fakeTransport{tools: sampleTools()} + reg := registry.New() + r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft)) + if err := r.Register(context.Background()); err != nil { + t.Fatalf("Register: %v", err) + } + // The server's tool list changes: sub is dropped, mul appears. + ft.setTools([]ToolInfo{ + {Name: "add", Description: "Adds numbers", InputSchema: map[string]any{"type": "object"}}, + {Name: "mul", Description: "Multiplies numbers", InputSchema: map[string]any{"type": "object"}}, + }) + ft.notifyServer(notificationToolsListChanged) + + waitFor(t, time.Second, func() bool { + _, err := reg.Get("mul") + return err == nil + }) + if _, err := reg.Get("add"); err != nil { + t.Fatalf("Get(add) after refresh = %v, want present", err) + } + if _, err := reg.Get("sub"); !errors.Is(err, registry.ErrNotFound) { + t.Fatalf("Get(sub) after refresh = %v, want ErrNotFound", err) + } +} + +func TestRegistry_ToolListChangedPreservesForeignTools(t *testing.T) { + t.Parallel() + ft := &fakeTransport{tools: sampleTools()} + reg := registry.New() + _ = reg.RegisterTool(&testTool{name: "add"}) // a foreign 'add' before refresh + r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft)) + if err := r.Register(context.Background()); err != nil { + // 'add' collides with the foreign tool; Register must surface the dup. + if !errors.Is(err, registry.ErrDuplicate) { + t.Fatalf("Register error = %v, want registry.ErrDuplicate", err) + } + } + // After the collision, sub is registered, 'add' stays foreign. + ft.setTools([]ToolInfo{{Name: "sub", Description: "Subtracts", InputSchema: map[string]any{}}}) + ft.notifyServer(notificationToolsListChanged) + waitFor(t, time.Second, func() bool { + _, err := reg.Get("sub") + return err == nil + }) + // The foreign 'add' must be preserved even though it is not in the server list. + if _, err := reg.Get("add"); err != nil { + t.Fatalf("Get(add) after refresh = %v, want foreign tool preserved", err) + } +} + // testTool is a minimal sdk.Tool used to pre-register a colliding name. type testTool struct { name string diff --git a/internal/capability/mcp/rpc.go b/internal/capability/mcp/rpc.go index 8918a45..0406d71 100644 --- a/internal/capability/mcp/rpc.go +++ b/internal/capability/mcp/rpc.go @@ -15,10 +15,13 @@ type rpcRequest struct { } // rpcResponse is a JSON-RPC 2.0 response. Result is kept raw until the caller -// decodes it into the concrete result type. +// decodes it into the concrete result type. Method is populated only for +// server-initiated notifications (no id), so transports can route them to the +// client's notification handler. type rpcResponse struct { JSONRPC string `json:"jsonrpc"` ID any `json:"id,omitempty"` + Method string `json:"method,omitempty"` Result json.RawMessage `json:"result,omitempty"` Error *rpcError `json:"error,omitempty"` } diff --git a/internal/capability/mcp/sse.go b/internal/capability/mcp/sse.go index d98560d..83f6384 100644 --- a/internal/capability/mcp/sse.go +++ b/internal/capability/mcp/sse.go @@ -19,18 +19,27 @@ import ( // requests to the message endpoint (discovered via an endpoint event, falling // back to the configured URL). type sseTransport struct { - cfg MCPConfig - log *slog.Logger - url string - endpoint string - mu sync.Mutex - nextID int - pending map[string]chan []byte - resp *http.Response - body io.ReadCloser - closeCh chan struct{} - closeOnce sync.Once - started bool + cfg MCPConfig + log *slog.Logger + url string + endpoint string + mu sync.Mutex + nextID int + pending map[string]chan []byte + resp *http.Response + body io.ReadCloser + closeCh chan struct{} + closeOnce sync.Once + started bool + onNotification func(method string) +} + +// setOnNotification records the handler invoked when the server sends a +// server-initiated notification (a message with no id) over the SSE stream. +func (s *sseTransport) setOnNotification(fn func(method string)) { + s.mu.Lock() + defer s.mu.Unlock() + s.onNotification = fn } // maxSSEResponse bounds how much of a POST reply body is read before falling @@ -120,6 +129,19 @@ func (s *sseTransport) dispatch(event, data string) { if err := json.Unmarshal([]byte(data), &resp); err != nil { return } + if resp.ID == nil { + // A server-initiated notification; hand it to the client's handler. + // The handler must not re-enter this transport synchronously. + if resp.Method != "" { + s.mu.Lock() + fn := s.onNotification + s.mu.Unlock() + if fn != nil { + fn(resp.Method) + } + } + return + } id, ok := resp.ID.(string) if !ok { return diff --git a/internal/capability/mcp/stdio.go b/internal/capability/mcp/stdio.go index 7154f18..aca18f1 100644 --- a/internal/capability/mcp/stdio.go +++ b/internal/capability/mcp/stdio.go @@ -17,13 +17,24 @@ import ( // stdioTransport runs an MCP server as a child process and speaks // newline-delimited JSON-RPC 2.0 to its stdin/stdout. type stdioTransport struct { - cfg MCPConfig - log *slog.Logger - line *lineRPC - cmd *exec.Cmd - stdin io.WriteCloser - closeOnce sync.Once - closed bool + cfg MCPConfig + log *slog.Logger + line *lineRPC + cmd *exec.Cmd + stdin io.WriteCloser + closeOnce sync.Once + closed bool + onNotification func(method string) +} + +// setOnNotification records the handler invoked when the server sends a +// notification frame (a JSON-RPC message with no id). It is wired by the client +// before initialize so server notifications are not lost. +func (s *stdioTransport) setOnNotification(fn func(method string)) { + s.onNotification = fn + if s.line != nil { + s.line.onNotification = fn + } } // newStdioTransport builds a stdio transport for cfg. @@ -57,7 +68,7 @@ func (s *stdioTransport) Start(ctx context.Context) error { } s.cmd = cmd s.stdin = stdin - s.line = &lineRPC{in: stdin, out: bufio.NewReader(stdout), log: s.log} + s.line = &lineRPC{in: stdin, out: bufio.NewReader(stdout), log: s.log, onNotification: s.onNotification} go s.drainStderr(stderr) return nil } @@ -106,11 +117,12 @@ func (s *stdioTransport) drainStderr(r io.Reader) { // serialized so at most one request is in flight at a time, which keeps id // matching simple and deterministic. type lineRPC struct { - in io.Writer - out *bufio.Reader - log *slog.Logger - mu sync.Mutex - nextID int + in io.Writer + out *bufio.Reader + log *slog.Logger + mu sync.Mutex + nextID int + onNotification func(method string) } // call writes a request and reads until the matching response arrives. @@ -151,6 +163,12 @@ func (l *lineRPC) readResponse(id, method string, result any) error { continue } if resp.ID == nil { + // A server-initiated notification (no id); hand it to the client's + // handler if one is installed. The handler must not re-enter this + // transport synchronously (see the serverNotifier contract). + if resp.Method != "" && l.onNotification != nil { + l.onNotification(resp.Method) + } continue } if !decodeIDString(resp, id) {