diff --git a/go.mod b/go.mod index fd9d2e7..5de73ac 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,10 @@ module git.yeij.top/AskaEth/Cyrene_v2 go 1.26.0 require ( + github.com/google/jsonschema-go v0.4.3 github.com/gorilla/websocket v1.5.3 github.com/jackc/pgx/v5 v5.10.0 + github.com/modelcontextprotocol/go-sdk v1.8.0 github.com/robfig/cron/v3 v3.0.1 github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 gopkg.in/yaml.v3 v3.0.1 @@ -22,9 +24,14 @@ require ( github.com/ncruces/go-strftime v1.0.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/rogpeppe/go-internal v1.16.0 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/oauth2 v0.35.0 // indirect golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.29.0 // indirect + golang.org/x/time v0.15.0 // indirect modernc.org/libc v1.75.6 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.12.1 // indirect diff --git a/go.sum b/go.sum index 2049d04..93eac46 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,12 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo= github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -26,6 +32,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= +github.com/modelcontextprotocol/go-sdk v1.8.0 h1:KIvahhYqwtbeniWVPs3TcXEA7b8jEtwfBpOTAI+Urx4= +github.com/modelcontextprotocol/go-sdk v1.8.0/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= @@ -38,19 +46,29 @@ github.com/rogpeppe/go-internal v1.16.0 h1:O9DK+vNMDVGLr2BeZqmpLeMjiMNkuXfcqntWb github.com/rogpeppe/go-internal v1.16.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 h1:lZUw3E0/J3roVtGQ+SCrUrg3ON6NgVqpn3+iol9aGu4= github.com/santhosh-tekuri/jsonschema/v5 v5.3.1/go.mod h1:uToXkOrWAZ6/Oc07xWQrPOhJotwFIyu2bBVN41fcDUY= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk= golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40= +golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ= +golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE= golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/capability/mcp/capabilities.go b/internal/capability/mcp/capabilities.go new file mode 100644 index 0000000..1307e6c --- /dev/null +++ b/internal/capability/mcp/capabilities.go @@ -0,0 +1,135 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// hostFeatures 表示宿主侧配置了哪些客户端能力（决定向服务器宣告什么）。 +type hostFeatures struct { + sampling bool + elicitation bool + roots bool +} + +// hostFeatures 返回宿主侧已配置的客户端能力。 +func (c *Client) hostFeatures() hostFeatures { + return hostFeatures{ + sampling: c.sampling != nil, + elicitation: c.elicitation != nil, + roots: len(c.roots) > 0, + } +} + +// clientCapabilities 构造客户端能力声明：只宣告宿主真的具备的能力。 +// 显式传非 nil 的 Capabilities 会关掉 SDK 默认宣告的 roots——本包默认不宣告任何 +// 能力，避免对未声明能力的服务器暴露工作区根目录或代跑 LLM。 +func (c *Client) clientCapabilities() *mcpsdk.ClientCapabilities { + caps := &mcpsdk.ClientCapabilities{} + if c.sampling != nil { + caps.Sampling = &mcpsdk.SamplingCapabilities{} + } + if c.elicitation != nil { + caps.Elicitation = &mcpsdk.ElicitationCapabilities{} + } + if len(c.roots) > 0 { + caps.RootsV2 = &mcpsdk.RootCapabilities{} + } + return caps +} + +// serverInfoFrom 把 SDK 会话的 initialize 结果整理成本包的 ServerInfo， +// 其中 capabilities 原文保留、对客户端可读（design/09 §1.1）。 +func serverInfoFrom(session *mcpsdk.ClientSession, host hostFeatures) ServerInfo { + ir := session.InitializeResult() + if ir == nil { + return ServerInfo{Capabilities: map[string]any{}} + } + info := ServerInfo{Instructions: ir.Instructions, ProtocolVersion: ir.ProtocolVersion} + if ir.ServerInfo != nil { + info.Name = ir.ServerInfo.Name + info.Version = ir.ServerInfo.Version + } + info.Capabilities = capabilitiesMap(ir.Capabilities) + info.Features = featuresFrom(info.Capabilities, host) + return info +} + +// capabilitiesMap 把 SDK 解析出的服务器能力序列化回原文 map 保留下来。 +// 它只做无损转换，不做增删；极端情况下序列化失败则返回空 map（等价服务器未声明）。 +func capabilitiesMap(caps *mcpsdk.ServerCapabilities) map[string]any { + if caps == nil { + return map[string]any{} + } + data, err := json.Marshal(caps) + if err != nil { + return map[string]any{} + } + out := map[string]any{} + if err := json.Unmarshal(data, &out); err != nil { + return map[string]any{} + } + return out +} + +// featuresFrom 依据服务器 initialize 声明推断可用能力族（design/09 §1.1： +// 声明了才算可用，未声明不得调用）： +// +// - tools / prompts / resources / logging / completions 直接按服务器是否声明该字段； +// - 老式服务器完全不发 capabilities（nil 或空）时按 legacy 处理：只认 tools/list +// 是唯一发现入口（Tools=true），其余一律不可用； +// - sampling / elicitation / roots 属于客户端侧能力，服务器不会声明「自身支持」它们。 +// 因此判定为「宿主配置了回调」且「服务器声明过能力族（即确实是会主动交互的对端）」 +// 两者同时成立才可用——什么都没声明的服务器不得回调这些能力。 +func featuresFrom(caps map[string]any, host hostFeatures) ServerFeatures { + declared := len(caps) > 0 + f := ServerFeatures{ + Tools: !declared || hasCap(caps, "tools"), + Prompts: hasCap(caps, "prompts"), + Resources: hasCap(caps, "resources"), + Logging: hasCap(caps, "logging"), + Completions: hasCap(caps, "completions"), + } + f.Sampling = host.sampling && declared + f.Elicitation = host.elicitation && declared + f.Roots = host.roots && declared + return f +} + +// hasCap 报告原始 capabilities 里是否声明了名为 name 的能力。 +func hasCap(caps map[string]any, name string) bool { + v, ok := caps[name] + return ok && v != nil +} + +// wrapSampling 包装宿主 sampling 回调：服务器未声明能力族时拒绝， +// 不静默放行（design/09 §1.1）。 +func (c *Client) wrapSampling() func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) { + h := c.sampling + if h == nil { + return nil + } + return func(ctx context.Context, req *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) { + if !c.Features().Sampling { + return nil, fmt.Errorf("mcp sampling %q: %w", c.cfg.Name, ErrCapabilityNotDeclared) + } + return h(ctx, req) + } +} + +// wrapElicitation 包装宿主 elicitation 回调，判定同 wrapSampling。 +func (c *Client) wrapElicitation() func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) { + h := c.elicitation + if h == nil { + return nil + } + return func(ctx context.Context, req *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) { + if !c.Features().Elicitation { + return nil, fmt.Errorf("mcp elicitation %q: %w", c.cfg.Name, ErrCapabilityNotDeclared) + } + return h(ctx, req) + } +} diff --git a/internal/capability/mcp/capabilities_test.go b/internal/capability/mcp/capabilities_test.go new file mode 100644 index 0000000..86d366c --- /dev/null +++ b/internal/capability/mcp/capabilities_test.go @@ -0,0 +1,119 @@ +package mcp + +import ( + "context" + "errors" + "testing" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// TestFeaturesFrom_DeclarationMatrix 覆盖「声明了才算可用」的判定矩阵。 +func TestFeaturesFrom_DeclarationMatrix(t *testing.T) { + t.Parallel() + host := hostFeatures{sampling: true, elicitation: true, roots: true} + tests := []struct { + name string + caps map[string]any + wantTools bool + wantPrompts bool + wantSampling bool + }{ + {name: "未声明任何能力（legacy 服务器）", caps: map[string]any{}, wantTools: true, wantSampling: false}, + {name: "只有 tools", caps: map[string]any{"tools": map[string]any{}}, wantTools: true, wantSampling: true}, + {name: "只有 logging", caps: map[string]any{"logging": map[string]any{}}, wantTools: false, wantSampling: true}, + {name: "tools+prompts", caps: map[string]any{"tools": map[string]any{}, "prompts": map[string]any{}}, wantTools: true, wantPrompts: true, wantSampling: true}, + {name: "tools 为 null 视为未声明", caps: map[string]any{"tools": nil}, wantTools: false, wantSampling: true}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := featuresFrom(tt.caps, host) + if got.Tools != tt.wantTools { + t.Fatalf("Tools = %v, 期望 %v（caps=%v）", got.Tools, tt.wantTools, tt.caps) + } + if got.Prompts != tt.wantPrompts { + t.Fatalf("Prompts = %v, 期望 %v", got.Prompts, tt.wantPrompts) + } + if got.Sampling != tt.wantSampling { + t.Fatalf("Sampling = %v, 期望 %v", got.Sampling, tt.wantSampling) + } + }) + } +} + +// TestFeaturesFrom_HostNotConfigured 宿主没配回调时，即使服务器声明了能力也不可用。 +func TestFeaturesFrom_HostNotConfigured(t *testing.T) { + t.Parallel() + got := featuresFrom(map[string]any{"tools": map[string]any{}}, hostFeatures{}) + if got.Sampling || got.Elicitation || got.Roots { + t.Fatalf("宿主未配置任何客户端能力，却推断为可用：%+v", got) + } +} + +// TestClient_SamplingHandler_RefusedWhenNotDeclared 服务器未声明能力族时， +// 包装后的 sampling 回调必须直接拒绝，不放行给宿主。 +func TestClient_SamplingHandler_RefusedWhenNotDeclared(t *testing.T) { + t.Parallel() + called := false + c := New(MCPConfig{Name: "s"}, WithSamplingHandler(func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) { + called = true + return &mcpsdk.CreateMessageResult{}, nil + })) + h := c.wrapSampling() + if h == nil { + t.Fatal("配了 sampling 回调，wrapSampling 却返回 nil") + } + // 未连接（等价服务器什么都没声明）：必须拒绝。 + if _, err := h(context.Background(), &mcpsdk.CreateMessageRequest{}); !errors.Is(err, ErrCapabilityNotDeclared) { + t.Fatalf("未声明时 error = %v, 期望 ErrCapabilityNotDeclared", err) + } + if called { + t.Fatal("宿主回调在服务器未声明能力时被调用了") + } + + // 服务器声明过能力族 → 放行。 + c.mu.Lock() + c.info.Features = featuresFrom(map[string]any{"tools": map[string]any{}}, c.hostFeatures()) + c.mu.Unlock() + if _, err := h(context.Background(), &mcpsdk.CreateMessageRequest{}); err != nil { + t.Fatalf("已声明时 error = %v, 期望放行", err) + } + if !called { + t.Fatal("已声明时宿主回调没有被调用") + } +} + +// TestClient_SamplingHandler_RefusedWhenHostNotConfigured 没配回调就不包装（返回 nil）。 +func TestClient_SamplingHandler_RefusedWhenHostNotConfigured(t *testing.T) { + t.Parallel() + c := New(MCPConfig{Name: "s"}) + if c.wrapSampling() != nil { + t.Fatal("未配置 sampling 回调时不应注册 SDK handler") + } + if c.wrapElicitation() != nil { + t.Fatal("未配置 elicitation 回调时不应注册 SDK handler") + } +} + +// TestClient_ClientCapabilities_OnlyConfigured 客户端只宣告宿主真正具备的能力。 +func TestClient_ClientCapabilities_OnlyConfigured(t *testing.T) { + t.Parallel() + bare := New(MCPConfig{Name: "s"}).clientCapabilities() + if bare.Sampling != nil || bare.Elicitation != nil || bare.RootsV2 != nil { + t.Fatalf("默认客户端能力应为空：%+v", bare) + } + full := New(MCPConfig{Name: "s"}, + WithSamplingHandler(func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) { + return nil, nil + }), + WithElicitationHandler(func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) { + return nil, nil + }), + WithRoots(&mcpsdk.Root{URI: "file:///workspace"}), + ).clientCapabilities() + if full.Sampling == nil || full.Elicitation == nil || full.RootsV2 == nil { + t.Fatalf("配置齐全时客户端能力不全：%+v", full) + } +} diff --git a/internal/capability/mcp/client.go b/internal/capability/mcp/client.go index 9840121..6849162 100644 --- a/internal/capability/mcp/client.go +++ b/internal/capability/mcp/client.go @@ -2,210 +2,155 @@ package mcp import ( "context" - "encoding/json" "errors" "fmt" - "log/slog" - "time" + + sdkjsonrpc "github.com/modelcontextprotocol/go-sdk/jsonrpc" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" "git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk" ) -// initializeParams is the MCP initialize request body. -type initializeParams struct { - ProtocolVersion string `json:"protocolVersion"` - Capabilities map[string]any `json:"capabilities"` - ClientInfo clientInfo `json:"clientInfo"` -} - -// initializeResult is the MCP initialize response body. -type initializeResult struct { - ProtocolVersion string `json:"protocolVersion"` - Capabilities map[string]any `json:"capabilities"` - ServerInfo serverInfo `json:"serverInfo"` -} - -// clientInfo identifies the client to the server. -type clientInfo struct { - Name string `json:"name"` - Version string `json:"version"` -} - -// serverInfo identifies the server to the client. -type serverInfo struct { - Name string `json:"name"` - Version string `json:"version"` -} - -// Connect selects the transport (unless one was injected), starts it, completes -// the MCP initialize handshake, and marks the client connected. It is -// idempotent; a second call while connected returns nil. +// Connect 选择传输（未注入时）、建立官方 SDK 会话并完成 initialize 握手。 +// 幂等：已连接时直接返回 nil；已关闭时返回 ErrClosed。 func (c *Client) Connect(ctx context.Context) error { - c.mu.Lock() - switch c.state { - case stateClosed: - c.mu.Unlock() - return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, ErrClosed) - case stateConnected: - c.mu.Unlock() + same, err := c.begin() + if err != nil { + return err + } + if same { return nil } - c.mu.Unlock() tr := c.tr if tr == nil { - var err error - if tr, err = newTransport(c.cfg, c.log); err != nil { + if tr, err = newTransport(c.cfg); err != nil { return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err) } } - if err := tr.Start(ctx); err != nil { - 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) + session, err := c.newSDKClient().Connect(ctx, tr, c.sessionOptions()) if err != nil { - tr.Close() - return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err) + return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, wrapConnectErr(err)) } - // Let the transport echo the negotiated protocol version (Streamable HTTP). - attachProtocolProvider(tr, c.protocolVersion) - c.sendInitialized(ctx, tr) + info := serverInfoFrom(session, c.hostFeatures()) c.mu.Lock() + c.session = session c.tr = tr - c.info = info c.state = stateConnected + c.info = info c.mu.Unlock() - c.startKeepalive() - c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name, "protocol", c.protocolVersion()) + c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name, + "version", info.Version, "protocol", info.ProtocolVersion, + "tools_declared", info.Features.Tools) return nil } -// Initialize performs the MCP initialize handshake on an already-connected -// client. It returns the server identity. -func (c *Client) Initialize(ctx context.Context) (ServerInfo, error) { - tr := c.transport() - if tr == nil { - return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w", c.cfg.Name, ErrNotConnected) +// Close 关闭 SDK 会话并标记客户端已关闭。幂等。 +func (c *Client) Close() error { + c.mu.Lock() + if c.state == stateClosed { + c.mu.Unlock() + return nil } - return c.initialize(ctx, tr) + session := c.session + c.session = nil + c.state = stateClosed + c.mu.Unlock() + + if session != nil { + if err := session.Close(); err != nil { + return fmt.Errorf("mcp close %q: %w", c.cfg.Name, err) + } + } + return nil } -// 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 +// ListTools 返回服务器暴露的工具（SDK 自动翻页）。 +// 服务器在 initialize 里没有声明 tools 能力族时拒绝调用、不发出 RPC： +// design/09 §1.1「声明了才算可用，未声明不得调用」。 +func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) { + session := c.connected() + if session == nil { + return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrNotConnected) } - var lastErr error - for attempt := 0; attempt < maxInitializeAttempts; attempt++ { - params := initializeParams{ - ProtocolVersion: sent, - Capabilities: c.capabilities(sent), - ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"}, - } - var res initializeResult - err := tr.Call(ctx, "initialize", params, &res) + if !c.Features().Tools { + return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrCapabilityNotDeclared) + } + out := make([]ToolInfo, 0, 8) + for tool, err := range session.Tools(ctx, nil) { 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 + return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, 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 + out = append(out, toolInfoFrom(tool)) } - return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w: %w", c.cfg.Name, ErrUnsupportedProtocolVersion, lastErr) + return out, nil } -// 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) +// CallTool 调用服务器上的工具。args 必须是 JSON 对象（map[string]any）。 +// 调用前用 schema 缓存里的 input schema 校验参数并补默认值。 +func (c *Client) CallTool(ctx context.Context, name string, args any) (ToolResult, error) { + arguments, err := normalizeArgs(args) if err != nil { - return nil, false + return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err) } - 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 + return c.callTool(ctx, name, arguments, nil) } -// sendInitialized best-effort sends the notifications/initialized advisory ack. -func (c *Client) sendInitialized(ctx context.Context, tr Transport) { - n, ok := tr.(notifier) - if !ok { - return +// callTool 是 CallTool 的内部实现，fallback 是缓存未命中时可用的 input schema +// （来自 tools/list 的发现结果）。 +func (c *Client) callTool(ctx context.Context, name string, args map[string]any, fallback map[string]any) (ToolResult, error) { + session := c.connected() + if session == nil { + return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, ErrNotConnected) } - if err := n.notify(ctx, "notifications/initialized", map[string]any{}); err != nil { - c.log.Debug("mcp initialized notification failed", "name", c.cfg.Name, "err", err) + prepared, err := c.schemas.Prepare(c.cfg.Name, name, fallback, args) + if err != nil { + return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err) + } + res, err := session.CallTool(ctx, &mcpsdk.CallToolParams{Name: name, Arguments: prepared}) + if err != nil { + return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err) + } + return ToolResult{ + Content: toContentBlocks(res.Content), + Structured: res.StructuredContent, + IsError: res.IsError, + }, nil +} + +// DiscoverTools 列出服务器工具，把每把工具包装成 sdk.Tool 并把 input schema +// 放进缓存，使 MCP server 看起来就是一个普通工具提供者。 +func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) { + tools, err := c.ListTools(ctx) + if err != nil { + return nil, err + } + out := make([]sdk.Tool, 0, len(tools)) + for _, t := range tools { + c.cacheSchema(t) + out = append(out, &mcpTool{client: c, info: t}) + } + return out, nil +} + +// cacheSchema 把一次发现到的 input schema 解析并放进缓存；解析失败只降级为 +// 「调用时不校验」，不影响工具接入。 +func (c *Client) cacheSchema(t ToolInfo) { + if err := c.schemas.Put(c.cfg.Name, t.Name, t.InputSchema); err != nil { + c.log.Warn("mcp tool input schema not resolvable; validation skipped", + "name", c.cfg.Name, "tool", t.Name, "err", err) } } -// 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 { +// handleToolListChanged 是 SDK 的 notifications/tools/list_changed 回调： +// 先让 schema 缓存失效（design/09 §1.1），再合并触发宿主的增量刷新。 +// 它可能运行在 SDK 的读路径上，因此真正的刷新放到新 goroutine 里。 +func (c *Client) handleToolListChanged(context.Context, *mcpsdk.ToolListChangedRequest) { + c.schemas.Invalidate(c.cfg.Name) + if c.onToolListChanged == nil { return } if !c.refreshPending.CompareAndSwap(false, true) { @@ -217,163 +162,41 @@ func (c *Client) handleServerNotification(method string) { }() } -// 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) +// newSDKClient 按配置构建官方 SDK 客户端：只宣告宿主真正具备的客户端能力。 +func (c *Client) newSDKClient() *mcpsdk.Client { + client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: clientName, Version: clientVersion}, &mcpsdk.ClientOptions{ + Logger: c.log, + Capabilities: c.clientCapabilities(), + CreateMessageHandler: c.wrapSampling(), + ElicitationHandler: c.wrapElicitation(), + ToolListChangedHandler: c.handleToolListChanged, + KeepAlive: c.keepAlive, + KeepAliveFailureThreshold: c.keepAliveFailures, + }) + if len(c.roots) > 0 { + client.AddRoots(c.roots...) } - if err := tr.Call(ctx, "ping", map[string]any{}, nil); err != nil { - return fmt.Errorf("mcp ping: %w", err) - } - return nil + return client } -// 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() - if tr == nil { - return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrNotConnected) - } - var res listToolsResult - if err := tr.Call(ctx, "tools/list", map[string]any{}, &res); err != nil { - return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, err) - } - out := make([]ToolInfo, 0, len(res.Tools)) - for _, t := range res.Tools { - out = append(out, ToolInfo{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema}) - } - return out, nil -} - -// CallTool invokes a tool on the connected MCP server. -func (c *Client) CallTool(ctx context.Context, name string, args any) (ToolResult, error) { - tr := c.transport() - if tr == nil { - return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, ErrNotConnected) - } - arguments, err := normalizeArgs(args) - if err != nil { - return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err) - } - params := callToolParams{Name: name, Arguments: arguments} - var res callToolResult - if err := tr.Call(ctx, "tools/call", params, &res); err != nil { - return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err) - } - return ToolResult{Content: toContentBlocks(res.Content), IsError: res.IsError}, nil -} - -// DiscoverTools lists the server tools and maps each to an sdk.Tool so the MCP -// server appears as a normal tool provider. -func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) { - tools, err := c.ListTools(ctx) - if err != nil { - return nil, err - } - out := make([]sdk.Tool, 0, len(tools)) - for _, t := range tools { - out = append(out, &mcpTool{client: c, info: t}) - } - return out, nil -} - -// 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 { - c.mu.Unlock() +// sessionOptions 返回 SDK 会话选项；宿主固定了协议版本时才返回非 nil。 +func (c *Client) sessionOptions() *mcpsdk.ClientSessionOptions { + if c.preferredProtocolVersion == "" { return nil } - c.state = stateClosed - tr := c.tr - if c.keepaliveCancel != nil { - c.keepaliveCancel() - c.keepaliveCancel = nil - } - c.mu.Unlock() - if tr != nil { - return tr.Close() - } - return nil + return &mcpsdk.ClientSessionOptions{ProtocolVersion: c.preferredProtocolVersion} } -// 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 - case cfg.URL != "": - return newSSETransport(cfg, log), nil - default: - return nil, fmt.Errorf("mcp: %w", ErrNoTransport) +// wrapConnectErr 把 SDK 报的「协议版本不支持」映射到本包哨兵错误，便于调用方判定。 +func wrapConnectErr(err error) error { + var wire *sdkjsonrpc.Error + if errors.As(err, &wire) && wire.Code == mcpsdk.CodeUnsupportedProtocolVersion { + return fmt.Errorf("%w: %w", ErrUnsupportedProtocolVersion, err) } + return err } -// normalizeArgs converts an agent-supplied argument value into the MCP arguments -// object, requiring a JSON object. +// normalizeArgs 把 agent 传来的参数值规整成 MCP arguments 对象（必须是 JSON 对象）。 func normalizeArgs(args any) (map[string]any, error) { if args == nil { return map[string]any{}, nil @@ -383,33 +206,3 @@ func normalizeArgs(args any) (map[string]any, error) { } return nil, fmt.Errorf("arguments must be a JSON object, got %T", args) } - -// listToolsResult is the MCP tools/list response body. -type listToolsResult struct { - Tools []toolSpec `json:"tools"` -} - -// toolSpec is a single tool in a tools/list response. -type toolSpec struct { - Name string `json:"name"` - Description string `json:"description"` - InputSchema map[string]any `json:"inputSchema"` -} - -// callToolParams is the MCP tools/call request body. -type callToolParams struct { - Name string `json:"name"` - Arguments any `json:"arguments,omitempty"` -} - -// callToolResult is the MCP tools/call response body. -type callToolResult struct { - Content []contentBlock `json:"content"` - IsError bool `json:"isError"` -} - -// contentBlock is a wire content block in a tools/call response. -type contentBlock struct { - Type string `json:"type"` - Text string `json:"text"` -} diff --git a/internal/capability/mcp/client_test.go b/internal/capability/mcp/client_test.go index 75946ae..7a52c64 100644 --- a/internal/capability/mcp/client_test.go +++ b/internal/capability/mcp/client_test.go @@ -4,244 +4,99 @@ import ( "context" "errors" "testing" + + sdkjsonrpc "github.com/modelcontextprotocol/go-sdk/jsonrpc" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" ) -func TestClient_New(t *testing.T) { +// TestClient_Connect_NegotiatesOlderProtocolVersion 协议版本不硬编码： +// 服务器只支持旧版本时，SDK 回退协商成功，协商结果保留在 ServerInfo 里。 +func TestClient_Connect_NegotiatesOlderProtocolVersion(t *testing.T) { t.Parallel() - c := New(MCPConfig{Name: "calc", Command: "mcp-calc"}) - if c.cfg.Name != "calc" { - t.Fatalf("cfg.Name = %q, want calc", c.cfg.Name) - } - if c.state != stateNew { - t.Fatalf("state = %v, want stateNew", c.state) - } -} + srv := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "old-mcp", Version: "0.9.0"}, &mcpsdk.ServerOptions{ + SupportedProtocolVersions: []string{"2024-11-05"}, + }) + addEchoAndAdd(srv) -func TestClient_Connect_FakeTransport(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 ft.starts != 1 { - t.Fatalf("starts = %d, want 1", ft.starts) - } - if !ft.hasMethod("initialize") { - t.Fatalf("initialize was not called; methods = %v", ft.calledMethods()) - } - if c.state != stateConnected { - t.Fatalf("state = %v, want stateConnected", c.state) - } - if c.info.Name != "fake-server" { - t.Fatalf("info.Name = %q, want fake-server", c.info.Name) - } -} - -func TestClient_Connect_Idempotent(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("first Connect: %v", err) - } - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("second Connect: %v", err) - } - if ft.starts != 1 { - t.Fatalf("starts = %d, want 1 (second connect must not restart)", ft.starts) - } -} - -func TestClient_Connect_NoTransport(t *testing.T) { - t.Parallel() - c := New(MCPConfig{Name: "empty"}) - err := c.Connect(context.Background()) - if !errors.Is(err, ErrNoTransport) { - t.Fatalf("Connect error = %v, want ErrNoTransport", err) - } -} - -func TestClient_Connect_Closed(t *testing.T) { - t.Parallel() - ft := &fakeTransport{} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Close(); err != nil { - t.Fatalf("Close: %v", err) - } - err := c.Connect(context.Background()) - if !errors.Is(err, ErrClosed) { - t.Fatalf("Connect error = %v, want ErrClosed", err) - } -} - -func TestClient_Connect_StartError(t *testing.T) { - t.Parallel() - ft := &fakeTransport{startErr: errors.New("boom")} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err == nil { - t.Fatal("Connect succeeded, want error") - } - if ft.closes != 1 { - t.Fatalf("closes = %d, want 1 (failed Start must Close the transport)", ft.closes) - } -} - -func TestClient_Connect_InitError(t *testing.T) { - t.Parallel() - ft := &fakeTransport{initErr: errors.New("init boom")} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err == nil { - t.Fatal("Connect succeeded, want error") - } - if ft.closes != 1 { - t.Fatalf("closes = %d, want 1 (failed initialize must Close the transport)", ft.closes) - } -} - -func TestClient_ListTools(t *testing.T) { - t.Parallel() - ft := &fakeTransport{ - tools: []ToolInfo{{Name: "add", Description: "Adds", InputSchema: map[string]any{"type": "object"}}}, - } - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) + c := connectFakeClient(t, srv) + if got := c.ServerInfo().ProtocolVersion; got != "2024-11-05" { + t.Fatalf("协商后的协议版本 = %q, 期望 2024-11-05", got) } tools, err := c.ListTools(context.Background()) if err != nil { t.Fatalf("ListTools: %v", err) } - if len(tools) != 1 || tools[0].Name != "add" || tools[0].Description != "Adds" { - t.Fatalf("tools = %+v, want [add/Adds]", tools) + if len(tools) != 2 { + t.Fatalf("旧协议下发现 %d 把工具, 期望 2", len(tools)) } } -func TestClient_ListTools_NotConnected(t *testing.T) { +// TestWrapConnectErr_MapsUnsupportedVersion SDK 报的 -32022 必须映射到本包哨兵错误， +// 其它错误原样透传（不误判）。 +func TestWrapConnectErr_MapsUnsupportedVersion(t *testing.T) { t.Parallel() - c := New(MCPConfig{Name: "s", Command: "x"}) - _, err := c.ListTools(context.Background()) - if !errors.Is(err, ErrNotConnected) { - t.Fatalf("ListTools error = %v, want ErrNotConnected", err) + mapped := wrapConnectErr(&sdkjsonrpc.Error{Code: mcpsdk.CodeUnsupportedProtocolVersion, Message: "unsupported protocol version"}) + if !errors.Is(mapped, ErrUnsupportedProtocolVersion) { + t.Fatalf("wrapConnectErr = %v, 期望 ErrUnsupportedProtocolVersion", mapped) + } + if other := wrapConnectErr(errors.New("boom")); errors.Is(other, ErrUnsupportedProtocolVersion) { + t.Fatalf("无关错误被误判：%v", other) + } + if other := wrapConnectErr(&sdkjsonrpc.Error{Code: -32601, Message: "method not found"}); errors.Is(other, ErrUnsupportedProtocolVersion) { + t.Fatalf("其它 JSON-RPC 错误码被误判：%v", other) } } -func TestClient_CallTool(t *testing.T) { +// TestClient_Connect_IdempotentAndClosed Connect 幂等；Close 后不可再连接。 +func TestClient_Connect_IdempotentAndClosed(t *testing.T) { t.Parallel() - ft := &fakeTransport{} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) + c := connectFakeClient(t, newFakeServer()) if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - res, err := c.CallTool(context.Background(), "add", map[string]any{"x": 1}) - if err != nil { - t.Fatalf("CallTool: %v", err) - } - if res.IsError { - t.Fatal("IsError = true, want false") - } - if len(res.Content) != 1 || res.Content[0].Text != "hello" { - t.Fatalf("Content = %+v, want [hello]", res.Content) - } -} - -func TestClient_CallTool_BadArgs(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) - } - _, err := c.CallTool(context.Background(), "add", "not-a-map") - if err == nil { - t.Fatal("CallTool with non-object args succeeded, want error") - } -} - -func TestClient_CallTool_NotConnected(t *testing.T) { - t.Parallel() - c := New(MCPConfig{Name: "s", Command: "x"}) - _, err := c.CallTool(context.Background(), "add", map[string]any{}) - if !errors.Is(err, ErrNotConnected) { - t.Fatalf("CallTool error = %v, want ErrNotConnected", err) - } -} - -func TestClient_DiscoverTools_MapsToSDK(t *testing.T) { - t.Parallel() - ft := &fakeTransport{ - tools: []ToolInfo{{Name: "add", Description: "Adds", InputSchema: map[string]any{"type": "object"}}}, - } - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - tools, err := c.DiscoverTools(context.Background()) - if err != nil { - t.Fatalf("DiscoverTools: %v", err) - } - if len(tools) != 1 { - t.Fatalf("len = %d, want 1", len(tools)) - } - if tools[0].Name() != "add" || tools[0].Description() != "Adds" { - t.Fatalf("tool = %q/%q, want add/Adds", tools[0].Name(), tools[0].Description()) - } - params, err := tools[0].Parameters() - if err != nil { - t.Fatalf("Parameters: %v", err) - } - if params["type"] != "object" { - t.Fatalf("params[type] = %v, want object", params["type"]) - } - out, err := tools[0].Execute(context.Background(), map[string]any{"x": 1}) - if err != nil { - t.Fatalf("Execute: %v", err) - } - if out != "hello" { - t.Fatalf("Execute result = %v, want hello", out) - } -} - -func TestClient_DiscoverTools_NilSchema(t *testing.T) { - t.Parallel() - ft := &fakeTransport{tools: []ToolInfo{{Name: "nop", InputSchema: nil}}} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - tools, err := c.DiscoverTools(context.Background()) - if err != nil { - t.Fatalf("DiscoverTools: %v", err) - } - params, err := tools[0].Parameters() - if err != nil { - t.Fatalf("Parameters: %v", err) - } - if params["type"] != "object" { - t.Fatalf("params[type] = %v, want object", params["type"]) - } -} - -func TestClient_Close_ClosesTransportAndForbidsUse(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) + t.Fatalf("重复 Connect 应幂等返回 nil, 得到 %v", err) } if err := c.Close(); err != nil { t.Fatalf("Close: %v", err) } - if ft.closes != 1 { - t.Fatalf("closes = %d, want 1", ft.closes) - } if err := c.Close(); err != nil { - t.Fatalf("second Close: %v", err) + t.Fatalf("重复 Close 应幂等, 得到 %v", err) } - if ft.closes != 1 { - t.Fatalf("closes = %d, want 1 after idempotent Close", ft.closes) + if err := c.Connect(context.Background()); !errors.Is(err, ErrClosed) { + t.Fatalf("Close 后 Connect error = %v, 期望 ErrClosed", err) } +} + +// TestClient_Connect_NoTransport 没有任何传输信息时报 ErrNoTransport。 +func TestClient_Connect_NoTransport(t *testing.T) { + t.Parallel() + c := New(MCPConfig{Name: "none"}) + if err := c.Connect(context.Background()); !errors.Is(err, ErrNoTransport) { + t.Fatalf("Connect error = %v, 期望 ErrNoTransport", err) + } +} + +// TestClient_UseBeforeConnect 未连接时各入口都返回 ErrNotConnected，不 panic。 +func TestClient_UseBeforeConnect(t *testing.T) { + t.Parallel() + c := New(MCPConfig{Name: "none", Command: "unused"}) if _, err := c.ListTools(context.Background()); !errors.Is(err, ErrNotConnected) { - t.Fatalf("ListTools after Close = %v, want ErrNotConnected", err) + t.Fatalf("ListTools error = %v, 期望 ErrNotConnected", err) + } + if _, err := c.CallTool(context.Background(), "echo", map[string]any{}); !errors.Is(err, ErrNotConnected) { + t.Fatalf("CallTool error = %v, 期望 ErrNotConnected", err) + } + if c.Session() != nil { + t.Fatal("未连接时 Session() 应返回 nil") + } + if c.Name() != "none" { + t.Fatalf("Name() = %q", c.Name()) + } +} + +// TestClient_CallTool_RequiresJSONObject 参数不是 JSON 对象时直接拒绝，不发请求。 +func TestClient_CallTool_RequiresJSONObject(t *testing.T) { + t.Parallel() + c := connectFakeClient(t, newFakeServer()) + if _, err := c.CallTool(context.Background(), "echo", "not-an-object"); err == nil { + t.Fatal("字符串参数没有被拒绝") } } diff --git a/internal/capability/mcp/discovery_test.go b/internal/capability/mcp/discovery_test.go new file mode 100644 index 0000000..f68dd9c --- /dev/null +++ b/internal/capability/mcp/discovery_test.go @@ -0,0 +1,294 @@ +package mcp + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "os" + "slices" + "sync/atomic" + "testing" + "time" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry" + "git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk" +) + +// TestClient_DiscoverTools_InMemory 走真实 SDK 内存传输打通「连接 → 工具自动发现」。 +func TestClient_DiscoverTools_InMemory(t *testing.T) { + t.Parallel() + c := connectFakeClient(t, newFakeServer()) + + info := c.ServerInfo() + if info.Name != "fake-mcp" || info.Version != "1.0.0" { + t.Fatalf("ServerInfo = %+v, 期望 fake-mcp/1.0.0", info) + } + if info.ProtocolVersion == "" { + t.Fatal("ProtocolVersion 为空，协商结果没有保留") + } + if _, ok := info.Capabilities["tools"]; !ok { + t.Fatalf("服务器声明的 capabilities 没有保留 tools 字段：%v", info.Capabilities) + } + if !info.Features.Tools { + t.Fatalf("Features.Tools = false，服务器已声明 tools：%v", info.Capabilities) + } + + tools, err := c.ListTools(context.Background()) + if err != nil { + t.Fatalf("ListTools: %v", err) + } + got := toolNames(tools) + slices.Sort(got) + if !slices.Equal(got, []string{"add", "echo"}) { + t.Fatalf("tools = %v, 期望 [add echo]", got) + } + + discovered, err := c.DiscoverTools(context.Background()) + if err != nil { + t.Fatalf("DiscoverTools: %v", err) + } + if len(discovered) != 2 { + t.Fatalf("DiscoverTools 返回 %d 把工具，期望 2", len(discovered)) + } + if c.SchemaCache().Len() != 2 { + t.Fatalf("schema 缓存条目 = %d，期望 2（发现即缓存）", c.SchemaCache().Len()) + } +} + +// TestClient_DiscoverTools_StdioProcess 用真实子进程（stdio 打桩）证明 stdio 传输下 +// 工具自动发现能跑通：子进程就是本测试二进制，以 MCP server 身份讲协议。 +func TestClient_DiscoverTools_StdioProcess(t *testing.T) { + t.Parallel() + self, err := os.Executable() + if err != nil { + t.Fatalf("os.Executable: %v", err) + } + c := New(MCPConfig{ + Name: "stdio-fake", + Command: self, + Args: []string{"-test.run=^TestMCPStdioHelper$"}, + Env: map[string]string{stdioHelperEnv: "1"}, + }) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := c.Connect(ctx); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + + tools, err := c.DiscoverTools(ctx) + if err != nil { + t.Fatalf("DiscoverTools: %v", err) + } + got := toolNames(toToolInfos(tools)) + slices.Sort(got) + if !slices.Equal(got, []string{"add", "echo"}) { + t.Fatalf("stdio 发现的工具 = %v, 期望 [add echo]", got) + } + + res, err := c.CallTool(ctx, "echo", map[string]any{"name": "stdio"}) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if text := contentText(res.Content); text != "hello stdio" { + t.Fatalf("工具结果 = %q, 期望 hello stdio", text) + } +} + +// TestClient_DiscoverTools_StreamableHTTP 用 httptest 打桩 Streamable HTTP 端点， +// 验证 URL 型传输 + bearer 鉴权 + 工具自动发现整条链。 +func TestClient_DiscoverTools_StreamableHTTP(t *testing.T) { + t.Parallel() + srv := newFakeServer() + handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return srv }, nil) + var authorized atomic.Bool + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer sk-test" { + w.WriteHeader(http.StatusUnauthorized) + return + } + authorized.Store(true) + handler.ServeHTTP(w, r) + })) + defer ts.Close() + + c := New(MCPConfig{Name: "http-fake", URL: ts.URL, TransportType: "http", APIKey: "sk-test"}) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := c.Connect(ctx); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + + tools, err := c.DiscoverTools(ctx) + if err != nil { + t.Fatalf("DiscoverTools: %v", err) + } + if len(tools) != 2 { + t.Fatalf("Streamable HTTP 发现 %d 把工具，期望 2", len(tools)) + } + if !authorized.Load() { + t.Fatal("服务器没有收到带 bearer 凭据的请求") + } + res, err := c.CallTool(ctx, "add", map[string]any{"a": 1, "b": 2}) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if text := contentText(res.Content); text != "3" { + t.Fatalf("add 结果 = %q, 期望 3", text) + } +} + +// TestClient_ListTools_RefusesUndeclaredCapability 证明「未声明不得调用」： +// 服务器只声明 logging，客户端拒绝发 tools/list。 +func TestClient_ListTools_RefusesUndeclaredCapability(t *testing.T) { + t.Parallel() + c := connectFakeClient(t, newEmptyServer()) + if c.Features().Tools { + t.Fatalf("空服务器不应推断出 tools 能力：%v", c.ServerInfo().Capabilities) + } + _, err := c.ListTools(context.Background()) + if !errors.Is(err, ErrCapabilityNotDeclared) { + t.Fatalf("ListTools error = %v, 期望 ErrCapabilityNotDeclared", err) + } + if _, err := c.DiscoverTools(context.Background()); !errors.Is(err, ErrCapabilityNotDeclared) { + t.Fatalf("DiscoverTools error = %v, 期望 ErrCapabilityNotDeclared", err) + } +} + +// TestRegistry_Register_AutoRegisters 打通「自动发现 → 注册进 capability registry +// → 模型可调用」这条链（用 registry 里存的工具值真的发起一次调用）。 +func TestRegistry_Register_AutoRegisters(t *testing.T) { + t.Parallel() + reg := registry.New() + tr := startFakeServer(t, newFakeServer()) + r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(tr)) + + if err := r.Register(context.Background()); err != nil { + t.Fatalf("Register: %v", err) + } + names := r.ToolNames() + slices.Sort(names) + if !slices.Equal(names, []string{"add", "echo"}) { + t.Fatalf("ToolNames = %v, 期望 [add echo]", names) + } + + cap, err := reg.Get("echo") + if err != nil { + t.Fatalf("registry.Get(echo): %v", err) + } + if cap.Kind != registry.KindTool { + t.Fatalf("echo 的 Kind = %q, 期望 %q", cap.Kind, registry.KindTool) + } + tool, ok := cap.Value.(sdk.Tool) + if !ok { + t.Fatalf("echo 的 Value 类型 = %T, 期望 sdk.Tool", cap.Value) + } + + // 不给参数调用：默认值 world 由 input schema 补上（证明 schema 校验/默认值生效）。 + out, err := tool.Execute(context.Background(), nil) + if err != nil { + t.Fatalf("Execute(echo): %v", err) + } + if out != "hello world" { + t.Fatalf("echo 结果 = %v, 期望 hello world（默认值未生效？）", out) + } + + // 必填参数缺失或类型错误时必须被 input schema 拦下，不把坏参数发给服务器。 + addCap, err := reg.Get("add") + if err != nil { + t.Fatalf("registry.Get(add): %v", err) + } + addTool := addCap.Value.(sdk.Tool) + if _, err := addTool.Execute(context.Background(), map[string]any{"a": "not-a-number"}); err == nil { + t.Fatal("非法参数没有被 input schema 拦下") + } +} + +// TestRegistry_Register_PropagatesConnectError 连接失败要冒泡，不能静默注册空工具集。 +func TestRegistry_Register_PropagatesConnectError(t *testing.T) { + t.Parallel() + reg := registry.New() + r := NewRegistry(reg, MCPConfig{Name: "broken", Command: "definitely-not-a-real-binary-xyz"}) + if err := r.Register(context.Background()); err == nil { + t.Fatal("Register 成功，期望连接错误") + } +} + +// TestRegistry_Register_NilRegistry 空 registry 必须报错。 +func TestRegistry_Register_NilRegistry(t *testing.T) { + t.Parallel() + r := NewRegistry(nil, MCPConfig{Name: "calc"}) + if err := r.Register(context.Background()); err == nil { + t.Fatal("Register 成功，期望 nil registry 错误") + } +} + +// TestRegistry_Unregister_RemovesTools 卸载要真从 registry 移除并关会话，且幂等。 +func TestRegistry_Unregister_RemovesTools(t *testing.T) { + t.Parallel() + reg := registry.New() + tr := startFakeServer(t, newFakeServer()) + r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(tr)) + if err := r.Register(context.Background()); err != nil { + t.Fatalf("Register: %v", err) + } + if err := r.Unregister(context.Background()); err != nil { + t.Fatalf("Unregister: %v", err) + } + if err := r.Unregister(context.Background()); err != nil { + t.Fatalf("第二次 Unregister: %v", err) + } + for _, name := range []string{"add", "echo"} { + if _, err := reg.Get(name); !errors.Is(err, registry.ErrNotFound) { + t.Fatalf("Get(%q) 卸载后 = %v, 期望 ErrNotFound", name, err) + } + } +} + +// TestRegistry_ToolListChangedRefreshes 服务器运行中加工具 → 通知 → registry 增量刷新， +// 同时 schema 缓存按服务器整表失效（design/09 §1.1）。 +func TestRegistry_ToolListChangedRefreshes(t *testing.T) { + t.Parallel() + srv := newFakeServer() + reg := registry.New() + tr := startFakeServer(t, srv) + r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(tr)) + if err := r.Register(context.Background()); err != nil { + t.Fatalf("Register: %v", err) + } + + _, _ = reg.Get("echo") // 让缓存先有命中记录，便于确认失效计数 + srv.AddTool(&mcpsdk.Tool{ + Name: "mul", + Description: "两数相乘", + InputSchema: addToolSpec().InputSchema, + }, addHandler) + + waitFor(t, 5*time.Second, "registry 出现新工具 mul", func() bool { + _, err := reg.Get("mul") + return err == nil + }) + if _, err := reg.Get("echo"); err != nil { + t.Fatalf("刷新后 echo 丢失：%v", err) + } + if got := r.ToolNames(); !slices.Contains(got, "mul") { + t.Fatalf("ToolNames = %v, 期望含 mul", got) + } +} + +// toToolInfos 把 sdk.Tool 列表还原成 ToolInfo（仅测试断言工具名用）。 +func toToolInfos(tools []sdk.Tool) []ToolInfo { + out := make([]ToolInfo, 0, len(tools)) + for _, t := range tools { + info := ToolInfo{Name: t.Name(), Description: t.Description()} + if t2, ok := t.(*mcpTool); ok { + info = t2.info + } + out = append(out, info) + } + return out +} diff --git a/internal/capability/mcp/enhancements_test.go b/internal/capability/mcp/enhancements_test.go deleted file mode 100644 index 1154db4..0000000 --- a/internal/capability/mcp/enhancements_test.go +++ /dev/null @@ -1,194 +0,0 @@ -package mcp - -import ( - "context" - "errors" - "fmt" - "reflect" - "testing" - "time" -) - -func TestClient_Connect_NegotiatesDowngradeOnUnsupportedError(t *testing.T) { - t.Parallel() - ft := &fakeTransport{ - initErrors: []error{unsupportedProtocolError("2024-11-05")}, - initResponseVersion: protocolVersion20241105, - } - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - got := ft.initSentVersions() - if !reflect.DeepEqual(got, []string{"2025-11-25", "2024-11-05"}) { - t.Fatalf("init versions = %v, want [2025-11-25 2024-11-05]", got) - } - if c.protocolVersion() != protocolVersion20241105 { - t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105) - } -} - -func TestClient_Connect_AcceptsServerSupportedVersion(t *testing.T) { - t.Parallel() - ft := &fakeTransport{} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - if got := ft.initSentVersions(); !reflect.DeepEqual(got, []string{"2025-11-25"}) { - t.Fatalf("init versions = %v, want [2025-11-25]", got) - } - // The fake echoes the legacy version; the client must accept it as supported. - if c.protocolVersion() != protocolVersion20241105 { - t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105) - } -} - -func TestClient_Connect_NoMutualVersion(t *testing.T) { - t.Parallel() - // The server only supports a version the client cannot speak, so negotiation - // must exhaust and surface ErrUnsupportedProtocolVersion. - ft := &fakeTransport{initErrors: []error{unsupportedProtocolError("2000-01-01")}} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - err := c.Connect(context.Background()) - if !errors.Is(err, ErrUnsupportedProtocolVersion) { - t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err) - } -} - -func TestClient_Connect_ServerReturnsUnsupportedVersion(t *testing.T) { - t.Parallel() - // A server that echoes a version this client does not support (and cannot - // downgrade to) must fail the handshake with ErrUnsupportedProtocolVersion. - ft := &fakeTransport{initResponseVersion: "2099-01-01"} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft)) - err := c.Connect(context.Background()) - if !errors.Is(err, ErrUnsupportedProtocolVersion) { - t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err) - } -} - -func TestClient_ToolListChangedInvokesHandler(t *testing.T) { - t.Parallel() - called := make(chan struct{}, 1) - ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}} - c := New(MCPConfig{Name: "s", Command: "x"}, - WithTransport(ft), - WithToolListChangedHandler(func() { called <- struct{}{} })) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - // No handler invoked for unrelated notifications. - ft.notifyServer("notifications/prompts/list_changed") - select { - case <-called: - t.Fatal("handler invoked for a non-tool notification") - case <-time.After(50 * time.Millisecond): - } - // The real notification must invoke the handler. - ft.notifyServer(notificationToolsListChanged) - select { - case <-called: - case <-time.After(time.Second): - t.Fatal("tool list changed handler was not invoked") - } -} - -func TestClient_KeepAlive_SendsPings(t *testing.T) { - t.Parallel() - ft := &fakeTransport{} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 3)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - defer c.Close() - waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 }) -} - -func TestClient_KeepAlive_FailureClosesSession(t *testing.T) { - t.Parallel() - ft := &fakeTransport{pingErr: errors.New("dead peer")} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - // A single failed ping (threshold 1) must close the client/transport. - waitFor(t, time.Second, func() bool { return c.transport() == nil }) - waitFor(t, time.Second, func() bool { - ft.mu.Lock() - defer ft.mu.Unlock() - return ft.closes >= 1 - }) -} - -func TestClient_KeepAlive_StopsOnUnsupportedPing(t *testing.T) { - t.Parallel() - ft := &fakeTransport{pingErr: fmt.Errorf("mcp ping: %w: %w", ErrRPC, &rpcError{Code: codeMethodNotFound, Message: "method not found"})} - c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1)) - if err := c.Connect(context.Background()); err != nil { - t.Fatalf("Connect: %v", err) - } - defer c.Close() - // A method-not-found ping must stop keepalive without closing the session. - waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 }) - time.Sleep(80 * time.Millisecond) - if c.transport() == nil { - t.Fatal("client was closed after an unsupported ping, want it to stay connected") - } - if ft.pingCount() > 1 { - t.Fatalf("keepalive kept pinging after method-not-found; pings = %d", ft.pingCount()) - } -} - -func TestNewTransport_SelectsTransportType(t *testing.T) { - t.Parallel() - tests := []struct { - name string - cfg MCPConfig - want any - }{ - {name: "http", cfg: MCPConfig{URL: "http://x", TransportType: "http"}, want: &streamableHTTPTransport{}}, - {name: "streamable-alias", cfg: MCPConfig{URL: "http://x", TransportType: "streamable"}, want: &streamableHTTPTransport{}}, - {name: "sse", cfg: MCPConfig{URL: "http://x", TransportType: "sse"}, want: &sseTransport{}}, - {name: "stdio", cfg: MCPConfig{Command: "x", TransportType: "stdio"}, want: &stdioTransport{}}, - {name: "infer-stdio", cfg: MCPConfig{Command: "x"}, want: &stdioTransport{}}, - {name: "infer-sse", cfg: MCPConfig{URL: "http://x"}, want: &sseTransport{}}, - } - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - tr, err := newTransport(tt.cfg, defaultLogger()) - if err != nil { - t.Fatalf("newTransport: %v", err) - } - if reflect.TypeOf(tr) != reflect.TypeOf(tt.want) { - t.Fatalf("transport type = %T, want %T", tr, tt.want) - } - }) - } -} - -func TestNewTransport_ExplicitTypeSelectsEvenWithoutField(t *testing.T) { - t.Parallel() - // Selecting an explicit transport does not validate the field; Start does. - if tr, err := newTransport(MCPConfig{TransportType: "http"}, defaultLogger()); err != nil { - t.Fatalf("newTransport(http): %v", err) - } else if _, ok := tr.(*streamableHTTPTransport); !ok { - t.Fatalf("transport type = %T, want *streamableHTTPTransport", tr) - } -} - -// waitFor polls cond until it returns true or the deadline passes, failing the -// test on timeout. -func waitFor(t *testing.T, timeout time.Duration, cond func() bool) { - t.Helper() - deadline := time.Now().Add(timeout) - for time.Now().Before(deadline) { - if cond() { - return - } - time.Sleep(5 * time.Millisecond) - } - t.Fatal("condition not met before timeout") -} diff --git a/internal/capability/mcp/helpers_test.go b/internal/capability/mcp/helpers_test.go index c40fccc..9de8501 100644 --- a/internal/capability/mcp/helpers_test.go +++ b/internal/capability/mcp/helpers_test.go @@ -2,187 +2,152 @@ package mcp import ( "context" + "encoding/json" "fmt" - "sync" + "os" + "testing" + "time" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" ) -// 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 +// stdioHelperEnv 标记「测试二进制以 stdio MCP server 身份运行」。 +const stdioHelperEnv = "CYRENE_MCP_STDIO_HELPER" - // 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. -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 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 = 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.activeTools { - r.Tools = append(r.Tools, toolSpec{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema}) - } - } - case "tools/call": - if f.callErr != nil { - return f.callErr - } - 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 -} - -// notify records a notification method. -func (f *fakeTransport) notify(_ context.Context, method string, _ any) error { - f.mu.Lock() - defer f.mu.Unlock() - f.methods = append(f.methods, "notify:"+method) - 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) +// echoToolSpec 返回入参 name（string，默认 "world"）的回显工具定义。 +func echoToolSpec() *mcpsdk.Tool { + return &mcpsdk.Tool{ + Name: "echo", + Description: "回显问候语", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string", "default": "world"}, + }, + }, } } -// Close records a close. -func (f *fakeTransport) Close() error { - f.mu.Lock() - defer f.mu.Unlock() - f.closes++ - return f.closeErr +// addToolSpec 返回入参 a、b（均为必填 number）的加法工具定义。 +func addToolSpec() *mcpsdk.Tool { + return &mcpsdk.Tool{ + Name: "add", + Description: "两数相加", + InputSchema: map[string]any{ + "type": "object", + "required": []string{"a", "b"}, + "properties": map[string]any{ + "a": map[string]any{"type": "number"}, + "b": map[string]any{"type": "number"}, + }, + }, + } } -// calledMethods returns a snapshot of the recorded methods. -func (f *fakeTransport) calledMethods() []string { - f.mu.Lock() - defer f.mu.Unlock() - out := make([]string, len(f.methods)) - copy(out, f.methods) +// echoHandler 按 schema 解参并回显；arguments 在 SDK 里是 json.RawMessage。 +func echoHandler(_ context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { + var in struct { + Name string `json:"name"` + } + if len(req.Params.Arguments) > 0 { + if err := json.Unmarshal(req.Params.Arguments, &in); err != nil { + return nil, fmt.Errorf("decode echo args: %w", err) + } + } + return &mcpsdk.CallToolResult{ + Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "hello " + in.Name}}, + }, nil +} + +// addHandler 求和并同时返回文本与结构化结果。 +func addHandler(_ context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { + var in struct { + A float64 `json:"a"` + B float64 `json:"b"` + } + if err := json.Unmarshal(req.Params.Arguments, &in); err != nil { + return nil, fmt.Errorf("decode add args: %w", err) + } + sum := in.A + in.B + return &mcpsdk.CallToolResult{ + Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: fmt.Sprintf("%g", sum)}}, + StructuredContent: map[string]any{"sum": sum}, + }, nil +} + +// addEchoAndAdd 往 srv 挂上两把测试工具。 +func addEchoAndAdd(srv *mcpsdk.Server) { + srv.AddTool(echoToolSpec(), echoHandler) + srv.AddTool(addToolSpec(), addHandler) +} + +// newFakeServer 构建一台真实的 MCP server（官方 SDK 服务端），挂上 echo/add。 +func newFakeServer() *mcpsdk.Server { + srv := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "fake-mcp", Version: "1.0.0"}, nil) + addEchoAndAdd(srv) + return srv +} + +// newEmptyServer 构建一台不提供任何工具的服务器（只声明 logging 能力）。 +func newEmptyServer() *mcpsdk.Server { + return mcpsdk.NewServer(&mcpsdk.Implementation{Name: "empty-mcp", Version: "1.0.0"}, nil) +} + +// startFakeServer 在 SDK 内存传输上跑 srv，返回客户端侧传输。 +func startFakeServer(t *testing.T, srv *mcpsdk.Server) mcpsdk.Transport { + t.Helper() + serverT, clientT := mcpsdk.NewInMemoryTransports() + ctx, cancel := context.WithCancel(context.Background()) + go func() { _ = srv.Run(ctx, serverT) }() + t.Cleanup(cancel) + return clientT +} + +// connectFakeClient 通过内存传输连上一台 fake server。 +func connectFakeClient(t *testing.T, srv *mcpsdk.Server, opts ...Option) *Client { + t.Helper() + opts = append(opts, WithTransport(startFakeServer(t, srv))) + c := New(MCPConfig{Name: "fake"}, opts...) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + t.Cleanup(func() { _ = c.Close() }) + return c +} + +// toolNames 取出工具名列表。 +func toolNames(tools []ToolInfo) []string { + out := make([]string, 0, len(tools)) + for _, t := range tools { + out = append(out, t.Name) + } return out } -// hasMethod reports whether method was recorded. -func (f *fakeTransport) hasMethod(method string) bool { - for _, m := range f.calledMethods() { - if m == method { - return true +// waitFor 轮询 cond 直到成立或超时。 +func waitFor(t *testing.T, timeout time.Duration, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return } + time.Sleep(5 * time.Millisecond) } - return false + t.Fatalf("等待超时：%s", what) } -// 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}, - }) +// TestMCPStdioHelper 是 stdio 打桩用的辅助进程入口：以真实 MCP server 身份 +// 在 stdin/stdout 上讲协议。stdout 是协议通道，因此诊断只写 stderr，且退出时 +// 直接 os.Exit（避免测试框架往 stdout 打 PASS/ok 破坏协议帧）。 +func TestMCPStdioHelper(t *testing.T) { + if os.Getenv(stdioHelperEnv) != "1" { + t.Skip("仅作为 stdio 辅助进程入口") + } + srv := newFakeServer() + if err := srv.Run(context.Background(), &mcpsdk.StdioTransport{}); err != nil { + fmt.Fprintln(os.Stderr, "stdio helper server:", err) + os.Exit(1) + } + os.Exit(0) } diff --git a/internal/capability/mcp/http.go b/internal/capability/mcp/http.go deleted file mode 100644 index bcfd055..0000000 --- a/internal/capability/mcp/http.go +++ /dev/null @@ -1,422 +0,0 @@ -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.applyAuth(req) - 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) - } -} - -// applyAuth sets the bearer credential an authenticated server requires. It is -// a no-op when the config carries no API key. -func (s *streamableHTTPTransport) applyAuth(req *http.Request) { - if s.cfg.APIKey != "" { - req.Header.Set("Authorization", "Bearer "+s.cfg.APIKey) - } -} - -// 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.applyAuth(req) - s.mu.Lock() - session := s.session - proto := "" - if s.protocol != nil { - proto = s.protocol() - } - s.mu.Unlock() - if session != "" { - req.Header.Set(mcpSessionIDHeader, session) - } - if proto != "" { - req.Header.Set("Mcp-Protocol-Version", proto) - } - resp, err := s.http.Do(req) - if err != nil { - if ctx.Err() != nil { - return - } - s.log.Debug("mcp http get stream error", "name", s.cfg.Name, "err", err) - s.sleepBackoff(ctx) - continue - } - if resp.StatusCode != http.StatusOK { - _ = resp.Body.Close() - s.log.Debug("mcp http get stream not accepted", "name", s.cfg.Name, "status", resp.StatusCode) - return - } - s.readStandaloneEvents(ctx, resp) - if ctx.Err() != nil { - return - } - s.sleepBackoff(ctx) - } -} - -// readStandaloneEvents consumes one GET stream, routing notifications until it -// ends or the connection is closed. -func (s *streamableHTTPTransport) readStandaloneEvents(ctx context.Context, resp *http.Response) { - stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() }) - defer stopClose() - br := bufio.NewReader(resp.Body) - var event, data string - for { - line, err := br.ReadString('\n') - if err != nil { - return - } - line = strings.TrimRight(line, "\r\n") - switch { - case line == "": - if data != "" { - s.handleEvent(event, data, "", "", nil) - } - event, data = "", "" - case strings.HasPrefix(line, "event:"): - event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) - case strings.HasPrefix(line, "data:"): - d := strings.TrimSpace(strings.TrimPrefix(line, "data:")) - if data == "" { - data = d - } else { - data += "\n" + d - } - } - } -} - -// sleepBackoff waits between standalone stream reconnect attempts, exiting early -// when the connection is closed. -func (s *streamableHTTPTransport) sleepBackoff(ctx context.Context) { - select { - case <-time.After(time.Second): - case <-ctx.Done(): - } -} - -// Close closes the standalone stream and marks the transport closed. It is -// idempotent. -func (s *streamableHTTPTransport) Close() error { - s.closeOnce.Do(func() { - close(s.closeCh) - s.mu.Lock() - if s.getCancel != nil { - s.getCancel() - } - s.mu.Unlock() - }) - return nil -} diff --git a/internal/capability/mcp/http_test.go b/internal/capability/mcp/http_test.go deleted file mode 100644 index ebbb258..0000000 --- a/internal/capability/mcp/http_test.go +++ /dev/null @@ -1,197 +0,0 @@ -package mcp - -import ( - "context" - "encoding/json" - "errors" - "io" - "net/http" - "net/http/httptest" - "testing" - "time" -) - -// httpTestServer is a Streamable HTTP test server. Its GET handler keeps an SSE -// stream open until the request context is done (so it does not leak), while POST -// handling is supplied by postFn. -func httpTestServer(t *testing.T, postFn func(w http.ResponseWriter, r *http.Request)) *httptest.Server { - t.Helper() - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - if f, ok := w.(http.Flusher); ok { - f.Flush() - } - <-r.Context().Done() - return - } - postFn(w, r) - })) -} - -// httpJSONResponse writes a single JSON-RPC response echo back, using the request -// id so the client matches it. -func httpJSONResponse(w http.ResponseWriter, r *http.Request) { - body, _ := io.ReadAll(r.Body) - var req rpcRequest - _ = json.Unmarshal(body, &req) - resp := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"direct"}]}}` - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(resp)) -} - -func TestHTTP_Call_DirectJSON(t *testing.T) { - t.Parallel() - srv := httpTestServer(t, httpJSONResponse) - defer srv.Close() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - defer tr.Close() - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - var res callToolResult - if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil { - t.Fatalf("Call: %v", err) - } - if len(res.Content) != 1 || res.Content[0].Text != "direct" { - t.Fatalf("Content = %+v, want [direct]", res.Content) - } -} - -func TestHTTP_Call_EventStream(t *testing.T) { - t.Parallel() - srv := httpTestServer(t, func(w http.ResponseWriter, r *http.Request) { - body, _ := io.ReadAll(r.Body) - var req rpcRequest - _ = json.Unmarshal(body, &req) - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - f, _ := w.(http.Flusher) - msg := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"stream"}]}}` - _, _ = w.Write([]byte("event: message\ndata: " + msg + "\n\n")) - f.Flush() - }) - defer srv.Close() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - defer tr.Close() - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - var res callToolResult - if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil { - t.Fatalf("Call: %v", err) - } - if len(res.Content) != 1 || res.Content[0].Text != "stream" { - t.Fatalf("Content = %+v, want [stream]", res.Content) - } -} - -func TestHTTP_ServerNotification(t *testing.T) { - t.Parallel() - notif := make(chan string, 1) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - f, _ := w.(http.Flusher) - _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\n\n")) - f.Flush() - <-r.Context().Done() - return - } - w.WriteHeader(http.StatusAccepted) - })) - defer srv.Close() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - tr.setOnNotification(func(m string) { notif <- m }) - defer tr.Close() - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - select { - case m := <-notif: - if m != notificationToolsListChanged { - t.Fatalf("notification = %q, want %q", m, notificationToolsListChanged) - } - case <-time.After(time.Second): - t.Fatal("no server notification received") - } -} - -func TestHTTP_SendsProtocolVersionHeader(t *testing.T) { - t.Parallel() - got := make(chan string, 1) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - if f, ok := w.(http.Flusher); ok { - f.Flush() - } - <-r.Context().Done() - return - } - if v := r.Header.Get("Mcp-Protocol-Version"); v != "" { - select { - case got <- v: - default: - } - } - httpJSONResponse(w, r) - })) - defer srv.Close() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - tr.setProtocolProvider(func() string { return "2025-11-25" }) - defer tr.Close() - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - if err := tr.Call(context.Background(), "tools/list", map[string]any{}, &listToolsResult{}); err != nil { - t.Fatalf("Call: %v", err) - } - select { - case v := <-got: - if v != "2025-11-25" { - t.Fatalf("protocol header = %q, want 2025-11-25", v) - } - case <-time.After(time.Second): - t.Fatal("Mcp-Protocol-Version header not sent") - } -} - -func TestHTTP_Call_NotStarted(t *testing.T) { - t.Parallel() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger()) - err := tr.Call(context.Background(), "tools/list", nil, nil) - if !errors.Is(err, ErrNotConnected) { - t.Fatalf("Call error = %v, want ErrNotConnected", err) - } -} - -func TestHTTP_Start_NoURL(t *testing.T) { - t.Parallel() - tr := newHTTPTransport(MCPConfig{Name: "s"}, defaultLogger()) - if err := tr.Start(context.Background()); !errors.Is(err, ErrNoTransport) { - t.Fatalf("Start error = %v, want ErrNoTransport", err) - } -} - -func TestHTTP_Start_InvalidURL(t *testing.T) { - t.Parallel() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: "ftp://x"}, defaultLogger()) - if err := tr.Start(context.Background()); err == nil { - t.Fatal("Start accepted a non-http url, want error") - } -} - -func TestHTTP_Close_Idempotent(t *testing.T) { - t.Parallel() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger()) - if err := tr.Close(); err != nil { - t.Fatalf("first Close: %v", err) - } - if err := tr.Close(); err != nil { - t.Fatalf("second Close: %v", err) - } -} diff --git a/internal/capability/mcp/keepalive_test.go b/internal/capability/mcp/keepalive_test.go new file mode 100644 index 0000000..4ce7247 --- /dev/null +++ b/internal/capability/mcp/keepalive_test.go @@ -0,0 +1,109 @@ +package mcp + +import ( + "context" + "sync" + "testing" + "time" + + sdkjsonrpc "github.com/modelcontextprotocol/go-sdk/jsonrpc" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// TestClient_KeepAlive_SendsPings keepalive 接线：配了间隔后 ping 真的上了线路。 +// ping 本身由 SDK 负责，这里只验我们把配置传导到了 SDK 会话上。 +func TestClient_KeepAlive_SendsPings(t *testing.T) { + t.Parallel() + sn := &sniffTransport{inner: startFakeServer(t, newFakeServer())} + c := New(MCPConfig{Name: "fake"}, + WithTransport(sn), + // 固定到 keepalive 生效的旧协议版本，避免新版协议的订阅路径干扰。 + WithProtocolVersion("2025-11-25"), + WithKeepAlive(20*time.Millisecond, 1)) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + if got := c.ServerInfo().ProtocolVersion; got != "2025-11-25" { + t.Fatalf("协议版本 = %q, 期望 2025-11-25", got) + } + waitFor(t, 5*time.Second, "线路上出现 keepalive ping", func() bool { return sn.called("ping") }) +} + +// TestClient_NoKeepAlive_NoPings 不配 keepalive 时不发 ping（反证：ping 不是 SDK 默认行为）。 +func TestClient_NoKeepAlive_NoPings(t *testing.T) { + t.Parallel() + sn := &sniffTransport{inner: startFakeServer(t, newFakeServer())} + c := New(MCPConfig{Name: "fake"}, WithTransport(sn), WithProtocolVersion("2025-11-25")) + if err := c.Connect(context.Background()); err != nil { + t.Fatalf("Connect: %v", err) + } + defer c.Close() + if _, err := c.DiscoverTools(context.Background()); err != nil { + t.Fatalf("DiscoverTools: %v", err) + } + time.Sleep(150 * time.Millisecond) + if sn.called("ping") { + t.Fatal("没配 keepalive 却发了 ping") + } + if !sn.called("tools/list") { + t.Fatalf("线路上没有 tools/list, 已见方法 = %v", sn.methods()) + } +} + +// sniffTransport 包一层 SDK 传输，记录客户端发出的 JSON-RPC 请求方法， +// 用于在线路层断言客户端确实发了某个请求。 +type sniffTransport struct { + inner mcpsdk.Transport + mu sync.Mutex + seen []string +} + +// Connect 实现 mcpsdk.Transport。 +func (t *sniffTransport) Connect(ctx context.Context) (mcpsdk.Connection, error) { + conn, err := t.inner.Connect(ctx) + if err != nil { + return nil, err + } + return &sniffConn{Connection: conn, owner: t}, nil +} + +// called 报告某个方法是否出现过。 +func (t *sniffTransport) called(method string) bool { + t.mu.Lock() + defer t.mu.Unlock() + for _, m := range t.seen { + if m == method { + return true + } + } + return false +} + +// methods 返回已见方法快照。 +func (t *sniffTransport) methods() []string { + t.mu.Lock() + defer t.mu.Unlock() + return append([]string(nil), t.seen...) +} + +// record 记录一个发出的方法名。 +func (t *sniffTransport) record(method string) { + t.mu.Lock() + defer t.mu.Unlock() + t.seen = append(t.seen, method) +} + +// sniffConn 拦截写入的请求消息，记录方法名后交给内层连接。 +type sniffConn struct { + mcpsdk.Connection + owner *sniffTransport +} + +// Write 记录请求方法后转发。 +func (c *sniffConn) Write(ctx context.Context, msg sdkjsonrpc.Message) error { + if req, ok := msg.(*sdkjsonrpc.Request); ok { + c.owner.record(req.Method) + } + return c.Connection.Write(ctx, msg) +} diff --git a/internal/capability/mcp/mcp.go b/internal/capability/mcp/mcp.go index 7609691..3b8cca7 100644 --- a/internal/capability/mcp/mcp.go +++ b/internal/capability/mcp/mcp.go @@ -1,264 +1,172 @@ -// Package mcp provides a minimal Model Context Protocol (MCP) client and the -// auto-registration glue that surfaces MCP server tools as capabilities. +// Package mcp 用官方 Model Context Protocol Go SDK 接入外部 MCP server， +// 并把服务器暴露的工具自动注册进 capability registry。 // -// An MCP server is an external tool provider. This package speaks JSON-RPC 2.0 -// 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). +// 协议层（协议版本协商、stdio / SSE / Streamable HTTP 传输、keepalive ping、 +// notifications/tools/list_changed 订阅）全部交给官方 SDK +// （github.com/modelcontextprotocol/go-sdk/mcp）——本包不自造协议，只做四件事： // -// 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. +// 1. 配置 → 传输选择（MCPConfig → SDK Transport，含 APIKey 鉴权）； +// 2. 工具自动发现 → 注册进 capability registry（Registry）； +// 3. 保留服务器声明的 capabilities 并推断可用能力族（ServerFeatures， +// 裁决 10.4.6 / design/09 §1.1「声明了才算可用」）； +// 4. tools/list input schema 缓存（SchemaCache）+ 调用前校验与默认值填充。 // -// 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 -// required. -// -// A caller usually wires it up as: +// 典型接法： // // reg := registry.New() -// mcpReg := mcp.New(reg, mcp.MCPConfig{Name: "calc", Command: "mcp-calc"}) +// mcpReg := mcp.NewRegistry(reg, mcp.MCPConfig{Name: "calc", Command: "mcp-calc"}) // if err := mcpReg.Register(ctx); err != nil { -// log.Fatal(err) +// log.Warn("mcp server skipped", "error", err) // } // defer mcpReg.Unregister(ctx) // -// which connects to the server, discovers its tools, and registers them. Call -// sites then invoke the tools by their MCP name through the registry. +// 注册后的工具就是普通 sdk.Tool：Dispatcher 每轮实时读 registry， +// 因此热接入的 MCP 工具下一轮模型请求即作为 Request.Tools 注入，模型可直接调用。 package mcp import ( "context" "errors" + "fmt" "io" "log/slog" - "slices" "sync" "sync/atomic" "time" + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "git.yeij.top/AskaEth/Cyrene_v2/internal/logger" ) -// 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). +// clientName / clientVersion 是 initialize 时上报给 MCP server 的 clientInfo。 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 + clientName = "cyrene" + clientVersion = "0.1.0" ) -// 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} +// keepAliveFailureThreshold 是 keepalive 关闭会话前容忍的连续 ping 失败次数默认值。 +const keepAliveFailureThreshold = 3 -// 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. +// 本包对外哨兵错误，一律用 errors.Is 判定。 var ( - // ErrNotConnected is returned when an operation requires an established - // connection but the client is not yet connected or is closed. + // ErrNotConnected 表示操作要求已建立的连接，但客户端尚未连接或已关闭。 ErrNotConnected = errors.New("mcp: not connected") - // ErrClosed is returned when operating on a client or transport that has - // been closed. + // ErrClosed 表示在已关闭的客户端上执行操作。 ErrClosed = errors.New("mcp: closed") - // ErrNoTransport is returned when the configuration selects no transport - // (neither Command nor URL is set). + // ErrNoTransport 表示配置既没有 Command 也没有 URL，选不出传输。 ErrNoTransport = errors.New("mcp: no transport configured") - // ErrRPC is returned when the MCP server replies with a JSON-RPC error - // 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 表示客户端与服务器无法就协议版本达成一致。 ErrUnsupportedProtocolVersion = errors.New("mcp: unsupported protocol version") + // ErrCapabilityNotDeclared 表示调用了服务器在 initialize 里未声明的能力族 + // （design/09 §1.1「声明了才算可用，未声明不得调用」）。 + ErrCapabilityNotDeclared = errors.New("mcp: capability not declared by server") ) -// 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. +// MCPConfig 描述一台 MCP server 的连接方式。 +// TransportType 非空时按其选择传输；为空则沿用历史推断： +// 有 Command 走 stdio，有 URL 走 SSE（TransportType 可显式改走 Streamable HTTP）。 type MCPConfig struct { - // Name is the human-readable label of the server, used in logs and as the - // prefix for registering its tools uniquely. + // Name 是服务器标签，用于日志与工具去重前缀。 Name string `env:"MCP_NAME"` - // Command is the executable path to spawn for a stdio MCP server. + // Command 是 stdio 模式下要拉起的可执行文件。 Command string `env:"MCP_COMMAND"` - // Args are the arguments passed to Command. + // Args 是传给 Command 的参数。 Args []string `env:"MCP_ARGS"` - // URL is the endpoint of an MCP server (SSE stream endpoint, or the - // Streamable HTTP endpoint when TransportType is "http"). + // URL 是 SSE 端点（TransportType="sse"）或 Streamable HTTP 端点（"http"/"streamable"）。 URL string `env:"MCP_URL"` - // TransportType selects the transport explicitly: "stdio", "sse" or - // "http" (Streamable HTTP). Empty means infer from Command/URL. + // TransportType 显式选择传输："stdio"、"sse"、"http"（等价 "streamable"）。 TransportType string `env:"MCP_TRANSPORT"` - // APIKey is the bearer credential sent to an http/sse server as an - // Authorization header. Empty means the server needs no auth. + // APIKey 是 http/sse 服务器用的 bearer 凭据，作为 Authorization 头发送。 APIKey string `env:"MCP_API_KEY"` - // Env is extra environment for a stdio subprocess, merged over the parent - // process environment (extra keys win). It is ignored by the HTTP transports. + // Env 是 stdio 子进程的额外环境变量，叠加在父进程环境之上（同名以这里为准）。 Env map[string]string `env:"MCP_ENV"` - // Enabled reports whether this server should be activated at startup. + // Enabled 表示该服务器是否在启动时接入。 Enabled bool `env:"MCP_ENABLED"` } -// Config holds the set of configured MCP servers. +// Config 是一组 MCP server 配置。 type Config struct { - // Servers is the ordered list of MCP server configurations. + // Servers 是按顺序排列的服务器配置。 Servers []MCPConfig } -// ServerInfo describes the MCP server reported by the initialize handshake. +// ServerFeatures 是从服务器 initialize 声明里推断出的能力族可用性。 +// 规则见 featuresFrom：声明了才算可用，未声明不得调用（design/09 §1.1）。 +type ServerFeatures struct { + // Tools 表示服务器声明提供工具（tools/list、tools/call 可用）。 + Tools bool + // Prompts 表示服务器声明提供提示模板。 + Prompts bool + // Resources 表示服务器声明提供资源。 + Resources bool + // Logging 表示服务器声明会发日志消息。 + Logging bool + // Completions 表示服务器声明支持参数自动补全。 + Completions bool + // Sampling 表示允许服务器回调宿主代跑 LLM：宿主配了回调， + // 且服务器声明过能力族（即确实是会主动交互的对端）才为真。 + Sampling bool + // Elicitation 表示允许服务器回调宿主向用户追加提问，判定同 Sampling。 + Elicitation bool + // Roots 表示允许向服务器暴露工作区根目录，判定同 Sampling。 + Roots bool +} + +// ServerInfo 是 initialize 握手的结果，服务器声明的 capabilities 原样保留。 type ServerInfo struct { - // Name is the server-provided name. + // Name 是服务器自报名称。 Name string - // Version is the server-provided version string. + // Version 是服务器自报版本。 Version string + // Instructions 是服务器给出的使用说明（可注入系统提示）。 + Instructions string + // ProtocolVersion 是协商后实际使用的 MCP 协议版本。 + ProtocolVersion string + // Capabilities 是服务器声明的 capabilities 原文（未做删改），对客户端可读。 + Capabilities map[string]any + // Features 是按 Capabilities 推断出的可用能力族。 + Features ServerFeatures } -// ToolInfo is a tool exposed by an MCP server. -type ToolInfo struct { - // Name is the tool identifier used to invoke it. - Name string - // Description explains what the tool does. - Description string - // InputSchema is the tool's JSON Schema, typically - // {"type":"object","properties":{...}}. - InputSchema map[string]any -} +// SamplingHandler 由宿主提供：服务器请求客户端代跑一次 LLM（sampling/createMessage）时调用。 +type SamplingHandler func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) -// ContentBlock is a single item in a tool call result (MCP content block). -type ContentBlock struct { - // Type is the block type, most commonly "text". - Type string `json:"type"` - // Text is the textual payload. - Text string `json:"text"` -} +// ElicitationHandler 由宿主提供：服务器向用户索要补充输入（elicitation/create）时调用。 +type ElicitationHandler func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) -// ToolResult is the structured result of a tools/call. -type ToolResult struct { - // Content holds the response content blocks. - Content []ContentBlock - // IsError reports whether the server flagged the call as an error. - IsError bool -} - -// Client is a minimal MCP client bound to a single server described by cfg. It -// is safe for concurrent use; all calls are serialized internally. +// Client 是一台 MCP server 的客户端。它内部持有一个官方 SDK 会话， +// 并发安全：Connect/Close 互斥，其余调用由 SDK 串行化。 type Client struct { - cfg MCPConfig - log *slog.Logger - tr Transport - mu sync.Mutex - state state - info ServerInfo + cfg MCPConfig + log *slog.Logger - // 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 + mu sync.Mutex + state state + session *mcpsdk.ClientSession + info ServerInfo - // preferredProtocolVersion, when non-empty, overrides latestProtocolVersion - // as the version requested first during initialize. Tests use it to drive the - // negotiation path deterministically. + // tr 是注入或按 cfg 选出的传输；schemas 是 tools/list input schema 缓存。 + tr mcpsdk.Transport + schemas *SchemaCache + + // preferredProtocolVersion 非空时固定 initialize 请求的协议版本（测试/排障用）。 preferredProtocolVersion string - - // keepAlive is the ping interval for the connection keepalive; zero disables - // it. keepAliveFailures is the tolerated consecutive-failure count. + // keepAlive 是 keepalive ping 间隔，零值关闭；keepAliveFailures 是容忍失败数。 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. + sampling SamplingHandler + elicitation ElicitationHandler + roots []*mcpsdk.Root + + // onToolListChanged 在服务器发 notifications/tools/list_changed 时被调用 + // （新 goroutine 里，已被 refreshPending 合并），用于增量刷新 registry。 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 是客户端生命周期。 type state int const ( @@ -267,63 +175,10 @@ const ( stateClosed ) -// Transport is the JSON-RPC message exchange with an MCP server over a concrete -// 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. - Start(ctx context.Context) error - // Call sends a JSON-RPC request and decodes the matching response into - // result. Params is marshaled into the request body; result is unmarshaled - // from the response's result field. - Call(ctx context.Context, method string, params any, result any) error - // Close releases every underlying resource (subprocess, streams, sockets). - Close() error -} - -// notifier is implemented by transports that can send a JSON-RPC notification -// (a request with no id and no response). -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. +// Option 定制 New 构建的 Client。 type Option func(*Client) -// WithLogger overrides the client's logger. The given logger is tagged with the -// source capability/mcp. +// WithLogger 覆盖客户端日志器；给定的日志器会被标上 source=capability/mcp。 func WithLogger(l *slog.Logger) Option { return func(c *Client) { if l != nil { @@ -332,19 +187,15 @@ func WithLogger(l *slog.Logger) Option { } } -// WithTransport overrides the transport used by the client. This is primarily -// for tests that inject a fake transport instead of spawning a subprocess or -// dialing SSE. -func WithTransport(t Transport) Option { +// WithTransport 注入传输（官方 SDK 的 Transport），主要用于测试。 +func WithTransport(t mcpsdk.Transport) Option { return func(c *Client) { c.tr = t } } -// 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. +// WithKeepAlive 打开 keepalive：每 interval 发一次 ping，连续失败 keepAliveFailures +// 次即关闭会话（由 SDK 执行）。间隔为零则关闭 keepalive；失败次数缺省用包默认值。 func WithKeepAlive(interval time.Duration, keepAliveFailures ...int) Option { failures := keepAliveFailureThreshold if len(keepAliveFailures) > 0 && keepAliveFailures[0] > 0 { @@ -356,33 +207,60 @@ func WithKeepAlive(interval time.Duration, keepAliveFailures ...int) Option { } } -// 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. +// WithToolListChangedHandler 注册工具清单变更回调，由 registry 用于增量刷新。 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. +// WithProtocolVersion 固定 initialize 请求的协议版本；不设置则用 SDK 最新版本， +// 服务器拒绝时由 SDK 自动回退到双方都支持的版本。 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. +// WithSchemaCache 复用外部 schema 缓存；不设置则每个客户端自建一份。 +func WithSchemaCache(cache *SchemaCache) Option { + return func(c *Client) { + if cache != nil { + c.schemas = cache + } + } +} + +// WithSamplingHandler 注册 sampling 回调。设了才向服务器宣告 sampling 能力； +// 服务器未声明能力族时回调会被拒绝（ErrCapabilityNotDeclared）。 +func WithSamplingHandler(h SamplingHandler) Option { + return func(c *Client) { + c.sampling = h + } +} + +// WithElicitationHandler 注册 elicitation 回调。设了才向服务器宣告 elicitation 能力。 +func WithElicitationHandler(h ElicitationHandler) Option { + return func(c *Client) { + c.elicitation = h + } +} + +// WithRoots 注册工作区根目录。设了才向服务器宣告 roots 能力。 +func WithRoots(roots ...*mcpsdk.Root) Option { + return func(c *Client) { + c.roots = append([]*mcpsdk.Root(nil), roots...) + } +} + +// New 构建 cfg 对应的 Client，默认用丢弃日志器并自动选传输； +// 用 WithLogger / WithTransport 覆盖。 func New(cfg MCPConfig, opts ...Option) *Client { c := &Client{ - cfg: cfg, - log: defaultLogger(), - state: stateNew, + cfg: cfg, + log: defaultLogger(), + state: stateNew, + schemas: NewSchemaCache(), } for _, opt := range opts { opt(c) @@ -390,19 +268,56 @@ func New(cfg MCPConfig, opts ...Option) *Client { return c } -// defaultLogger returns a discard logger tagged with the mcp source. Hosts -// should prefer WithLogger so records are observable. +// defaultLogger 返回带 capability/mcp 来源标签的丢弃日志器；宿主应传 WithLogger。 func defaultLogger() *slog.Logger { base := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelWarn})) return logger.WithSource(base, "capability/mcp") } -// transport returns the established transport, or nil when not connected. -func (c *Client) transport() Transport { +// Name 返回服务器标签。 +func (c *Client) Name() string { return c.cfg.Name } + +// ServerInfo 返回 initialize 握手结果；未连接时为零值。 +func (c *Client) ServerInfo() ServerInfo { c.mu.Lock() defer c.mu.Unlock() - if c.state != stateConnected || c.tr == nil { + return c.info +} + +// Features 返回按服务器声明推断出的能力族可用性；未连接时全为 false。 +func (c *Client) Features() ServerFeatures { + c.mu.Lock() + defer c.mu.Unlock() + return c.info.Features +} + +// SchemaCache 返回该客户端使用的 input schema 缓存。 +func (c *Client) SchemaCache() *SchemaCache { return c.schemas } + +// Session 返回底层官方 SDK 会话（未连接时为 nil），供需要 SDK 原生能力的调用方使用。 +func (c *Client) Session() *mcpsdk.ClientSession { + c.mu.Lock() + defer c.mu.Unlock() + if c.state != stateConnected { return nil } - return c.tr + return c.session +} + +// connected 返回已建立的 SDK 会话，未连接时返回 nil。 +func (c *Client) connected() *mcpsdk.ClientSession { + return c.Session() +} + +// begin 做 Connect 前的状态检查：已连接返回 same=true（幂等），已关闭返回错误。 +func (c *Client) begin() (same bool, err error) { + c.mu.Lock() + defer c.mu.Unlock() + switch c.state { + case stateClosed: + return false, fmt.Errorf("mcp connect %q: %w", c.cfg.Name, ErrClosed) + case stateConnected: + return true, nil + } + return false, nil } diff --git a/internal/capability/mcp/mcp_auth_test.go b/internal/capability/mcp/mcp_auth_test.go deleted file mode 100644 index bb2366c..0000000 --- a/internal/capability/mcp/mcp_auth_test.go +++ /dev/null @@ -1,47 +0,0 @@ -package mcp - -import ( - "context" - "net/http" - "testing" -) - -func TestEnvWith_OverridesAndPreserves(t *testing.T) { - t.Parallel() - base := []string{"PATH=/usr/bin", "HOME=/root"} - got := envWith(base, map[string]string{"HOME": "/tmp", "MCP_TOKEN": "abc"}) - if len(got) != 3 { - t.Fatalf("len = %d, want 3 (%v)", len(got), got) - } - seen := map[string]bool{} - for _, kv := range got { - seen[kv] = true - } - if !seen["PATH=/usr/bin"] || !seen["HOME=/tmp"] || !seen["MCP_TOKEN=abc"] { - t.Fatalf("merged env = %v", got) - } - if len(envWith(base, nil)) != len(base) { - t.Fatal("empty extra changed the base environment") - } -} - -func TestHTTPTransport_SendsAPIKey(t *testing.T) { - t.Parallel() - got := make(chan string, 1) - srv := httpTestServer(t, func(w http.ResponseWriter, r *http.Request) { - got <- r.Header.Get("Authorization") - httpJSONResponse(w, r) - }) - defer srv.Close() - tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL, APIKey: "sk-1234"}, defaultLogger()) - defer tr.Close() - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - if err := tr.Call(context.Background(), "tools/list", nil, &listToolsResult{}); err != nil { - t.Fatalf("Call: %v", err) - } - if auth := <-got; auth != "Bearer sk-1234" { - t.Fatalf("Authorization = %q, want bearer key", auth) - } -} diff --git a/internal/capability/mcp/registry_test.go b/internal/capability/mcp/registry_test.go deleted file mode 100644 index 6f4b255..0000000 --- a/internal/capability/mcp/registry_test.go +++ /dev/null @@ -1,171 +0,0 @@ -package mcp - -import ( - "context" - "errors" - "testing" - "time" - - "git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry" -) - -func sampleTools() []ToolInfo { - return []ToolInfo{ - {Name: "add", Description: "Adds numbers", InputSchema: map[string]any{"type": "object"}}, - {Name: "sub", Description: "Subtracts numbers", InputSchema: map[string]any{"type": "object"}}, - } -} - -func TestRegistry_Register_AutoRegisters(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) - } - for _, name := range []string{"add", "sub"} { - cap, err := reg.Get(name) - if err != nil { - t.Fatalf("Get(%q): %v", name, err) - } - if cap.Kind != registry.KindTool { - t.Fatalf("Kind(%q) = %q, want %q", name, cap.Kind, registry.KindTool) - } - } - if !ft.hasMethod("tools/list") { - t.Fatalf("tools/list was not called; methods = %v", ft.calledMethods()) - } -} - -func TestRegistry_Register_PropagatesConnectError(t *testing.T) { - t.Parallel() - ft := &fakeTransport{startErr: errors.New("boom")} - reg := registry.New() - r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft)) - if err := r.Register(context.Background()); err == nil { - t.Fatal("Register succeeded, want connect error") - } -} - -func TestRegistry_Register_DuplicateTool(t *testing.T) { - t.Parallel() - ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}} - reg := registry.New() - // Pre-register a tool with the same name so the MCP registration collides. - _ = reg.RegisterTool(&testTool{name: "add"}) - r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft)) - err := r.Register(context.Background()) - if !errors.Is(err, registry.ErrDuplicate) { - t.Fatalf("Register error = %v, want registry.ErrDuplicate", err) - } -} - -func TestRegistry_Register_NilRegistry(t *testing.T) { - t.Parallel() - ft := &fakeTransport{} - r := NewRegistry(nil, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft)) - if err := r.Register(context.Background()); err == nil { - t.Fatal("Register with nil registry succeeded, want error") - } -} - -func TestRegistry_Unregister_RemovesTools(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) - } - if err := r.Unregister(context.Background()); err != nil { - t.Fatalf("Unregister: %v", err) - } - for _, name := range []string{"add", "sub"} { - if _, err := reg.Get(name); !errors.Is(err, registry.ErrNotFound) { - t.Fatalf("Get(%q) after Unregister = %v, want ErrNotFound", name, err) - } - } - if ft.closes != 1 { - t.Fatalf("closes = %d, want 1 (Unregister must close the client)", ft.closes) - } -} - -func TestRegistry_Unregister_Idempotent(t *testing.T) { - t.Parallel() - ft := &fakeTransport{} - 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) - } - if err := r.Unregister(context.Background()); err != nil { - t.Fatalf("first Unregister: %v", err) - } - if err := r.Unregister(context.Background()); err != nil { - t.Fatalf("second Unregister: %v", err) - } -} - -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 -} - -func (t *testTool) Name() string { return t.name } -func (t *testTool) Description() string { return "desc " + t.name } -func (t *testTool) Parameters() (map[string]any, error) { return map[string]any{}, nil } -func (t *testTool) Execute(_ context.Context, _ any) (any, error) { return nil, nil } diff --git a/internal/capability/mcp/rpc.go b/internal/capability/mcp/rpc.go deleted file mode 100644 index 0406d71..0000000 --- a/internal/capability/mcp/rpc.go +++ /dev/null @@ -1,53 +0,0 @@ -package mcp - -import ( - "encoding/json" - "fmt" -) - -// rpcRequest is a JSON-RPC 2.0 request. Notifications omit ID so the field is -// tagged omitempty. Params is any so callers pass a struct or a JSON object. -type rpcRequest struct { - JSONRPC string `json:"jsonrpc"` - ID any `json:"id,omitempty"` - Method string `json:"method"` - Params any `json:"params,omitempty"` -} - -// rpcResponse is a JSON-RPC 2.0 response. Result is kept raw until the caller -// 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"` -} - -// rpcError is a JSON-RPC error object. It satisfies error so it can be wrapped -// and surfaced with %w. -type rpcError struct { - Code int `json:"code"` - Message string `json:"message"` - Data any `json:"data,omitempty"` -} - -// Error implements error. -func (e *rpcError) Error() string { - return fmt.Sprintf("code %d: %s", e.Code, e.Message) -} - -// encodeRequest marshals a JSON-RPC request. An id of nil produces a -// notification (no id field). -func encodeRequest(id any, method string, params any) ([]byte, error) { - return json.Marshal(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}) -} - -// decodeIDString reports whether res carries the string id expected for a JSON -// string id request. -func decodeIDString(res rpcResponse, id string) bool { - s, ok := res.ID.(string) - return ok && s == id -} diff --git a/internal/capability/mcp/schema.go b/internal/capability/mcp/schema.go new file mode 100644 index 0000000..ce527ce --- /dev/null +++ b/internal/capability/mcp/schema.go @@ -0,0 +1,84 @@ +package mcp + +import ( + "encoding/json" + "fmt" + + "github.com/google/jsonschema-go/jsonschema" +) + +// schemaMap 把 SDK 传来的 input schema 规整成 map[string]any；nil 或无法识别时 +// 退化为「无约束的对象 schema」，保证 Parameters() 永远能返回可用形状。 +func schemaMap(v any) map[string]any { + switch s := v.(type) { + case nil: + return emptyObjectSchema() + case map[string]any: + return s + case json.RawMessage: + return decodeSchema(s) + case []byte: + return decodeSchema(s) + default: + data, err := json.Marshal(v) + if err != nil { + return emptyObjectSchema() + } + return decodeSchema(data) + } +} + +// emptyObjectSchema 返回无约束的空对象 schema。 +func emptyObjectSchema() map[string]any { + return map[string]any{"type": "object", "properties": map[string]any{}} +} + +// decodeSchema 把 JSON 文本解成 schema map，失败时退回空对象 schema。 +func decodeSchema(data []byte) map[string]any { + out := map[string]any{} + if err := json.Unmarshal(data, &out); err != nil { + return emptyObjectSchema() + } + return out +} + +// resolveSchema 把原始 schema 解析成可校验、可补默认值的已解析 schema。 +func resolveSchema(raw map[string]any) (*jsonschema.Resolved, error) { + data, err := json.Marshal(raw) + if err != nil { + return nil, fmt.Errorf("marshal input schema: %w", err) + } + var s jsonschema.Schema + if err := json.Unmarshal(data, &s); err != nil { + return nil, fmt.Errorf("parse input schema: %w", err) + } + resolved, err := s.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) + if err != nil { + return nil, fmt.Errorf("resolve input schema: %w", err) + } + return resolved, nil +} + +// applySchema 按已解析的 input schema 先补默认值、再校验参数。 +// resolved 为 nil 表示 schema 不可解析，此时降级为不校验。 +// ApplyDefaults 对病态默认值会 panic，这里兜住转成错误，避免拖垮调用方。 +func applySchema(resolved *jsonschema.Resolved, args map[string]any) (err error) { + if resolved == nil { + return nil + } + if args == nil { + args = map[string]any{} + } + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("input schema default application panicked: %v", r) + } + }() + if err = resolved.ApplyDefaults(&args); err != nil { + return fmt.Errorf("apply input schema defaults: %w", err) + } + if err = resolved.Validate(args); err != nil { + return fmt.Errorf("arguments do not match input schema: %w", err) + } + return nil +} diff --git a/internal/capability/mcp/schemacache.go b/internal/capability/mcp/schemacache.go new file mode 100644 index 0000000..8f64b67 --- /dev/null +++ b/internal/capability/mcp/schemacache.go @@ -0,0 +1,142 @@ +package mcp + +import ( + "strings" + "sync" + + "github.com/google/jsonschema-go/jsonschema" +) + +// SchemaCache 缓存 tools/list 里已解析的 input schema（design/09 §1.1）。 +// +// 键 = 服务器名 + 工具名；收到 notifications/tools/list_changed 时按服务器整表失效。 +// 命中时 CallTool 直接用缓存 schema 校验参数并补默认值，不重复拉取/解析 schema。 +// 并发安全；命中/未命中计数对调用方可读，便于验证缓存真的生效。 +type SchemaCache struct { + mu sync.Mutex + entries map[string]*cachedSchema + hits int + misses int +} + +// cachedSchema 是一条缓存：原始 schema（模型可见的形状）+ 已解析结果。 +// resolved 为 nil 表示 schema 解析失败，调用时降级为不校验。 +type cachedSchema struct { + raw map[string]any + resolved *jsonschema.Resolved +} + +// NewSchemaCache 构建空缓存。 +func NewSchemaCache() *SchemaCache { + return &SchemaCache{entries: map[string]*cachedSchema{}} +} + +// Put 解析并缓存一条 input schema。解析失败时仍缓存原始 schema（校验降级）并返回 +// 解析错误，由调用方决定是否告警。 +func (c *SchemaCache) Put(server, tool string, schema map[string]any) error { + if schema == nil { + schema = emptyObjectSchema() + } + entry := &cachedSchema{raw: schema} + resolved, err := resolveSchema(schema) + entry.resolved = resolved + + c.mu.Lock() + c.entries[cacheKey(server, tool)] = entry + c.mu.Unlock() + return err +} + +// Get 返回缓存的原始 input schema。 +func (c *SchemaCache) Get(server, tool string) (map[string]any, bool) { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[cacheKey(server, tool)] + if !ok { + return nil, false + } + return entry.raw, true +} + +// Prepare 返回「已补默认值、已校验」的参数： +// +// - 命中缓存：用缓存 schema 补默认值 + 校验（计一次命中），不再重复拉取/解析； +// - 未命中：有 fallback（本次 tools/list 发现结果）则解析后缓存再用；没有可用 +// schema 时原样返回（计一次未命中），即首次工具清单到达前的直接调用降级为不校验。 +// +// 传给服务器的是返回的参数值；args 为 nil 时按空对象处理。 +func (c *SchemaCache) Prepare(server, tool string, fallback, args map[string]any) (map[string]any, error) { + entry := c.lookup(server, tool) + if entry == nil { + if fallback == nil { + return args, nil + } + // 解析失败不阻断调用：Put 已把原始 schema 放进缓存，resolved 为 nil 即不校验。 + _ = c.Put(server, tool, fallback) + entry = c.lookup(server, tool) + } + if entry == nil || entry.resolved == nil { + return args, nil + } + if args == nil { + args = map[string]any{} + } + if err := applySchema(entry.resolved, args); err != nil { + return nil, err + } + return args, nil +} + +// lookup 取一条缓存并记账（命中/未命中）。 +func (c *SchemaCache) lookup(server, tool string) *cachedSchema { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[cacheKey(server, tool)] + if !ok { + c.misses++ + return nil + } + c.hits++ + return entry +} + +// Invalidate 让某台服务器的全部缓存失效；工具清单变化时调用。 +func (c *SchemaCache) Invalidate(server string) { + prefix := server + keySeparator + c.mu.Lock() + defer c.mu.Unlock() + for k := range c.entries { + if strings.HasPrefix(k, prefix) { + delete(c.entries, k) + } + } +} + +// Clear 清空所有缓存。 +func (c *SchemaCache) Clear() { + c.mu.Lock() + defer c.mu.Unlock() + c.entries = map[string]*cachedSchema{} +} + +// Len 返回当前缓存条目数。 +func (c *SchemaCache) Len() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.entries) +} + +// Stats 返回命中次数与未命中次数。 +func (c *SchemaCache) Stats() (hits, misses int) { + c.mu.Lock() + defer c.mu.Unlock() + return c.hits, c.misses +} + +// keySeparator 分隔服务器名与工具名，避免不同服务器同名工具互相覆盖。 +const keySeparator = "\x00" + +// cacheKey 拼接缓存键。 +func cacheKey(server, tool string) string { + return server + keySeparator + tool +} diff --git a/internal/capability/mcp/schemacache_test.go b/internal/capability/mcp/schemacache_test.go new file mode 100644 index 0000000..17dfa5b --- /dev/null +++ b/internal/capability/mcp/schemacache_test.go @@ -0,0 +1,149 @@ +package mcp + +import ( + "testing" +) + +// TestSchemaCache_PutGetInvalidate 覆盖缓存基本语义与按服务器失效。 +func TestSchemaCache_PutGetInvalidate(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + schema := map[string]any{"type": "object", "properties": map[string]any{"a": map[string]any{"type": "number"}}} + if err := cache.Put("s1", "tool", schema); err != nil { + t.Fatalf("Put: %v", err) + } + if err := cache.Put("s2", "tool", schema); err != nil { + t.Fatalf("Put(s2): %v", err) + } + if _, ok := cache.Get("s1", "tool"); !ok { + t.Fatal("Get(s1,tool) 未命中") + } + if cache.Len() != 2 { + t.Fatalf("Len = %d, 期望 2", cache.Len()) + } + + // 只失效 s1，s2 的同名工具必须留下。 + cache.Invalidate("s1") + if _, ok := cache.Get("s1", "tool"); ok { + t.Fatal("Invalidate(s1) 后仍然命中") + } + if _, ok := cache.Get("s2", "tool"); !ok { + t.Fatal("Invalidate(s1) 误删了 s2 的缓存") + } + + cache.Clear() + if cache.Len() != 0 { + t.Fatalf("Clear 后 Len = %d, 期望 0", cache.Len()) + } +} + +// TestSchemaCache_Prepare_AppliesDefaultsAndValidates 命中缓存时补默认值并校验。 +func TestSchemaCache_Prepare_AppliesDefaultsAndValidates(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string", "default": "world"}, + "n": map[string]any{"type": "number"}, + }, + "required": []string{"n"}, + } + if err := cache.Put("s", "t", schema); err != nil { + t.Fatalf("Put: %v", err) + } + + out, err := cache.Prepare("s", "t", nil, map[string]any{"n": 1}) + if err != nil { + t.Fatalf("Prepare: %v", err) + } + if out["name"] != "world" { + t.Fatalf("默认值未生效：%v", out) + } + + if _, err := cache.Prepare("s", "t", nil, map[string]any{"n": "x"}); err == nil { + t.Fatal("类型错误的参数没有被拦下") + } + if _, err := cache.Prepare("s", "t", nil, map[string]any{}); err == nil { + t.Fatal("缺必填参数没有被拦下") + } +} + +// TestSchemaCache_Prepare_HitsCacheWithoutReparsing 证明「命中缓存」这条要求： +// 发现阶段已缓存 → 后续调用全是命中，未命中计数为 0（没有重复拉取/解析）。 +func TestSchemaCache_Prepare_HitsCacheWithoutReparsing(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + if err := cache.Put("s", "t", emptyObjectSchema()); err != nil { + t.Fatalf("Put: %v", err) + } + for i := 0; i < 3; i++ { + if _, err := cache.Prepare("s", "t", nil, map[string]any{}); err != nil { + t.Fatalf("Prepare #%d: %v", i, err) + } + } + hits, misses := cache.Stats() + if hits != 3 || misses != 0 { + t.Fatalf("Stats = (%d hits, %d misses), 期望 (3, 0)", hits, misses) + } +} + +// TestSchemaCache_Prepare_MissFillsFromFallback 缓存未命中时用发现结果兜底填入， +// 且失效后第一次调用计一次未命中。 +func TestSchemaCache_Prepare_MissFillsFromFallback(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + fallback := map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string", "default": "fallback"}, + }, + } + out, err := cache.Prepare("s", "t", fallback, map[string]any{}) + if err != nil { + t.Fatalf("Prepare: %v", err) + } + if out["name"] != "fallback" { + t.Fatalf("兜底 schema 的默认值未生效：%v", out) + } + if _, misses := cache.Stats(); misses != 1 { + t.Fatalf("misses = %d, 期望 1", misses) + } + if cache.Len() != 1 { + t.Fatalf("未命中后没有回填缓存：Len = %d", cache.Len()) + } + + cache.Invalidate("s") + if _, err := cache.Prepare("s", "t", nil, map[string]any{"x": 1}); err != nil { + t.Fatalf("无 schema 时应降级为不校验：%v", err) + } +} + +// TestSchemaCache_Put_BadSchemaDegrades 无法解析的 schema 只降级为「不校验」，不阻断调用。 +func TestSchemaCache_Put_BadSchemaDegrades(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + bad := map[string]any{"type": "object", "properties": map[string]any{"a": map[string]any{"$ref": "#/$defs/does-not-exist"}}} + err := cache.Put("s", "t", bad) + if err == nil { + t.Fatal("病态 schema 未报错，期望解析失败可见") + } + if _, ok := cache.Get("s", "t"); !ok { + t.Fatal("解析失败时原始 schema 没有缓存下来（应降级保留）") + } + if _, err := cache.Prepare("s", "t", nil, map[string]any{"whatever": true}); err != nil { + t.Fatalf("降级路径不应报错：%v", err) + } +} + +// TestSchemaCache_Prepare_NilSchemaSafe nil schema 走空对象兜底，不 panic。 +func TestSchemaCache_Prepare_NilSchemaSafe(t *testing.T) { + t.Parallel() + cache := NewSchemaCache() + if err := cache.Put("s", "t", nil); err != nil { + t.Fatalf("Put(nil): %v", err) + } + if _, err := cache.Prepare("s", "t", nil, nil); err != nil { + t.Fatalf("Prepare(nil args): %v", err) + } +} diff --git a/internal/capability/mcp/sse.go b/internal/capability/mcp/sse.go deleted file mode 100644 index 9b483f1..0000000 --- a/internal/capability/mcp/sse.go +++ /dev/null @@ -1,283 +0,0 @@ -package mcp - -import ( - "bufio" - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "log/slog" - "net/http" - "strconv" - "strings" - "sync" -) - -// sseTransport speaks JSON-RPC 2.0 to an MCP server over SSE. It opens a -// long-lived GET stream for server-to-client messages and POSTs client-to-server -// requests to the message endpoint (discovered via an endpoint event, falling -// back to the configured URL). -type sseTransport struct { - cfg MCPConfig - log *slog.Logger - url string - endpoint string - mu sync.Mutex - nextID int - pending map[string]chan []byte - resp *http.Response - body io.ReadCloser - closeCh chan struct{} - closeOnce sync.Once - started bool - 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 -// back to the SSE stream. -const maxSSEResponse = 1 << 20 - -// newSSETransport builds an SSE transport for cfg. -func newSSETransport(cfg MCPConfig, log *slog.Logger) *sseTransport { - return &sseTransport{ - cfg: cfg, - log: log, - url: cfg.URL, - endpoint: cfg.URL, - pending: make(map[string]chan []byte), - closeCh: make(chan struct{}), - } -} - -// Start opens the SSE stream and launches the reader goroutine. -func (s *sseTransport) Start(ctx context.Context) error { - if err := ctx.Err(); err != nil { - return fmt.Errorf("mcp sse start: %w", err) - } - if s.cfg.URL == "" { - return fmt.Errorf("mcp sse: %w", ErrNoTransport) - } - req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, s.cfg.URL, nil) - if err != nil { - return fmt.Errorf("mcp sse get: %w", err) - } - req.Header.Set("Accept", "text/event-stream") - s.applyAuth(req) - resp, err := http.DefaultClient.Do(req) - if err != nil { - return fmt.Errorf("mcp sse connect %s: %w", s.cfg.URL, err) - } - if resp.StatusCode != http.StatusOK { - resp.Body.Close() - return fmt.Errorf("mcp sse connect %s: status %d", s.cfg.URL, resp.StatusCode) - } - s.resp = resp - s.body = resp.Body - s.started = true - go s.readLoop() - return nil -} - -// readLoop parses SSE events and dispatches message events to pending calls. -func (s *sseTransport) readLoop() { - r := bufio.NewReader(s.body) - var event, data string - for { - line, err := r.ReadString('\n') - if err != nil { - return - } - line = strings.TrimRight(line, "\r\n") - switch { - case line == "": - if data != "" { - s.dispatch(event, data) - } - 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 - } - } - } -} - -// dispatch handles a single SSE event. -func (s *sseTransport) dispatch(event, data string) { - switch event { - case "endpoint": - if e := strings.TrimSpace(data); e != "" { - s.mu.Lock() - s.endpoint = e - s.mu.Unlock() - } - case "message": - var resp rpcResponse - 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 - } - s.mu.Lock() - ch, has := s.pending[id] - delete(s.pending, id) - s.mu.Unlock() - if has { - ch <- []byte(data) - } - } -} - -// Call POSTs a request and waits for a response, either from the POST reply body -// or from the SSE stream. -func (s *sseTransport) Call(ctx context.Context, method string, params any, result any) error { - s.mu.Lock() - if !s.started { - s.mu.Unlock() - return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected) - } - s.nextID++ - id := strconv.Itoa(s.nextID) - ch := make(chan []byte, 1) - s.pending[id] = ch - endpoint := s.endpoint - s.mu.Unlock() - - defer func() { - s.mu.Lock() - delete(s.pending, id) - s.mu.Unlock() - }() - - reqBody, err := encodeRequest(id, method, params) - if err != nil { - return fmt.Errorf("mcp sse %s: %w", method, err) - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody)) - if err != nil { - return fmt.Errorf("mcp sse %s: %w", method, err) - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Accept", "application/json") - s.applyAuth(httpReq) - resp, err := http.DefaultClient.Do(httpReq) - if err != nil { - return fmt.Errorf("mcp sse %s post: %w", method, err) - } - body, _ := io.ReadAll(io.LimitReader(resp.Body, maxSSEResponse)) - resp.Body.Close() - - var direct rpcResponse - if json.Unmarshal(bytes.TrimSpace(body), &direct) == nil && decodeIDString(direct, id) { - return s.applyResponse(direct, method, result) - } - - select { - case data := <-ch: - var rr rpcResponse - if err := json.Unmarshal(data, &rr); err != nil { - return fmt.Errorf("mcp sse %s decode: %w", method, err) - } - return s.applyResponse(rr, method, result) - case <-ctx.Done(): - return fmt.Errorf("mcp sse %s: %w", method, ctx.Err()) - case <-s.closeCh: - return fmt.Errorf("mcp sse %s: %w", method, ErrClosed) - } -} - -// applyResponse decodes a JSON-RPC response into result, surfacing an error -// object if present. -func (s *sseTransport) applyResponse(resp rpcResponse, method string, result any) error { - if resp.Error != nil { - return fmt.Errorf("mcp sse %s: %w: %w", method, ErrRPC, resp.Error) - } - if result != nil { - if err := json.Unmarshal(resp.Result, result); err != nil { - return fmt.Errorf("mcp sse %s decode: %w", method, err) - } - } - return nil -} - -// notify POSTs a fire-and-forget JSON-RPC notification. -func (s *sseTransport) notify(ctx context.Context, method string, params any) error { - s.mu.Lock() - started := s.started - endpoint := s.endpoint - s.mu.Unlock() - if !started { - return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected) - } - reqBody, err := encodeRequest(nil, method, params) - if err != nil { - return fmt.Errorf("mcp notify %s: %w", method, err) - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody)) - if err != nil { - return fmt.Errorf("mcp notify %s: %w", method, err) - } - httpReq.Header.Set("Content-Type", "application/json") - s.applyAuth(httpReq) - resp, err := http.DefaultClient.Do(httpReq) - if err != nil { - return fmt.Errorf("mcp notify %s: %w", method, err) - } - defer resp.Body.Close() - _, _ = io.Copy(io.Discard, resp.Body) - return nil -} - -// applyAuth sets the bearer credential an authenticated server requires. It is -// a no-op when the config carries no API key. -func (s *sseTransport) applyAuth(req *http.Request) { - if s.cfg.APIKey != "" { - req.Header.Set("Authorization", "Bearer "+s.cfg.APIKey) - } -} - -// Close closes the SSE stream and any pending responses. -func (s *sseTransport) Close() error { - s.closeOnce.Do(func() { - close(s.closeCh) - if s.body != nil { - _ = s.body.Close() - } - if s.resp != nil && s.resp.Body != nil { - _ = s.resp.Body.Close() - } - s.mu.Lock() - s.pending = make(map[string]chan []byte) - s.mu.Unlock() - }) - return nil -} diff --git a/internal/capability/mcp/sse_test.go b/internal/capability/mcp/sse_test.go deleted file mode 100644 index e93a998..0000000 --- a/internal/capability/mcp/sse_test.go +++ /dev/null @@ -1,203 +0,0 @@ -package mcp - -import ( - "context" - "encoding/json" - "errors" - "io" - "net/http" - "net/http/httptest" - "testing" - "time" -) - -// postRespondServer serves a GET stream that ends immediately and answers every -// POST directly with a JSON-RPC response body (the non-streaming reply path). -func postRespondServer(t *testing.T) *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() - } - return - } - 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)) - })) -} - -// streamServer pushes responses to the POST caller back through the SSE stream. -func streamServer(t *testing.T) *httptest.Server { - t.Helper() - push := make(chan string, 8) - var srv *httptest.Server - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.Method { - case http.MethodGet: - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache") - f, _ := w.(http.Flusher) - _, _ = w.Write([]byte("event: endpoint\ndata: " + srv.URL + "\n\n")) - f.Flush() - for { - select { - case data := <-push: - _, _ = w.Write([]byte("event: message\ndata: " + data + "\n\n")) - f.Flush() - case <-r.Context().Done(): - return - } - } - default: - 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":"stream"}]}}` - push <- resp - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("ok")) - } - })) - return srv -} - -// idJSON renders a request id as a JSON string literal when it is a string. -func idJSON(id any) string { - if s, ok := id.(string); ok { - b, _ := json.Marshal(s) - return string(b) - } - return `null` -} - -func TestSSE_Call_DirectResponse(t *testing.T) { - t.Parallel() - srv := postRespondServer(t) - defer srv.Close() - tr := newSSETransport(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", Arguments: map[string]any{}}, &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 TestSSE_Call_MessageEvent(t *testing.T) { - t.Parallel() - srv := streamServer(t) - defer srv.Close() - tr := newSSETransport(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", Arguments: map[string]any{}}, &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 TestSSE_Start_Non200(t *testing.T) { - t.Parallel() - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - })) - defer srv.Close() - tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - if err := tr.Start(context.Background()); err == nil { - t.Fatal("Start succeeded on a non-200 response, want error") - } -} - -func TestSSE_Call_NotStarted(t *testing.T) { - t.Parallel() - tr := newSSETransport(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 TestSSE_Call_ContextTimeout(t *testing.T) { - t.Parallel() - // A server whose GET stream stays open and whose POST replies with no body - // and never pushes forces the Call to wait on its context deadline. - 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 - } - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("ok")) - })) - defer srv.Close() - - tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - defer tr.Close() - - ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) - defer cancel() - err := tr.Call(ctx, "tools/call", callToolParams{Name: "add"}, &callToolResult{}) - if !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("Call error = %v, want context.DeadlineExceeded", err) - } -} - -func TestSSE_Start_NoURL(t *testing.T) { - t.Parallel() - tr := newSSETransport(MCPConfig{Name: "s"}, defaultLogger()) - err := tr.Start(context.Background()) - if !errors.Is(err, ErrNoTransport) { - t.Fatalf("Start error = %v, want ErrNoTransport", err) - } -} - -func TestSSE_Close_ClosesPendingAndBody(t *testing.T) { - t.Parallel() - srv := postRespondServer(t) - defer srv.Close() - tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger()) - if err := tr.Start(context.Background()); err != nil { - t.Fatalf("Start: %v", err) - } - // Put a sentinel pending call in place. - tr.mu.Lock() - tr.pending["999"] = make(chan []byte, 1) - tr.mu.Unlock() - if err := tr.Close(); err != nil { - t.Fatalf("Close: %v", err) - } - tr.mu.Lock() - empty := len(tr.pending) == 0 - tr.mu.Unlock() - if !empty { - t.Fatalf("pending not cleared after Close") - } - if err := tr.Close(); err != nil { - t.Fatalf("second Close: %v", err) - } -} diff --git a/internal/capability/mcp/stdio.go b/internal/capability/mcp/stdio.go deleted file mode 100644 index f79163b..0000000 --- a/internal/capability/mcp/stdio.go +++ /dev/null @@ -1,225 +0,0 @@ -package mcp - -import ( - "bufio" - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "log/slog" - "os" - "os/exec" - "strconv" - "strings" - "sync" -) - -// stdioTransport runs an MCP server as a child process and speaks -// newline-delimited JSON-RPC 2.0 to its stdin/stdout. -type stdioTransport struct { - cfg MCPConfig - log *slog.Logger - line *lineRPC - cmd *exec.Cmd - stdin io.WriteCloser - closeOnce sync.Once - closed bool - 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. -func newStdioTransport(cfg MCPConfig, log *slog.Logger) *stdioTransport { - return &stdioTransport{cfg: cfg, log: log} -} - -// Start spawns the subprocess and wires up the request/response streams. -func (s *stdioTransport) Start(ctx context.Context) error { - if err := ctx.Err(); err != nil { - return fmt.Errorf("mcp stdio start: %w", err) - } - if s.cfg.Command == "" { - return fmt.Errorf("mcp stdio: %w", ErrNoTransport) - } - cmd := exec.Command(s.cfg.Command, s.cfg.Args...) - cmd.Env = envWith(os.Environ(), s.cfg.Env) - stdin, err := cmd.StdinPipe() - if err != nil { - return fmt.Errorf("mcp stdio stdin: %w", err) - } - stdout, err := cmd.StdoutPipe() - if err != nil { - return fmt.Errorf("mcp stdio stdout: %w", err) - } - stderr, err := cmd.StderrPipe() - if err != nil { - return fmt.Errorf("mcp stdio stderr: %w", err) - } - if err := cmd.Start(); err != nil { - return fmt.Errorf("mcp stdio start %q: %w", s.cfg.Command, err) - } - s.cmd = cmd - s.stdin = stdin - s.line = &lineRPC{in: stdin, out: bufio.NewReader(stdout), log: s.log, onNotification: s.onNotification} - go s.drainStderr(stderr) - return nil -} - -// Call sends a request over stdio and decodes the response. -func (s *stdioTransport) Call(ctx context.Context, method string, params any, result any) error { - if s.line == nil { - return fmt.Errorf("mcp stdio %s: %w", method, ErrNotConnected) - } - return s.line.call(ctx, method, params, result) -} - -// Close terminates the child process and closes its streams. -func (s *stdioTransport) Close() error { - s.closeOnce.Do(func() { - s.closed = true - if s.stdin != nil { - _ = s.stdin.Close() - } - if s.cmd != nil && s.cmd.Process != nil { - _ = s.cmd.Process.Kill() - } - if s.cmd != nil { - _ = s.cmd.Wait() - } - }) - return nil -} - -// drainStderr reads and logs the child's standard error at debug level so server -// diagnostics are visible without blocking the process. -func (s *stdioTransport) drainStderr(r io.Reader) { - buf := make([]byte, 4096) - for { - n, err := r.Read(buf) - if n > 0 { - s.log.Debug("mcp server stderr", "name", s.cfg.Name, "msg", strings.TrimSpace(string(buf[:n]))) - } - if err != nil { - return - } - } -} - -// lineRPC carries newline-delimited JSON-RPC 2.0 over a stream pair. Calls are -// serialized so at most one request is in flight at a time, which keeps id -// matching simple and deterministic. -type lineRPC struct { - in io.Writer - out *bufio.Reader - log *slog.Logger - mu sync.Mutex - nextID int - onNotification func(method string) -} - -// call writes a request and reads until the matching response arrives. -func (l *lineRPC) call(ctx context.Context, method string, params any, result any) error { - l.mu.Lock() - defer l.mu.Unlock() - if err := ctx.Err(); err != nil { - return fmt.Errorf("mcp %s: %w", method, err) - } - l.nextID++ - id := strconv.Itoa(l.nextID) - req, err := encodeRequest(id, method, params) - if err != nil { - return fmt.Errorf("mcp %s: %w", method, err) - } - req = append(req, '\n') - if _, err := l.in.Write(req); err != nil { - return fmt.Errorf("mcp %s write: %w", method, err) - } - return l.readResponse(id, method, result) -} - -// readResponse reads lines until a response matching id arrives, skipping blank -// lines, malformed frames and notifications (no id). -func (l *lineRPC) readResponse(id, method string, result any) error { - for { - line, err := l.out.ReadBytes('\n') - if err != nil { - return fmt.Errorf("mcp %s read: %w", method, err) - } - line = bytes.TrimSpace(line) - if len(line) == 0 { - continue - } - var resp rpcResponse - if err := json.Unmarshal(line, &resp); err != nil { - l.log.Debug("mcp ignoring malformed frame", "line", string(line)) - 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) { - continue - } - if resp.Error != nil { - return fmt.Errorf("mcp %s: %w: %w", method, ErrRPC, resp.Error) - } - if result != nil { - if err := json.Unmarshal(resp.Result, result); err != nil { - return fmt.Errorf("mcp %s decode: %w", method, err) - } - } - return nil - } -} - -// notify writes a fire-and-forget JSON-RPC notification (no id). -func (l *lineRPC) notify(ctx context.Context, method string, params any) error { - l.mu.Lock() - defer l.mu.Unlock() - req, err := encodeRequest(nil, method, params) - if err != nil { - return fmt.Errorf("mcp notify %s: %w", method, err) - } - req = append(req, '\n') - if _, err := l.in.Write(req); err != nil { - return fmt.Errorf("mcp notify %s write: %w", method, err) - } - return nil -} - -// envWith returns base with extra merged over it (an extra key overrides a -// same-named base entry). A nil/empty extra returns base unchanged. -func envWith(base []string, extra map[string]string) []string { - if len(extra) == 0 { - return base - } - out := make([]string, 0, len(base)+len(extra)) - for _, kv := range base { - if k, _, ok := strings.Cut(kv, "="); ok { - if _, dup := extra[k]; dup { - continue - } - } - out = append(out, kv) - } - for k, v := range extra { - out = append(out, k+"="+v) - } - return out -} diff --git a/internal/capability/mcp/stdio_test.go b/internal/capability/mcp/stdio_test.go deleted file mode 100644 index 9ed3ef4..0000000 --- a/internal/capability/mcp/stdio_test.go +++ /dev/null @@ -1,199 +0,0 @@ -package mcp - -import ( - "bufio" - "bytes" - "context" - "errors" - "io" - "strings" - "testing" -) - -// newLineRPC builds a lineRPC reading the supplied pre-encoded responses. -func newLineRPC(responses string) *lineRPC { - return &lineRPC{ - in: &bytes.Buffer{}, - out: bufio.NewReader(bytes.NewReader([]byte(responses))), - log: defaultLogger(), - } -} - -func TestLineRPC_Call_Success(t *testing.T) { - t.Parallel() - out := `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n" - l := newLineRPC(out) - var res struct { - OK bool `json:"ok"` - } - if err := l.call(context.Background(), "tools/call", map[string]any{}, &res); err != nil { - t.Fatalf("call: %v", err) - } - if !res.OK { - t.Fatal("res.OK = false, want true") - } -} - -func TestLineRPC_Call_SkipsNotificationThenMatches(t *testing.T) { - t.Parallel() - out := `{"jsonrpc":"2.0","method":"notifications/initialized"}` + "\n" + - `{"jsonrpc":"2.0","id":"1","result":{}}` + "\n" - l := newLineRPC(out) - var res map[string]any - if err := l.call(context.Background(), "initialize", map[string]any{}, &res); err != nil { - t.Fatalf("call: %v", err) - } -} - -func TestLineRPC_Call_SkipsMismatchedID(t *testing.T) { - t.Parallel() - out := `{"jsonrpc":"2.0","id":"2","result":{}}` + "\n" + - `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n" - l := newLineRPC(out) - var res struct { - OK bool `json:"ok"` - } - if err := l.call(context.Background(), "tools/call", nil, &res); err != nil { - t.Fatalf("call: %v", err) - } - if !res.OK { - t.Fatal("res.OK = false, want true") - } -} - -func TestLineRPC_Call_RPCError(t *testing.T) { - t.Parallel() - out := `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"boom"}}` + "\n" - l := newLineRPC(out) - err := l.call(context.Background(), "tools/call", nil, nil) - if !errors.Is(err, ErrRPC) { - t.Fatalf("err = %v, want ErrRPC", err) - } - var rpcErr *rpcError - if !errors.As(err, &rpcErr) { - t.Fatalf("err = %v, want *rpcError in chain", err) - } - if rpcErr.Message != "boom" { - t.Fatalf("rpcErr.Message = %q, want boom", rpcErr.Message) - } -} - -func TestLineRPC_Call_LastResponseDecodeError(t *testing.T) { - t.Parallel() - // A matching id whose result does not fit the target type must surface a - // decode error. - out := `{"jsonrpc":"2.0","id":"1","result":{"ok":"not-a-bool"}}` + "\n" - l := newLineRPC(out) - var res struct { - OK bool `json:"ok"` - } - err := l.call(context.Background(), "tools/list", nil, &res) - if err == nil { - t.Fatal("call with mismatched result succeeded, want decode error") - } -} - -func TestLineRPC_Call_EOF(t *testing.T) { - t.Parallel() - l := newLineRPC("") - err := l.call(context.Background(), "tools/list", nil, nil) - if !errors.Is(err, io.EOF) { - t.Fatalf("err = %v, want io.EOF", err) - } -} - -func TestLineRPC_Call_MalformedFrameIgnored(t *testing.T) { - t.Parallel() - out := `not json` + "\n" + `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n" - l := newLineRPC(out) - var res struct { - OK bool `json:"ok"` - } - if err := l.call(context.Background(), "tools/call", nil, &res); err != nil { - t.Fatalf("call: %v", err) - } - if !res.OK { - t.Fatal("res.OK = false, want true") - } -} - -func TestLineRPC_WriteError(t *testing.T) { - t.Parallel() - l := &lineRPC{in: failWriter{}, out: bufio.NewReader(bytes.NewReader(nil)), log: defaultLogger()} - err := l.call(context.Background(), "tools/list", nil, nil) - if err == nil { - t.Fatal("call with failing writer succeeded, want error") - } - if !strings.Contains(err.Error(), "write") { - t.Fatalf("err = %v, want write error", err) - } -} - -func TestLineRPC_Notify_WritesNoID(t *testing.T) { - t.Parallel() - var buf bytes.Buffer - l := &lineRPC{in: &buf, out: bufio.NewReader(bytes.NewReader(nil)), log: defaultLogger()} - if err := l.notify(context.Background(), "notifications/initialized", map[string]any{}); err != nil { - t.Fatalf("notify: %v", err) - } - body := buf.String() - if !strings.Contains(body, "notifications/initialized") { - t.Fatalf("body = %q, want notifications/initialized", body) - } - if strings.Contains(body, `"id"`) { - t.Fatalf("body = %q, must not contain id (notification)", body) - } -} - -func TestLineRPC_ContextCanceled(t *testing.T) { - t.Parallel() - ctx, cancel := context.WithCancel(context.Background()) - cancel() - l := newLineRPC(`{"jsonrpc":"2.0","id":"1","result":{}}` + "\n") - err := l.call(ctx, "initialize", nil, nil) - if !errors.Is(err, context.Canceled) { - t.Fatalf("err = %v, want context.Canceled", err) - } -} - -func TestStdioTransport_Start_NoCommand(t *testing.T) { - t.Parallel() - s := newStdioTransport(MCPConfig{Name: "s"}, defaultLogger()) - err := s.Start(context.Background()) - if !errors.Is(err, ErrNoTransport) { - t.Fatalf("Start error = %v, want ErrNoTransport", err) - } -} - -func TestStdioTransport_Start_InvalidCommand(t *testing.T) { - t.Parallel() - s := newStdioTransport(MCPConfig{Name: "s", Command: "cyrene-no-such-binary-xyz"}, defaultLogger()) - if err := s.Start(context.Background()); err == nil { - t.Fatal("Start with bogus command succeeded, want error") - } -} - -func TestStdioTransport_Call_NotStarted(t *testing.T) { - t.Parallel() - s := newStdioTransport(MCPConfig{Name: "s", Command: "x"}, defaultLogger()) - err := s.Call(context.Background(), "tools/list", nil, nil) - if !errors.Is(err, ErrNotConnected) { - t.Fatalf("Call error = %v, want ErrNotConnected", err) - } -} - -func TestStdioTransport_Close_NoCmd(t *testing.T) { - t.Parallel() - s := newStdioTransport(MCPConfig{Name: "s"}, defaultLogger()) - if err := s.Close(); err != nil { - t.Fatalf("Close: %v", err) - } -} - -// failWriter returns an error on every write. -type failWriter struct{} - -// Write always fails. -func (failWriter) Write(_ []byte) (int, error) { - return 0, errors.New("write boom") -} diff --git a/internal/capability/mcp/tool.go b/internal/capability/mcp/tool.go index 49eb945..272cbc5 100644 --- a/internal/capability/mcp/tool.go +++ b/internal/capability/mcp/tool.go @@ -2,47 +2,91 @@ package mcp import ( "context" + "encoding/json" + "fmt" "strings" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" ) -// mcpTool adapts an MCP tool to the pkg/plugins/sdk Tool interface so it can be -// registered and invoked like any built-in tool. It is immutable and safe for -// concurrent use. +// ToolInfo 是 MCP 服务器暴露的一把工具。 +type ToolInfo struct { + // Name 是调用该工具用的标识。 + Name string + // Description 说明工具做什么。 + Description string + // InputSchema 是工具的 JSON Schema，通常是 + // {"type":"object","properties":{...}}。 + InputSchema map[string]any +} + +// ContentBlock 是工具调用结果里的一个内容块（MCP content block）。 +type ContentBlock struct { + // Type 是块类型，最常见的是 "text"；非文本块只保留类型标记。 + Type string `json:"type"` + // Text 是文本载荷。 + Text string `json:"text"` +} + +// ToolResult 是一次 tools/call 的结构化结果。 +type ToolResult struct { + // Content 是服务器返回的内容块。 + Content []ContentBlock + // Structured 是服务器返回的结构化结果（structuredContent），可能为 nil。 + Structured any + // IsError 表示服务器把这次调用标成了错误。 + IsError bool +} + +// toolInfoFrom 把 SDK 的工具描述转成本包的 ToolInfo。 +func toolInfoFrom(t *mcpsdk.Tool) ToolInfo { + return ToolInfo{ + Name: t.Name, + Description: t.Description, + InputSchema: schemaMap(t.InputSchema), + } +} + +// mcpTool 把一把 MCP 工具适配成 pkg/plugins/sdk 的 Tool 接口， +// 使其能被注册进 registry、被 agent 当作普通工具调用。字段不可变，可并发使用。 type mcpTool struct { client *Client info ToolInfo } -// Name returns the MCP tool identifier. -func (t *mcpTool) Name() string { - return t.info.Name -} +// Name 返回 MCP 工具标识。 +func (t *mcpTool) Name() string { return t.info.Name } -// Description returns the MCP tool description. -func (t *mcpTool) Description() string { - return t.info.Description -} +// Description 返回 MCP 工具描述。 +func (t *mcpTool) Description() string { return t.info.Description } -// Parameters returns the MCP tool input schema, defaulting to an empty object -// schema when the server omits one. +// Parameters 返回 MCP input schema；服务器未给 schema 时返回无约束空对象 schema。 func (t *mcpTool) Parameters() (map[string]any, error) { if t.info.InputSchema == nil { - return map[string]any{"type": "object", "properties": map[string]any{}}, nil + return emptyObjectSchema(), nil } return t.info.InputSchema, nil } -// Execute invokes the MCP tool and returns its concatenated text content. +// Execute 调用 MCP 工具并返回其结果文本。 +// 参数先经 input schema 缓存校验 + 补默认值（design/09 §1.1）， +// 文本块全为空时退回结构化结果的 JSON 文本，保证模型总能看见东西。 func (t *mcpTool) Execute(ctx context.Context, args any) (any, error) { - res, err := t.client.CallTool(ctx, t.info.Name, args) + params, err := normalizeArgs(args) if err != nil { return nil, err } - return contentText(res.Content), nil + res, err := t.client.callTool(ctx, t.info.Name, params, t.info.InputSchema) + if err != nil { + return nil, err + } + if text := contentText(res.Content); text != "" { + return text, nil + } + return structuredText(res.Structured) } -// contentText concatenates the text blocks of a tool result into a single -// string, ignoring non-text blocks. +// contentText 把工具结果里的文本块拼成一个字符串，忽略非文本块。 func contentText(content []ContentBlock) string { var b strings.Builder for _, blk := range content { @@ -53,11 +97,47 @@ func contentText(content []ContentBlock) string { return b.String() } -// toContentBlocks converts wire content blocks to the public content block type. -func toContentBlocks(wire []contentBlock) []ContentBlock { - out := make([]ContentBlock, 0, len(wire)) - for _, c := range wire { - out = append(out, ContentBlock{Type: c.Type, Text: c.Text}) +// structuredText 把结构化结果序列化成 JSON 文本；无结构化结果时返回空串。 +func structuredText(structured any) (string, error) { + if structured == nil { + return "", nil + } + data, err := json.Marshal(structured) + if err != nil { + return "", fmt.Errorf("marshal structured tool result: %w", err) + } + return string(data), nil +} + +// toContentBlocks 把 SDK 内容块转成本包结构。非文本块只保留类型标记， +// 不把 base64 大块塞进模型上下文。 +func toContentBlocks(content []mcpsdk.Content) []ContentBlock { + out := make([]ContentBlock, 0, len(content)) + for _, blk := range content { + if text, ok := blk.(*mcpsdk.TextContent); ok { + out = append(out, ContentBlock{Type: "text", Text: text.Text}) + continue + } + out = append(out, ContentBlock{Type: contentTypeOf(blk)}) } return out } + +// contentTypeOf 返回内容块的线上类型名（如 image / audio / resource）； +// 未知类型统一记作 non-text。 +func contentTypeOf(blk mcpsdk.Content) string { + switch blk.(type) { + case *mcpsdk.ImageContent: + return "image" + case *mcpsdk.AudioContent: + return "audio" + case *mcpsdk.ResourceLink, *mcpsdk.EmbeddedResource: + return "resource" + case *mcpsdk.ToolUseContent: + return "tool_use" + case *mcpsdk.ToolResultContent: + return "tool_result" + default: + return "non-text" + } +} diff --git a/internal/capability/mcp/transport.go b/internal/capability/mcp/transport.go new file mode 100644 index 0000000..19e14b1 --- /dev/null +++ b/internal/capability/mcp/transport.go @@ -0,0 +1,98 @@ +package mcp + +import ( + "fmt" + "net/http" + "os" + "os/exec" + "strings" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// newTransport 按 cfg 选传输。显式 TransportType 优先（stdio / sse / http，http 与 +// streamable 等价，对齐 design/09 §2.1 的枚举）；为空则沿用历史推断： +// 有 Command 走 stdio，有 URL 走 SSE（既有 URL 配置多为 /sse 端点）。 +// APIKey 通过自定义 RoundTripper 以 Authorization: Bearer 发送。 +func newTransport(cfg MCPConfig) (mcpsdk.Transport, error) { + switch cfg.TransportType { + case "http", "streamable": + return httpURLTransport(cfg, true) + case "sse": + return httpURLTransport(cfg, false) + case "stdio": + return commandTransport(cfg) + } + switch { + case cfg.Command != "": + return commandTransport(cfg) + case cfg.URL != "": + return httpURLTransport(cfg, false) + default: + return nil, fmt.Errorf("mcp: %w", ErrNoTransport) + } +} + +// httpURLTransport 构建 URL 型传输：streamable=true 走 Streamable HTTP（2025-11-25 +// 规范首选），否则走旧式 SSE。 +func httpURLTransport(cfg MCPConfig, streamable bool) (mcpsdk.Transport, error) { + if cfg.URL == "" { + return nil, fmt.Errorf("mcp transport %s: %w", cfg.TransportType, ErrNoTransport) + } + client := httpClientFor(cfg) + if streamable { + return &mcpsdk.StreamableClientTransport{Endpoint: cfg.URL, HTTPClient: client}, nil + } + return &mcpsdk.SSEClientTransport{Endpoint: cfg.URL, HTTPClient: client}, nil +} + +// commandTransport 拉起 stdio 子进程，按换行分隔的 JSON-RPC 通信。 +func commandTransport(cfg MCPConfig) (mcpsdk.Transport, error) { + if cfg.Command == "" { + return nil, fmt.Errorf("mcp transport stdio: %w", ErrNoTransport) + } + cmd := exec.Command(cfg.Command, cfg.Args...) + cmd.Env = envWith(os.Environ(), cfg.Env) + return &mcpsdk.CommandTransport{Command: cmd}, nil +} + +// httpClientFor 返回带鉴权头的 HTTP 客户端；没有 APIKey 时返回 nil（用 SDK 默认）。 +func httpClientFor(cfg MCPConfig) *http.Client { + if cfg.APIKey == "" { + return nil + } + return &http.Client{Transport: &bearerTransport{base: http.DefaultTransport, token: cfg.APIKey}} +} + +// bearerTransport 给每个请求加 Authorization: Bearer。凭据不写日志、不进错误消息。 +type bearerTransport struct { + base http.RoundTripper + token string +} + +// RoundTrip 实现 http.RoundTripper。 +func (t *bearerTransport) RoundTrip(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + clone.Header.Set("Authorization", "Bearer "+t.token) + return t.base.RoundTrip(clone) +} + +// envWith 返回 base 叠加 extra 后的环境变量（同名以 extra 为准）；extra 为空时原样返回。 +func envWith(base []string, extra map[string]string) []string { + if len(extra) == 0 { + return base + } + out := make([]string, 0, len(base)+len(extra)) + for _, kv := range base { + if k, _, ok := strings.Cut(kv, "="); ok { + if _, dup := extra[k]; dup { + continue + } + } + out = append(out, kv) + } + for k, v := range extra { + out = append(out, k+"="+v) + } + return out +} diff --git a/internal/capability/mcp/transport_test.go b/internal/capability/mcp/transport_test.go new file mode 100644 index 0000000..d386b0a --- /dev/null +++ b/internal/capability/mcp/transport_test.go @@ -0,0 +1,123 @@ +package mcp + +import ( + "net/http" + "net/http/httptest" + "testing" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// TestNewTransport_SelectsByType 覆盖显式传输类型与历史推断。 +func TestNewTransport_SelectsByType(t *testing.T) { + t.Parallel() + tests := []struct { + name string + cfg MCPConfig + want any + }{ + {name: "http", cfg: MCPConfig{URL: "http://x", TransportType: "http"}, want: &mcpsdk.StreamableClientTransport{}}, + {name: "streamable 别名", cfg: MCPConfig{URL: "http://x", TransportType: "streamable"}, want: &mcpsdk.StreamableClientTransport{}}, + {name: "sse", cfg: MCPConfig{URL: "http://x", TransportType: "sse"}, want: &mcpsdk.SSEClientTransport{}}, + {name: "stdio", cfg: MCPConfig{Command: "x", TransportType: "stdio"}, want: &mcpsdk.CommandTransport{}}, + {name: "推断 stdio", cfg: MCPConfig{Command: "x"}, want: &mcpsdk.CommandTransport{}}, + {name: "推断 sse", cfg: MCPConfig{URL: "http://x"}, want: &mcpsdk.SSEClientTransport{}}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + tr, err := newTransport(tt.cfg) + if err != nil { + t.Fatalf("newTransport: %v", err) + } + if got, want := typeName(tr), typeName(tt.want); got != want { + t.Fatalf("传输类型 = %s, 期望 %s", got, want) + } + }) + } +} + +// TestNewTransport_NoTransport 两个字段都空时必须报 ErrNoTransport。 +func TestNewTransport_NoTransport(t *testing.T) { + t.Parallel() + if _, err := newTransport(MCPConfig{Name: "s"}); err == nil { + t.Fatal("没有 Command/URL 却选出了传输") + } + // 显式类型但缺必要字段同样要报错（不静默构造坏传输）。 + if _, err := newTransport(MCPConfig{TransportType: "http"}); err == nil { + t.Fatal("TransportType=http 但 URL 为空却没有报错") + } + if _, err := newTransport(MCPConfig{TransportType: "stdio"}); err == nil { + t.Fatal("TransportType=stdio 但 Command 为空却没有报错") + } +} + +// TestBearerTransport_SetsAuthorization 凭据以 Authorization: Bearer 下发， +// 并且不污染原始请求。 +func TestBearerTransport_SetsAuthorization(t *testing.T) { + t.Parallel() + got := make(chan string, 1) + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got <- r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + client := httpClientFor(MCPConfig{APIKey: "sk-1234"}) + if client == nil { + t.Fatal("配了 APIKey 却没有返回自定义 http client") + } + req, err := http.NewRequest(http.MethodGet, ts.URL, nil) + if err != nil { + t.Fatalf("NewRequest: %v", err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + _ = resp.Body.Close() + if auth := <-got; auth != "Bearer sk-1234" { + t.Fatalf("Authorization = %q, 期望 Bearer sk-1234", auth) + } + if req.Header.Get("Authorization") != "" { + t.Fatal("原始请求被就地改写了") + } + if httpClientFor(MCPConfig{}) != nil { + t.Fatal("无 APIKey 时应返回 nil（用 SDK 默认客户端）") + } +} + +// TestEnvWith_OverridesAndPreserves stdio 子进程环境变量叠加语义。 +func TestEnvWith_OverridesAndPreserves(t *testing.T) { + t.Parallel() + base := []string{"PATH=/usr/bin", "HOME=/root"} + got := envWith(base, map[string]string{"HOME": "/tmp", "TOKEN": "abc"}) + if len(got) != 3 { + t.Fatalf("len = %d, 期望 3（%v）", len(got), got) + } + seen := map[string]bool{} + for _, kv := range got { + seen[kv] = true + } + if !seen["PATH=/usr/bin"] || !seen["HOME=/tmp"] || !seen["TOKEN=abc"] { + t.Fatalf("合并后的环境 = %v", got) + } + if len(envWith(base, nil)) != len(base) { + t.Fatal("空 extra 改变了基础环境") + } +} + +// typeName 返回传输的类别名，避免在断言里依赖具体指针比较。 +func typeName(v any) string { + switch v.(type) { + case *mcpsdk.StreamableClientTransport: + return "streamable" + case *mcpsdk.SSEClientTransport: + return "sse" + case *mcpsdk.CommandTransport: + return "stdio" + default: + return "unknown" + } +}