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
+1517
-66
No files matched your search
@@ -2,8 +2,11 @@ package mcp
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
|
"time"
|
||||||
|
|
||||||
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
)
|
)
|
||||||
@@ -61,11 +64,17 @@ func (c *Client) Connect(ctx context.Context) error {
|
|||||||
tr.Close()
|
tr.Close()
|
||||||
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
|
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)
|
info, err := c.initialize(ctx, tr)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
tr.Close()
|
tr.Close()
|
||||||
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
|
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.sendInitialized(ctx, tr)
|
||||||
|
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
@@ -73,7 +82,9 @@ func (c *Client) Connect(ctx context.Context) error {
|
|||||||
c.info = info
|
c.info = info
|
||||||
c.state = stateConnected
|
c.state = stateConnected
|
||||||
c.mu.Unlock()
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -87,19 +98,94 @@ func (c *Client) Initialize(ctx context.Context) (ServerInfo, error) {
|
|||||||
return c.initialize(ctx, tr)
|
return c.initialize(ctx, tr)
|
||||||
}
|
}
|
||||||
|
|
||||||
// initialize sends the initialize request over tr and records the server
|
// initialize performs the MCP initialize handshake with protocol version
|
||||||
// identity.
|
// 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) {
|
func (c *Client) initialize(ctx context.Context, tr Transport) (ServerInfo, error) {
|
||||||
params := initializeParams{
|
sent := c.preferredProtocolVersion
|
||||||
ProtocolVersion: protocolVersion,
|
if sent == "" {
|
||||||
Capabilities: map[string]any{},
|
sent = latestProtocolVersion
|
||||||
ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"},
|
|
||||||
}
|
}
|
||||||
var res initializeResult
|
var lastErr error
|
||||||
if err := tr.Call(ctx, "initialize", params, &res); err != nil {
|
for attempt := 0; attempt < maxInitializeAttempts; attempt++ {
|
||||||
return ServerInfo{}, err
|
params := initializeParams{
|
||||||
|
ProtocolVersion: sent,
|
||||||
|
Capabilities: c.capabilities(sent),
|
||||||
|
ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"},
|
||||||
|
}
|
||||||
|
var res initializeResult
|
||||||
|
err := tr.Call(ctx, "initialize", params, &res)
|
||||||
|
if err != nil {
|
||||||
|
// Server rejected the requested version and may advertise the versions
|
||||||
|
// it supports; downgrade to the highest mutually supported one.
|
||||||
|
if supported, ok := unsupportedVersions(err); ok {
|
||||||
|
if v := negotiateMutuallySupportedVersion(supported); v != "" && v != sent {
|
||||||
|
c.log.Debug("mcp downgrading protocol version",
|
||||||
|
"name", c.cfg.Name, "requested", sent, "version", v)
|
||||||
|
sent = v
|
||||||
|
lastErr = err
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, err)
|
||||||
|
}
|
||||||
|
return ServerInfo{}, err
|
||||||
|
}
|
||||||
|
negotiated := res.ProtocolVersion
|
||||||
|
if negotiated == "" {
|
||||||
|
negotiated = sent
|
||||||
|
}
|
||||||
|
if !isSupportedVersion(negotiated) {
|
||||||
|
// The server replied with a version this client does not speak;
|
||||||
|
// downgrade once to the highest version we both accept.
|
||||||
|
if v := highestCommonVersion(negotiated); v != "" && v != sent {
|
||||||
|
c.log.Debug("mcp renegotiating protocol version",
|
||||||
|
"name", c.cfg.Name, "server", negotiated, "version", v)
|
||||||
|
sent = v
|
||||||
|
lastErr = err
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: server advertised %q", c.cfg.Name, ErrUnsupportedProtocolVersion, negotiated)
|
||||||
|
}
|
||||||
|
c.setProtocolVersion(negotiated)
|
||||||
|
return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil
|
||||||
}
|
}
|
||||||
return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil
|
return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, lastErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
// isMethodNotFound reports whether err wraps a JSON-RPC "method not found" error
|
||||||
|
// (code -32601), which a peer returns when it does not implement a method such as
|
||||||
|
// the keepalive ping.
|
||||||
|
func isMethodNotFound(err error) bool {
|
||||||
|
var re *rpcError
|
||||||
|
return errors.As(err, &re) && re.Code == codeMethodNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
// unsupportedVersions reports the protocol versions the server supports, extracted
|
||||||
|
// from the data field of a codeUnsupportedProtocolVersion JSON-RPC error. It
|
||||||
|
// reports ok=false when err does not carry such an error or the data is unusable.
|
||||||
|
func unsupportedVersions(err error) ([]string, bool) {
|
||||||
|
var re *rpcError
|
||||||
|
if !errors.As(err, &re) || re.Code != codeUnsupportedProtocolVersion {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
if re.Data == nil {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(re.Data)
|
||||||
|
if err != nil {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
var payload struct {
|
||||||
|
Supported []string `json:"supported"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(data, &payload); err != nil || len(payload.Supported) == 0 {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
return payload.Supported, true
|
||||||
}
|
}
|
||||||
|
|
||||||
// sendInitialized best-effort sends the notifications/initialized advisory ack.
|
// 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.
|
// ListTools returns the tools exposed by the connected MCP server.
|
||||||
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
|
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
|
||||||
tr := c.transport()
|
tr := c.transport()
|
||||||
@@ -162,7 +329,8 @@ func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) {
|
|||||||
return out, nil
|
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 {
|
func (c *Client) Close() error {
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
if c.state == stateClosed {
|
if c.state == stateClosed {
|
||||||
@@ -171,6 +339,10 @@ func (c *Client) Close() error {
|
|||||||
}
|
}
|
||||||
c.state = stateClosed
|
c.state = stateClosed
|
||||||
tr := c.tr
|
tr := c.tr
|
||||||
|
if c.keepaliveCancel != nil {
|
||||||
|
c.keepaliveCancel()
|
||||||
|
c.keepaliveCancel = nil
|
||||||
|
}
|
||||||
c.mu.Unlock()
|
c.mu.Unlock()
|
||||||
if tr != nil {
|
if tr != nil {
|
||||||
return tr.Close()
|
return tr.Close()
|
||||||
@@ -178,9 +350,18 @@ func (c *Client) Close() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// newTransport selects the transport matching cfg: stdio when Command is set, SSE
|
// newTransport selects the transport matching cfg. An explicit TransportType
|
||||||
// when URL is set.
|
// ("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) {
|
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 {
|
switch {
|
||||||
case cfg.Command != "":
|
case cfg.Command != "":
|
||||||
return newStdioTransport(cfg, log), nil
|
return newStdioTransport(cfg, log), nil
|
||||||
|
|||||||
@@ -74,15 +74,17 @@ func readOne(env map[string]string, prefix string) (MCPConfig, bool) {
|
|||||||
name := env[p+"_NAME"]
|
name := env[p+"_NAME"]
|
||||||
command := env[p+"_COMMAND"]
|
command := env[p+"_COMMAND"]
|
||||||
url := env[p+"_URL"]
|
url := env[p+"_URL"]
|
||||||
|
transp := env[p+"_TRANSPORT"]
|
||||||
if name == "" && command == "" && url == "" {
|
if name == "" && command == "" && url == "" {
|
||||||
return MCPConfig{}, false
|
return MCPConfig{}, false
|
||||||
}
|
}
|
||||||
return MCPConfig{
|
return MCPConfig{
|
||||||
Name: name,
|
Name: name,
|
||||||
Command: command,
|
Command: command,
|
||||||
Args: splitArgs(env[p+"_ARGS"]),
|
Args: splitArgs(env[p+"_ARGS"]),
|
||||||
URL: url,
|
URL: url,
|
||||||
Enabled: parseBool(env[p+"_ENABLED"], false),
|
TransportType: transp,
|
||||||
|
Enabled: parseBool(env[p+"_ENABLED"], false),
|
||||||
}, true
|
}, 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) {
|
func TestFromEnv_NoServers(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
cfg, err := FromEnv([]string{"MCP_ENABLED=true"})
|
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 (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"fmt"
|
||||||
"sync"
|
"sync"
|
||||||
)
|
)
|
||||||
|
|
||||||
// fakeTransport is a scripted Transport for tests. It records starts, closes and
|
// 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.
|
// 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 {
|
type fakeTransport struct {
|
||||||
startErr error
|
startErr error
|
||||||
closeErr error
|
closeErr error
|
||||||
initErr error
|
initErr error
|
||||||
listErr error
|
listErr error
|
||||||
callErr error
|
callErr error
|
||||||
|
pingErr error
|
||||||
tools []ToolInfo
|
tools []ToolInfo
|
||||||
|
|
||||||
mu sync.Mutex
|
// initErrors is a per-call queue of initialize errors. When non-empty one
|
||||||
starts int
|
// entry is consumed per initialize call; once exhausted the handshake
|
||||||
closes int
|
// succeeds. initErr is returned for every call when initErrors is empty.
|
||||||
methods []string
|
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.
|
// Start records a start.
|
||||||
@@ -26,29 +43,42 @@ func (f *fakeTransport) Start(_ context.Context) error {
|
|||||||
f.mu.Lock()
|
f.mu.Lock()
|
||||||
defer f.mu.Unlock()
|
defer f.mu.Unlock()
|
||||||
f.starts++
|
f.starts++
|
||||||
|
f.activeTools = f.tools
|
||||||
return f.startErr
|
return f.startErr
|
||||||
}
|
}
|
||||||
|
|
||||||
// Call answers the three MCP methods the client issues.
|
// Call answers the MCP methods the client issues.
|
||||||
func (f *fakeTransport) Call(_ context.Context, method string, _ any, result any) error {
|
func (f *fakeTransport) Call(_ context.Context, method string, params any, result any) error {
|
||||||
f.mu.Lock()
|
f.mu.Lock()
|
||||||
defer f.mu.Unlock()
|
defer f.mu.Unlock()
|
||||||
f.methods = append(f.methods, method)
|
f.methods = append(f.methods, method)
|
||||||
switch method {
|
switch method {
|
||||||
case "initialize":
|
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 {
|
if f.initErr != nil {
|
||||||
return f.initErr
|
return f.initErr
|
||||||
}
|
}
|
||||||
if r, ok := result.(*initializeResult); ok {
|
if r, ok := result.(*initializeResult); ok {
|
||||||
r.ServerInfo = serverInfo{Name: "fake-server", Version: "1.0"}
|
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":
|
case "tools/list":
|
||||||
if f.listErr != nil {
|
if f.listErr != nil {
|
||||||
return f.listErr
|
return f.listErr
|
||||||
}
|
}
|
||||||
if r, ok := result.(*listToolsResult); ok {
|
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})
|
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 {
|
if r, ok := result.(*callToolResult); ok {
|
||||||
r.Content = []contentBlock{{Type: "text", Text: "hello"}}
|
r.Content = []contentBlock{{Type: "text", Text: "hello"}}
|
||||||
}
|
}
|
||||||
|
case "ping":
|
||||||
|
f.pings++
|
||||||
|
if f.pingErr != nil {
|
||||||
|
return f.pingErr
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -71,6 +106,25 @@ func (f *fakeTransport) notify(_ context.Context, method string, _ any) error {
|
|||||||
return nil
|
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.
|
// Close records a close.
|
||||||
func (f *fakeTransport) Close() error {
|
func (f *fakeTransport) Close() error {
|
||||||
f.mu.Lock()
|
f.mu.Lock()
|
||||||
@@ -97,3 +151,38 @@ func (f *fakeTransport) hasMethod(method string) bool {
|
|||||||
}
|
}
|
||||||
return false
|
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.
|
// 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
|
// 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
|
// to it over stdio (a spawned subprocess), SSE or Streamable HTTP, discovers the
|
||||||
// exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and registers
|
// tools it exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and
|
||||||
// them into the capability registry so any character can invoke them
|
// registers them into the capability registry so any character can invoke them
|
||||||
// (DevDocs/design/09 插件生态 §1).
|
// (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
|
// It is deliberately dependency-light: only the standard library and the
|
||||||
// project packages (plugins/sdk, capability/registry, internal/logger) are
|
// project packages (plugins/sdk, capability/registry, internal/logger) are
|
||||||
// used. The wire protocol is implemented directly, so no third-party MCP SDK is
|
// used. The wire protocol is implemented directly, so no third-party MCP SDK is
|
||||||
@@ -30,14 +37,78 @@ import (
|
|||||||
"errors"
|
"errors"
|
||||||
"io"
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
|
"slices"
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
|
||||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||||
)
|
)
|
||||||
|
|
||||||
// protocolVersion is the MCP protocol version this client advertises during the
|
// MCP protocol versions this client can speak, newest first. The client advertises
|
||||||
// initialize handshake.
|
// its preferred version (latestProtocolVersion) during initialize and accepts a
|
||||||
const protocolVersion = "2024-11-05"
|
// 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.
|
// Sentinel errors returned by this package. Test with errors.Is.
|
||||||
var (
|
var (
|
||||||
@@ -54,10 +125,16 @@ var (
|
|||||||
// object. The concrete server error is wrapped and accessible via
|
// object. The concrete server error is wrapped and accessible via
|
||||||
// errors.Is/errors.As on an *rpcError.
|
// errors.Is/errors.As on an *rpcError.
|
||||||
ErrRPC = errors.New("mcp: rpc error")
|
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
|
// MCPConfig describes a single MCP server connection. The transport is selected
|
||||||
// selected: stdio when Command is set, SSE when URL is set.
|
// 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 {
|
type MCPConfig struct {
|
||||||
// Name is the human-readable label of the server, used in logs and as the
|
// Name is the human-readable label of the server, used in logs and as the
|
||||||
// prefix for registering its tools uniquely.
|
// prefix for registering its tools uniquely.
|
||||||
@@ -66,8 +143,12 @@ type MCPConfig struct {
|
|||||||
Command string `env:"MCP_COMMAND"`
|
Command string `env:"MCP_COMMAND"`
|
||||||
// Args are the arguments passed to Command.
|
// Args are the arguments passed to Command.
|
||||||
Args []string `env:"MCP_ARGS"`
|
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"`
|
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 reports whether this server should be activated at startup.
|
||||||
Enabled bool `env:"MCP_ENABLED"`
|
Enabled bool `env:"MCP_ENABLED"`
|
||||||
}
|
}
|
||||||
@@ -122,6 +203,53 @@ type Client struct {
|
|||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
state state
|
state state
|
||||||
info ServerInfo
|
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.
|
// 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 is the JSON-RPC message exchange with an MCP server over a concrete
|
||||||
// transport (stdio subprocess or SSE). Implementations must be safe for
|
// transport (stdio subprocess, SSE or Streamable HTTP). Implementations must be
|
||||||
// concurrent use. This is exported so tests and hosts can inject a fake
|
// safe for concurrent use. This is exported so tests and hosts can inject a fake
|
||||||
// transport.
|
// transport.
|
||||||
type Transport interface {
|
type Transport interface {
|
||||||
// Start makes the transport ready to exchange messages.
|
// Start makes the transport ready to exchange messages.
|
||||||
@@ -154,6 +282,37 @@ type notifier interface {
|
|||||||
notify(ctx context.Context, method string, params any) error
|
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.
|
// Option customizes a Client built by New.
|
||||||
type Option func(*Client)
|
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
|
// 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
|
// stdio or SSE transport from cfg; pass WithLogger and/or WithTransport to
|
||||||
// override.
|
// override.
|
||||||
|
|||||||
@@ -24,10 +24,15 @@ type Registry struct {
|
|||||||
// NewRegistry builds a Registry that registers cfg's tools into reg. Options
|
// NewRegistry builds a Registry that registers cfg's tools into reg. Options
|
||||||
// apply to the underlying Client (logger, injected transport). It is named
|
// apply to the underlying Client (logger, injected transport). It is named
|
||||||
// NewRegistry rather than New because this package already exports New for
|
// 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 {
|
func NewRegistry(reg *registry.Registry, cfg MCPConfig, opts ...Option) *Registry {
|
||||||
|
r := &Registry{reg: reg}
|
||||||
|
opts = append(opts, WithToolListChangedHandler(r.refresh))
|
||||||
c := New(cfg, opts...)
|
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
|
// 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
|
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
|
// Unregister removes every tool this registry previously registered and closes
|
||||||
// the underlying client. It is safe to call multiple times.
|
// the underlying client. It is safe to call multiple times.
|
||||||
func (r *Registry) Unregister(ctx context.Context) error {
|
func (r *Registry) Unregister(ctx context.Context) error {
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
"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.
|
// testTool is a minimal sdk.Tool used to pre-register a colliding name.
|
||||||
type testTool struct {
|
type testTool struct {
|
||||||
name string
|
name string
|
||||||
|
|||||||
@@ -15,10 +15,13 @@ type rpcRequest struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// rpcResponse is a JSON-RPC 2.0 response. Result is kept raw until the caller
|
// 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 {
|
type rpcResponse struct {
|
||||||
JSONRPC string `json:"jsonrpc"`
|
JSONRPC string `json:"jsonrpc"`
|
||||||
ID any `json:"id,omitempty"`
|
ID any `json:"id,omitempty"`
|
||||||
|
Method string `json:"method,omitempty"`
|
||||||
Result json.RawMessage `json:"result,omitempty"`
|
Result json.RawMessage `json:"result,omitempty"`
|
||||||
Error *rpcError `json:"error,omitempty"`
|
Error *rpcError `json:"error,omitempty"`
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -19,18 +19,27 @@ import (
|
|||||||
// requests to the message endpoint (discovered via an endpoint event, falling
|
// requests to the message endpoint (discovered via an endpoint event, falling
|
||||||
// back to the configured URL).
|
// back to the configured URL).
|
||||||
type sseTransport struct {
|
type sseTransport struct {
|
||||||
cfg MCPConfig
|
cfg MCPConfig
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
url string
|
url string
|
||||||
endpoint string
|
endpoint string
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
nextID int
|
nextID int
|
||||||
pending map[string]chan []byte
|
pending map[string]chan []byte
|
||||||
resp *http.Response
|
resp *http.Response
|
||||||
body io.ReadCloser
|
body io.ReadCloser
|
||||||
closeCh chan struct{}
|
closeCh chan struct{}
|
||||||
closeOnce sync.Once
|
closeOnce sync.Once
|
||||||
started bool
|
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
|
// 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 {
|
if err := json.Unmarshal([]byte(data), &resp); err != nil {
|
||||||
return
|
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)
|
id, ok := resp.ID.(string)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -17,13 +17,24 @@ import (
|
|||||||
// stdioTransport runs an MCP server as a child process and speaks
|
// stdioTransport runs an MCP server as a child process and speaks
|
||||||
// newline-delimited JSON-RPC 2.0 to its stdin/stdout.
|
// newline-delimited JSON-RPC 2.0 to its stdin/stdout.
|
||||||
type stdioTransport struct {
|
type stdioTransport struct {
|
||||||
cfg MCPConfig
|
cfg MCPConfig
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
line *lineRPC
|
line *lineRPC
|
||||||
cmd *exec.Cmd
|
cmd *exec.Cmd
|
||||||
stdin io.WriteCloser
|
stdin io.WriteCloser
|
||||||
closeOnce sync.Once
|
closeOnce sync.Once
|
||||||
closed bool
|
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.
|
// newStdioTransport builds a stdio transport for cfg.
|
||||||
@@ -57,7 +68,7 @@ func (s *stdioTransport) Start(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
s.cmd = cmd
|
s.cmd = cmd
|
||||||
s.stdin = stdin
|
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)
|
go s.drainStderr(stderr)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -106,11 +117,12 @@ func (s *stdioTransport) drainStderr(r io.Reader) {
|
|||||||
// serialized so at most one request is in flight at a time, which keeps id
|
// serialized so at most one request is in flight at a time, which keeps id
|
||||||
// matching simple and deterministic.
|
// matching simple and deterministic.
|
||||||
type lineRPC struct {
|
type lineRPC struct {
|
||||||
in io.Writer
|
in io.Writer
|
||||||
out *bufio.Reader
|
out *bufio.Reader
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
nextID int
|
nextID int
|
||||||
|
onNotification func(method string)
|
||||||
}
|
}
|
||||||
|
|
||||||
// call writes a request and reads until the matching response arrives.
|
// 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
|
continue
|
||||||
}
|
}
|
||||||
if resp.ID == nil {
|
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
|
continue
|
||||||
}
|
}
|
||||||
if !decodeIDString(resp, id) {
|
if !decodeIDString(resp, id) {
|
||||||
|
|||||||
Reference in new issue
Block a user