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 }