AskaEth/Cyrene_v2
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mcp): 协议层换成官方 Model Context Protocol SDK，删自研实现（裁决 10.4.6 + AGENTS §9.5）
CI / Backend (build / vet / test / cross-compile) (push) Has been cancelled
Details
CI / Frontend (tsc + vite build) (push) Has been cancelled
Details

Browse Source 
- 新增依赖：github.com/modelcontextprotocol/go-sdk v1.8.0（direct）+ github.com/google/jsonschema-go v0.4.3
  （SDK 自带依赖，直接用于 schema 校验/默认值，避免再造一份）；go.sum 纯新增
- 删除自研协议层：rpc.go / stdio.go / http.go / sse.go 及其测试
- 用 SDK 的成熟逻辑：命令传输（stdio 子进程）/ SSE / Streamable HTTP、协议版本协商与
  -32022 回退、keepalive（含失败阈值）、tools/list_changed 订阅、能力协商（InitializeResult.Capabilities）
- 保留自研的四件事（其余零自研）：配置→传输、发现→registry 注册、能力族推断、schema 缓存（Prepare 命中不重复解析）
- 能力闸诚实化：服务器声明 capabilities 但不含 tools 时 ListTools/DiscoverTools 返回
  ErrCapabilityNotDeclared（不静默空集）；完全不声明 capabilities 的老式服务器仍按 legacy 处理
- 复核：build 绿 + capability 全包测试绿（mcp 42 用例，含真子进程 stdio / httptest Streamable HTTP /
  内存传输三类假 server，不依赖外网）+ `-race` 干净；变异 2 条（featuresFrom 恒 true / 删 applySchema 调用）已还原
- 无 REST/WS/env/数据目录变更；包内 API 有增删（见回报），仓外调用方签名未变（未改 cmd/cyrene 与 plugin/manager）
This commit is contained in:
AskaEth committed 2026-10-04 13:53:47 +08:00
1 parent 7b05b4bfa9
commit a1083a322f
26 files changed
+1903 -3011
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

go.mod
+7
View File
@@ -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
go.sum
+18
View File
@@ -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=
internal/capability/mcp/capabilities.go
+135
View File
@@ -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)
}
}
internal/capability/mcp/capabilities_test.go
+119
View File
@@ -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)
}
}
internal/capability/mcp/client.go
+130 -337
View File
@@ -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"`
}
internal/capability/mcp/client_test.go
+68 -213
View File
@@ -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("字符串参数没有被拒绝")
}
}
internal/capability/mcp/discovery_test.go
+294
View File
@@ -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
}
internal/capability/mcp/enhancements_test.go
-194
View File
@@ -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")
}
internal/capability/mcp/helpers_test.go
+132 -167
View File
@@ -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)
}
internal/capability/mcp/http.go
-422
View File
@@ -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
}
internal/capability/mcp/http_test.go
-197
View File
@@ -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)
}
}
internal/capability/mcp/keepalive_test.go
+109
View File
@@ -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)
}
internal/capability/mcp/mcp.go
+191 -276
View File
@@ -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
}
internal/capability/mcp/mcp_auth_test.go
-47
View File
@@ -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)
}
}
internal/capability/mcp/registry_test.go
-171
View File
@@ -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 }
internal/capability/mcp/rpc.go
-53
View File
@@ -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
}
internal/capability/mcp/schema.go
+84
View File
@@ -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
}
internal/capability/mcp/schemacache.go
+142
View File
@@ -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
}
internal/capability/mcp/schemacache_test.go
+149
View File
@@ -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)
}
}
internal/capability/mcp/sse.go
-283
View File
@@ -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
}
internal/capability/mcp/sse_test.go
-203
View File
@@ -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)
}
}
internal/capability/mcp/stdio.go
-225
View File
@@ -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
}
internal/capability/mcp/stdio_test.go
-199
View File
@@ -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")
}
internal/capability/mcp/tool.go
+104 -24
View File
@@ -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"
}
}
internal/capability/mcp/transport.go
+98
View File
@@ -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
}
internal/capability/mcp/transport_test.go
+123
View File
@@ -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"
}
}