feat(mcp): Top5#3 对齐 MCP Go SDK——协议版本协商 + list_changed 差量刷新 + keepalive + Streamable HTTP
- 协议版本协商: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:
1 parent
464ed2cad3
commit
b0bd93b15c
13 files changed
+1479
-28
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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},
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user