feat(mcp): Top5#3 对齐 MCP Go SDK——协议版本协商 + list_changed 差量刷新 + keepalive + Streamable HTTP
CI / Backend (build / vet / test / cross-compile) (push) Has been cancelled
CI / Frontend (tsc + vite build) (push) Has been cancelled

- 协议版本协商: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 交叉编译。
This commit is contained in:
Aska committed 2026-09-06 17:27:14 +08:00
1 parent 464ed2cad3
commit b0bd93b15c
13 files changed
+1479 -28

No files matched your search

+190 -9
View File
@@ -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) {
sent := c.preferredProtocolVersion
if sent == "" {
sent = latestProtocolVersion
}
var lastErr error
for attempt := 0; attempt < maxInitializeAttempts; attempt++ {
params := initializeParams{
ProtocolVersion: protocolVersion,
Capabilities: map[string]any{},
ProtocolVersion: sent,
Capabilities: c.capabilities(sent),
ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"},
}
var res initializeResult
if err := tr.Call(ctx, "initialize", params, &res); err != nil {
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{}, 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
+2
View File
@@ -74,6 +74,7 @@ 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
}
@@ -82,6 +83,7 @@ func readOne(env map[string]string, prefix string) (MCPConfig, bool) {
Command: command,
Args: splitArgs(env[p+"_ARGS"]),
URL: url,
TransportType: transp,
Enabled: parseBool(env[p+"_ENABLED"], false),
}, true
}
+19
View File
@@ -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"})
@@ -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")
}
+93 -4
View File
@@ -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
// 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},
})
}
+412
View File
@@ -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
}
+197
View File
@@ -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)
}
}
+204 -11
View File
@@ -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.
+70 -2
View File
@@ -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 {
+53
View File
@@ -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
+4 -1
View File
@@ -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"`
}
+22
View File
@@ -31,6 +31,15 @@ type sseTransport struct {
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
+19 -1
View File
@@ -24,6 +24,17 @@ type stdioTransport struct {
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
}
@@ -111,6 +122,7 @@ type lineRPC struct {
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) {