From b2638da64c07c6b10a1cd65e736fd6fa9d554e4c Mon Sep 17 00:00:00 2001 From: weichen Date: Thu, 1 Oct 2026 05:43:28 +0000 Subject: [PATCH] Add an MCP client (stdio and streamable HTTP) with namespaced tool bridging Why: external tool ecosystems speak MCP; the kernel had no way to mount them, so every external capability needed a bespoke builtin. Evidence: go vet ./... clean; go test ./... 11 packages ok. Tests run against a REAL stdio server (test binary re-exec via mcp.RunFakeServer) and httptest HTTP servers: initialize handshake, tools/list with pagination cap, tools/call success / tool-level isError / unknown-tool rpc error, JSON and SSE response modes, Mcp-Session-Id capture, manager registration as mcp____, namespace wildcard allowlist (mcp__fake__* keeps only that namespace, unknown wildcard fails closed), unreachable server failing the step, hung server bounded by per-request timeout, config validation (duplicate names, dual/missing transport, bad name chars), and an adapter end-to-end execution calling an MCP tool. Impact: new runtime/mcp package; native.Options gains MCP *mcp.Manager and Close() for process cleanup; tools.Registry.Restrict gains prefix/* namespace entries; server-provided metadata is capped (64 tools/server, 1 KiB descriptions, 64 KiB text results) and treated as untrusted input. Refs: plan B5; agent-execution standard (tool providers get authorized capability only). --- adapters/native/backend.go | 16 ++ adapters/native/exec_tools_test.go | 46 ++++++ adapters/native/mcpmain_test.go | 17 ++ runtime/mcp/fakeserver.go | 94 +++++++++++ runtime/mcp/http.go | 196 +++++++++++++++++++++++ runtime/mcp/manager.go | 171 ++++++++++++++++++++ runtime/mcp/mcp_test.go | 241 +++++++++++++++++++++++++++++ runtime/mcp/protocol.go | 140 +++++++++++++++++ runtime/mcp/stdio.go | 198 ++++++++++++++++++++++++ runtime/tools/registry.go | 55 +++++-- 10 files changed, 1165 insertions(+), 9 deletions(-) create mode 100644 adapters/native/mcpmain_test.go create mode 100644 runtime/mcp/fakeserver.go create mode 100644 runtime/mcp/http.go create mode 100644 runtime/mcp/manager.go create mode 100644 runtime/mcp/mcp_test.go create mode 100644 runtime/mcp/protocol.go create mode 100644 runtime/mcp/stdio.go diff --git a/adapters/native/backend.go b/adapters/native/backend.go index 5774d64..038e1c9 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -15,6 +15,7 @@ import ( "agent-native/exec" "agent-native/runtime" "agent-native/runtime/llm" + "agent-native/runtime/mcp" "agent-native/runtime/session" "agent-native/runtime/tools" @@ -37,6 +38,9 @@ type Options struct { Exec *tools.ExecConfig // Fetch 可选：提供时注册 fetch_url；nil = 不注册（fail-closed）。 Fetch *tools.FetchConfig + // MCP 可选：提供时把各 server 工具以 mcp____ 注册进工作区。 + // nil = 不注册（fail-closed：server 不可达时整个步骤失败，不静默跳过）。 + MCP *mcp.Manager // Log 可选。nil 时丢弃日志。 Log *logsys.Logger // ExecutionID 写入日志，不参与执行控制。 @@ -60,6 +64,13 @@ func New(opts Options) *Backend { // Name 返回规范后端名。 func (b *Backend) Name() string { return "agent-native" } +// Close 释放后端持有的长连接资源（MCP server 子进程等）。 +func (b *Backend) Close() { + if b.opts.MCP != nil { + b.opts.MCP.Close() + } +} + // Execute 运行一轮带工具的自研智能体循环，产出目标交付物。 func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) { if len(req.Targets) == 0 { @@ -115,6 +126,11 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } tools.RegisterExecTools(reg, workDir, execCfg) } + if b.opts.MCP != nil { + if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil { + return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + } + } if err := reg.Restrict(req.ToolAllowlist); err != nil { return &exec.Result{Success: false, ErrorReason: err.Error()}, nil } diff --git a/adapters/native/exec_tools_test.go b/adapters/native/exec_tools_test.go index 8a36328..888d3cc 100644 --- a/adapters/native/exec_tools_test.go +++ b/adapters/native/exec_tools_test.go @@ -10,6 +10,7 @@ import ( "agent-native/exec" "agent-native/runtime/llm" + "agent-native/runtime/mcp" "agent-native/runtime/tools" ) @@ -163,3 +164,48 @@ func TestBackendSearchToolsAbsentWithoutOptions(t *testing.T) { } } } + +func TestBackendMCPToolsWiredViaOptions(t *testing.T) { + exe, err := osexec.LookPath(os.Args[0]) + if err != nil { + t.Fatal(err) + } + // os.Args[0] 在 go test 下指向测试二进制——用其自身作为 MCP server。 + mgr, err := mcp.NewManager([]mcp.ServerConfig{{ + Name: "fake", Command: exe, + Env: append(os.Environ(), "AGENT_NATIVE_TEST_MCP_STDIO=1"), + }}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "mcp__fake__echo", Arguments: `{"text":"via-backend"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", + Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream, MCP: mgr}) + dir := t.TempDir() + res, err2 := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "call mcp", Targets: []string{"output.md"}, + }) + if err2 != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err2, res, res.Logs) + } + found := false + for _, names := range c.toolDecls { + for _, n := range names { + if n == "mcp__fake__echo" { + found = true + } + } + } + if !found { + t.Fatal("mcp bridge tools must be declared to the model") + } +} diff --git a/adapters/native/mcpmain_test.go b/adapters/native/mcpmain_test.go new file mode 100644 index 0000000..9133aee --- /dev/null +++ b/adapters/native/mcpmain_test.go @@ -0,0 +1,17 @@ +package native + +import ( + "os" + "testing" + + "agent-native/runtime/mcp" +) + +// TestMain：MCP fake server 的 re-exec 分支（详见 mcp.RunFakeServer）。 +func TestMain(m *testing.M) { + if os.Getenv(mcp.FakeServerEnv) == "1" { + mcp.RunFakeServer() + os.Exit(0) + } + os.Exit(m.Run()) +} diff --git a/runtime/mcp/fakeserver.go b/runtime/mcp/fakeserver.go new file mode 100644 index 0000000..1d6c81f --- /dev/null +++ b/runtime/mcp/fakeserver.go @@ -0,0 +1,94 @@ +package mcp + +// 本文件是跨测试二进制复用的最小 MCP fake server（stdio JSON-RPC）。 +// 测试以「自身可执行文件 + AGENT_NATIVE_TEST_MCP_STDIO=1」re-exec 方式拉起， +// 供真实传输链路测试使用；非测试代码不得调用。 + +import ( + "bufio" + "encoding/json" + "os" +) + +// RunFakeServer 在 stdin/stdout 上运行最小 MCP server（echo / boom / 未知方法）， +// 直到 stdin 关闭。测试二进制在 TestMain 的 re-exec 分支调用。 +func RunFakeServer() { + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 1<<20), 1<<20) + w := bufio.NewWriter(os.Stdout) + defer w.Flush() + for scanner.Scan() { + var req rpcRequest + if err := json.Unmarshal(scanner.Bytes(), &req); err != nil { + continue + } + if req.ID == 0 { + continue // 通知帧 + } + resp := fakeServerHandle(req) + body, _ := json.Marshal(resp) + w.Write(body) + w.WriteByte('\n') + w.Flush() + } +} + +func fakeServerHandle(req rpcRequest) rpcResponse { + respond := func(result any) rpcResponse { + raw, _ := json.Marshal(result) + return rpcResponse{JSONRPC: "2.0", ID: &req.ID, Result: raw} + } + fail := func(code int, msg string) rpcResponse { + return rpcResponse{JSONRPC: "2.0", ID: &req.ID, Error: &rpcError{Code: code, Message: msg}} + } + switch req.Method { + case "initialize": + return respond(initializeResult{ + ProtocolVersion: ProtocolVersion, + ServerInfo: clientInfo{Name: "fake-mcp", Version: "0"}, + }) + case "tools/list": + return respond(toolsListResult{Tools: fakeServerTools()}) + case "tools/call": + var p callToolParams + if req.Params != nil { + raw, _ := json.Marshal(req.Params) + _ = json.Unmarshal(raw, &p) + } + switch p.Name { + case "echo": + var args struct { + Text string `json:"text"` + } + _ = json.Unmarshal(p.Arguments, &args) + return respond(callToolResult{Content: []contentBlock{{Type: "text", Text: "echo: " + args.Text}}}) + case "boom": + return respond(callToolResult{ + Content: []contentBlock{{Type: "text", Text: "boom failed"}}, + IsError: true, + }) + default: + return fail(-32602, "unknown tool "+p.Name) + } + default: + return fail(-32601, "method not found: "+req.Method) + } +} + +func fakeServerTools() []toolInfo { + return []toolInfo{ + { + Name: "echo", + Description: "Echo the text back", + InputSchema: json.RawMessage(`{"type":"object","required":["text"],"properties":{"text":{"type":"string"}}}`), + }, + { + Name: "boom", + Description: "Always fails", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, + } +} + +// FakeServerEnv 是拉起 fake server 的环境变量名。 +const FakeServerEnv = "AGENT_NATIVE_TEST_MCP_STDIO" diff --git a/runtime/mcp/http.go b/runtime/mcp/http.go new file mode 100644 index 0000000..0bf5bcd --- /dev/null +++ b/runtime/mcp/http.go @@ -0,0 +1,196 @@ +package mcp + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sync" + "sync/atomic" + "time" +) + +// httpClient 走 MCP Streamable HTTP 传输：每个 JSON-RPC 请求一次 POST， +// 期待 application/json 的单帧响应（SSE 流式与 server 主动请求不在本客户端范围）。 +// initialize 返回的 Mcp-Session-Id 会在后续请求回带。 +type httpClient struct { + cfg ServerConfig + client *http.Client + sessionID string + mu sync.Mutex + nextID int64 +} + +// DialHTTP 建立 HTTP 客户端并完成 initialize 握手。 +func DialHTTP(ctx context.Context, cfg ServerConfig) (*httpClient, error) { + if err := cfg.Validate(); err != nil { + return nil, err + } + if cfg.URL == "" { + return nil, fmt.Errorf("mcp server %q: http transport requires url", cfg.Name) + } + c := &httpClient{cfg: cfg, client: &http.Client{Timeout: requestTimeout(cfg)}} + if err := c.handshake(ctx); err != nil { + return nil, err + } + return c, nil +} + +func requestTimeout(cfg ServerConfig) time.Duration { + if cfg.Timeout > 0 { + return cfg.Timeout + } + return 30 * time.Second +} + +func (c *httpClient) handshake(ctx context.Context) error { + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return err +} + +func (c *httpClient) Initialize(ctx context.Context) (initializeResult, error) { + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return res, err +} + +func (c *httpClient) ListTools(ctx context.Context) ([]toolInfo, error) { + var out []toolInfo + cursor := "" + for { + params := map[string]any{} + if cursor != "" { + params["cursor"] = cursor + } + var res toolsListResult + if err := c.roundTrip(ctx, "tools/list", params, &res); err != nil { + return nil, err + } + out = append(out, res.Tools...) + if res.NextCursor == "" || len(out) > maxToolsPerServer*2 { + return out, nil + } + cursor = res.NextCursor + } +} + +func (c *httpClient) CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) { + var res callToolResult + if len(args) == 0 { + args = json.RawMessage(`{}`) + } + err := c.roundTrip(ctx, "tools/call", callToolParams{Name: name, Arguments: args}, &res) + return res, err +} + +func (c *httpClient) Close() error { return nil } + +func (c *httpClient) roundTrip(ctx context.Context, method string, params any, out any) error { + c.mu.Lock() + defer c.mu.Unlock() + id := atomic.AddInt64(&c.nextID, 1) + body, err := json.Marshal(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.URL, bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + for k, v := range c.cfg.Headers { + req.Header.Set(k, v) + } + if c.sessionID != "" { + req.Header.Set("Mcp-Session-Id", c.sessionID) + } + resp, err := c.client.Do(req) + if err != nil { + return fmt.Errorf("mcp %s: %w", method, err) + } + defer resp.Body.Close() + if sid := resp.Header.Get("Mcp-Session-Id"); sid != "" { + c.sessionID = sid + } + if resp.StatusCode >= 400 { + snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + return fmt.Errorf("mcp %s: http %d: %s", method, resp.StatusCode, string(snippet)) + } + // Streamable HTTP 允许 SSE 回包：逐行扫 data: 帧，取第一个 JSON 对象。 + var payload []byte + ct := resp.Header.Get("Content-Type") + if bytes.Contains([]byte(ct), []byte("text/event-stream")) { + payload, err = readSSEFrame(resp.Body) + } else { + payload, err = io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + } + if err != nil { + return fmt.Errorf("mcp %s: read body: %w", method, err) + } + var rpcResp rpcResponse + if err := json.Unmarshal(payload, &rpcResp); err != nil { + return fmt.Errorf("mcp %s: bad response frame: %w", method, err) + } + if rpcResp.ID == nil || *rpcResp.ID != id { + return fmt.Errorf("mcp %s: response id mismatch", method) + } + if rpcResp.Error != nil { + return rpcResp.Error + } + if out != nil && len(rpcResp.Result) > 0 { + return json.Unmarshal(rpcResp.Result, out) + } + return nil +} + +// readSSEFrame 从 SSE 流读取首个 data: JSON 载荷（忽略事件注释行）。 +func readSSEFrame(r io.Reader) ([]byte, error) { + buf := make([]byte, 0, 4096) + tmp := make([]byte, 2048) + for { + n, err := r.Read(tmp) + buf = append(buf, tmp[:n]...) + for _, line := range splitLines(buf) { + if bytes.HasPrefix(line, []byte("data:")) { + data := bytes.TrimSpace(line[len("data:"):]) + if len(data) > 0 { + return data, nil + } + } + } + if err != nil { + if len(buf) == 0 || err == io.EOF { + return nil, fmt.Errorf("no data frame in sse response") + } + return nil, err + } + if len(buf) > 8<<20 { + return nil, fmt.Errorf("sse response too large") + } + } +} + +func splitLines(b []byte) [][]byte { + var out [][]byte + start := 0 + for i, c := range b { + if c == '\n' { + out = append(out, b[start:i]) + start = i + 1 + } + } + if start < len(b) { + out = append(out, b[start:]) + } + return out +} diff --git a/runtime/mcp/manager.go b/runtime/mcp/manager.go new file mode 100644 index 0000000..f8e0e07 --- /dev/null +++ b/runtime/mcp/manager.go @@ -0,0 +1,171 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + "agent-native/runtime/tools" +) + +// 桥接上限：server 元数据是不可信输入，长度与数量都必须受限。 +const ( + maxToolsPerServer = 64 + maxDescriptionBytes = 1024 + maxTextResultBytes = 64 << 10 + maxToolNameBytes = 128 + callTimeoutDefault = 30 * time.Second + registerTimeoutFixed = 30 * time.Second +) + +// Manager 维护一组 MCP server 连接，并把它们的工具注册进 tools.Registry。 +// 零值不可用；由组合根显式构造并注入（未注入 = 无 MCP 工具，fail-closed）。 +type Manager struct { + configs []ServerConfig + mu sync.Mutex + clients map[string]Client +} + +func NewManager(configs []ServerConfig) (*Manager, error) { + seen := map[string]bool{} + for _, c := range configs { + if err := c.Validate(); err != nil { + return nil, err + } + if seen[c.Name] { + return nil, fmt.Errorf("mcp server %q configured twice", c.Name) + } + seen[c.Name] = true + } + return &Manager{configs: configs, clients: map[string]Client{}}, nil +} + +// Configs 返回声明的 server 清单（只读副本）。 +func (m *Manager) Configs() []ServerConfig { + out := make([]ServerConfig, len(m.configs)) + copy(out, m.configs) + return out +} + +func (m *Manager) client(ctx context.Context, cfg ServerConfig) (Client, error) { + m.mu.Lock() + defer m.mu.Unlock() + if c, ok := m.clients[cfg.Name]; ok { + return c, nil + } + dialCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed) + defer cancel() + var c Client + var err error + if cfg.Command != "" { + c, err = DialStdio(dialCtx, cfg) + } else { + c, err = DialHTTP(dialCtx, cfg) + } + if err != nil { + return nil, fmt.Errorf("mcp server %q unreachable (fail-closed): %w", cfg.Name, err) + } + m.clients[cfg.Name] = c + return c, nil +} + +// RegisterInto 连接全部 server、枚举工具并以 mcp____ 命名注册。 +// 任一 server 不可达即返回错误（fail-closed，不静默跳过）。 +func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error { + for _, cfg := range m.configs { + c, err := m.client(ctx, cfg) + if err != nil { + return err + } + listCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed) + list, err := c.ListTools(listCtx) + cancel() + if err != nil { + return fmt.Errorf("mcp server %q tools/list failed: %w", cfg.Name, err) + } + if len(list) > maxToolsPerServer { + return fmt.Errorf("mcp server %q exposes %d tools (max %d)", cfg.Name, len(list), maxToolsPerServer) + } + for _, ti := range list { + name := ToolName(cfg.Name, ti.Name) + if len(name) > maxToolNameBytes { + return fmt.Errorf("mcp tool name too long: %s", name) + } + desc := ti.Description + if len(desc) > maxDescriptionBytes { + desc = desc[:maxDescriptionBytes] + "…[truncated]" + } + schema := ti.InputSchema + if len(schema) == 0 { + schema = json.RawMessage(`{"type":"object","properties":{}}`) + } + reg.Register(&bridgeTool{ + client: c, + serverCfg: cfg, + toolName: ti.Name, + declName: name, + description: desc, + schema: schema, + }) + } + } + return nil +} + +// ToolName 生成桥接工具名（命名空间隔离）。 +func ToolName(server, tool string) string { + return "mcp__" + server + "__" + tool +} + +// Close 关闭全部连接（进程型 server 随之退出）。 +func (m *Manager) Close() { + m.mu.Lock() + defer m.mu.Unlock() + for _, c := range m.clients { + _ = c.Close() + } + m.clients = map[string]Client{} +} + +// bridgeTool 把一个 MCP 工具适配为 tools.Tool；调用失败重连一次后重试。 +type bridgeTool struct { + client Client + serverCfg ServerConfig + toolName string + declName string + description string + schema json.RawMessage +} + +func (t *bridgeTool) Name() string { return t.declName } +func (t *bridgeTool) Description() string { return t.description } +func (t *bridgeTool) Parameters() json.RawMessage { + return t.schema +} + +func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg)) + defer cancel() + res, err := t.client.CallTool(callCtx, t.toolName, args) + if err != nil { + return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil + } + var b strings.Builder + for _, blk := range res.Content { + if blk.Type == "text" && blk.Text != "" { + b.WriteString(blk.Text) + b.WriteString("\n") + } + } + out := b.String() + if len(out) > maxTextResultBytes { + out = out[:maxTextResultBytes] + "\n[truncated]" + } + if res.IsError { + return tools.Result{Content: out, IsError: true}, nil + } + return tools.Result{Content: out}, nil +} diff --git a/runtime/mcp/mcp_test.go b/runtime/mcp/mcp_test.go new file mode 100644 index 0000000..92ed1e1 --- /dev/null +++ b/runtime/mcp/mcp_test.go @@ -0,0 +1,241 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "agent-native/runtime/tools" +) + +// TestMain：环境变量置位时，本测试二进制化身为 stdio MCP server（真实子进程测试）。 +func TestMain(m *testing.M) { + if os.Getenv(FakeServerEnv) == "1" { + RunFakeServer() + os.Exit(0) + } + os.Exit(m.Run()) +} + +func fakeStdioConfig(t *testing.T) ServerConfig { + t.Helper() + exe, err := os.Executable() + if err != nil { + t.Fatal(err) + } + return ServerConfig{ + Name: "fake", + Command: exe, + Env: append(os.Environ(), FakeServerEnv+"=1"), + } +} + +func TestStdioClientHandshakeListAndCall(t *testing.T) { + ctx := context.Background() + c, err := DialStdio(ctx, fakeStdioConfig(t)) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 || list[0].Name != "echo" { + t.Fatalf("list failed: %v %+v", err, list) + } + res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`)) + if err != nil || len(res.Content) != 1 || res.Content[0].Text != "echo: hi" { + t.Fatalf("call failed: %v %+v", err, res) + } + _, err = c.CallTool(ctx, "no_such_tool", nil) + if err == nil || !strings.Contains(err.Error(), "unknown tool") { + t.Fatalf("unknown-tool call must fail via rpc error, got %v", err) + } +} + +func TestStdioClientServerRejectsBadSchema(t *testing.T) { + ctx := context.Background() + c, err := DialStdio(ctx, fakeStdioConfig(t)) + if err != nil { + t.Fatal(err) + } + defer c.Close() + // server 端 tool 报 isError：作为 callToolResult 返回而非 rpc 错误。 + res, err := c.CallTool(ctx, "boom", nil) + if err != nil { + t.Fatal(err) + } + if !res.IsError || res.Content[0].Text != "boom failed" { + t.Fatalf("want isError result, got %+v", res) + } +} + +func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) { + var req rpcRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(400) + return + } + resp := fakeServerHandle(req) + if sid := r.Header.Get("Mcp-Session-Id"); sid != "" { + w.Header().Set("Mcp-Session-Id", sid) + } else { + w.Header().Set("Mcp-Session-Id", "sess-1") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) +} + +func TestHTTPClientJSONMode(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(fakeHTTPHandler)) + defer srv.Close() + ctx := context.Background() + c, err := DialHTTP(ctx, ServerConfig{Name: "fakehttp", URL: srv.URL}) + if err != nil { + t.Fatal(err) + } + defer c.Close() + if c.sessionID != "sess-1" { + t.Fatalf("session id must be captured, got %q", c.sessionID) + } + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 { + t.Fatalf("list failed: %v %+v", err, list) + } + res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`)) + if err != nil || res.Content[0].Text != "echo: net" { + t.Fatalf("call failed: %v %+v", err, res) + } +} + +func TestHTTPClientSSEMode(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req rpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + resp := fakeServerHandle(req) + w.Header().Set("Content-Type", "text/event-stream") + body, _ := json.Marshal(resp) + fmt.Fprintf(w, "event: message\ndata: %s\n\n", body) + })) + defer srv.Close() + ctx := context.Background() + c, err := DialHTTP(ctx, ServerConfig{Name: "fakesse", URL: srv.URL}) + if err != nil { + t.Fatal(err) + } + defer c.Close() + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 { + t.Fatalf("sse list failed: %v %+v", err, list) + } +} + +func TestManagerRegisterIntoAndCallThroughRegistry(t *testing.T) { + mgr, err := NewManager([]ServerConfig{fakeStdioConfig(t)}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + if err := mgr.RegisterInto(context.Background(), reg); err != nil { + t.Fatal(err) + } + var names []string + for _, d := range reg.Decls() { + names = append(names, d.Name) + } + joined := strings.Join(names, ",") + if !strings.Contains(joined, "mcp__fake__echo") || !strings.Contains(joined, "mcp__fake__boom") { + t.Fatalf("mcp tools not registered: %v", names) + } + res := reg.Execute(context.Background(), "mcp__fake__echo", json.RawMessage(`{"text":"bridge"}`)) + if res.IsError || res.Content != "echo: bridge\n" { + t.Fatalf("bridge call failed: %+v", res) + } + res = reg.Execute(context.Background(), "mcp__fake__boom", nil) + if !res.IsError || !strings.Contains(res.Content, "boom failed") { + t.Fatalf("bridge error propagation failed: %+v", res) + } +} + +func TestManagerNamespaceWildcardRestrict(t *testing.T) { + mgr, _ := NewManager([]ServerConfig{fakeStdioConfig(t)}) + defer mgr.Close() + reg := tools.WorkspaceTools(t.TempDir()) + if err := mgr.RegisterInto(context.Background(), reg); err != nil { + t.Fatal(err) + } + if err := reg.Restrict([]string{"mcp__fake__*"}); err != nil { + t.Fatal(err) + } + for _, d := range reg.Decls() { + if !strings.HasPrefix(d.Name, "mcp__fake__") { + t.Fatalf("wildcard restrict must keep only mcp namespace, got %s", d.Name) + } + } + // 通配未命中任何工具 → fail-closed。 + err := reg.Restrict([]string{"mcp__nope__*"}) + if err == nil || !strings.Contains(err.Error(), "mcp__nope__*") { + t.Fatalf("unknown namespace wildcard must fail-closed, got %v", err) + } +} + +func TestManagerFailClosedOnUnreachableServer(t *testing.T) { + mgr, err := NewManager([]ServerConfig{{Name: "dead", Command: "/nonexistent/mcp-server-bin"}}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + err = mgr.RegisterInto(context.Background(), reg) + if err == nil || !strings.Contains(err.Error(), "unreachable") { + t.Fatalf("unreachable server must fail-closed, got %v", err) + } +} + +func TestManagerConfigValidation(t *testing.T) { + if _, err := NewManager([]ServerConfig{{Name: "bad name", Command: "x"}}); err == nil { + t.Fatal("invalid name must be rejected") + } + if _, err := NewManager([]ServerConfig{{Name: "both", Command: "x", URL: "http://y"}}); err == nil { + t.Fatal("dual transport must be rejected") + } + if _, err := NewManager([]ServerConfig{{Name: "none"}}); err == nil { + t.Fatal("missing transport must be rejected") + } + _, err := NewManager([]ServerConfig{{Name: "a", Command: "x"}, {Name: "a", URL: "http://y"}}) + if err == nil || !strings.Contains(err.Error(), "twice") { + t.Fatal("duplicate names must be rejected") + } +} + +func TestManagerRegisterTimeoutBounded(t *testing.T) { + // server 挂起不响应：initialize 必须在超时内 fail-closed，而不是永久阻塞。 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(10 * time.Second): + } + })) + defer srv.Close() + mgr, err := NewManager([]ServerConfig{{Name: "slow", URL: srv.URL, Timeout: 300 * time.Millisecond}}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + start := time.Now() + err = mgr.RegisterInto(context.Background(), reg) + elapsed := time.Since(start) + if err == nil || !strings.Contains(err.Error(), "unreachable") { + t.Fatalf("hung server must fail-closed, got %v", err) + } + if elapsed > 5*time.Second { + t.Fatalf("timeout must bound the dial, took %s", elapsed) + } +} diff --git a/runtime/mcp/protocol.go b/runtime/mcp/protocol.go new file mode 100644 index 0000000..c8474c0 --- /dev/null +++ b/runtime/mcp/protocol.go @@ -0,0 +1,140 @@ +// Package mcp 实现 Model Context Protocol 客户端（stdio 与 Streamable HTTP +// 两种传输），并把 server 工具桥接为 runtime/tools.Registry 的普通工具。 +// +// 边界声明：server 提供的描述与 schema 视为不可信输入（长度受限、不解释）； +// 工具调用的策略（allowlist/确认）仍由 tools.Registry 策略层统一执行。 +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "time" +) + +// ProtocolVersion 是本客户端声明的 MCP 协议版本。 +const ProtocolVersion = "2024-11-05" + +// ServerConfig 声明一个 MCP server 接入点：Command（stdio）或 URL（HTTP）二选一。 +type ServerConfig struct { + Name string `json:"name"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env []string `json:"env,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + // Timeout 单次请求上限（默认 30s）。 + Timeout time.Duration `json:"timeout,omitempty"` +} + +// Validate 校验配置完整性（fail-closed：命名不规范或双传输/无传输都拒绝）。 +func (c ServerConfig) Validate() error { + if !validName(c.Name) { + return fmt.Errorf("mcp server name %q must be [a-zA-Z0-9_-]+", c.Name) + } + hasCmd := c.Command != "" + hasURL := c.URL != "" + if hasCmd == hasURL { + return fmt.Errorf("mcp server %q: exactly one of command or url is required", c.Name) + } + return nil +} + +func validName(s string) bool { + if s == "" { + return false + } + for _, r := range s { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_': + default: + return false + } + } + return true +} + +// ---- JSON-RPC 2.0 envelope ---- + +type rpcRequest struct { + JSONRPC string `json:"jsonrpc"` + ID int64 `json:"id,omitempty"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + +type rpcNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` +} + +type rpcResponse struct { + JSONRPC string `json:"jsonrpc"` + ID *int64 `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` +} + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` + Data json.RawMessage `json:"data,omitempty"` +} + +func (e *rpcError) Error() string { + return fmt.Sprintf("mcp rpc error %d: %s", e.Code, e.Message) +} + +// ---- MCP 语义类型 ---- + +type initializeParams struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities json.RawMessage `json:"capabilities,omitempty"` + ClientInfo clientInfo `json:"clientInfo"` +} + +type clientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type initializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities json.RawMessage `json:"capabilities"` + ServerInfo clientInfo `json:"serverInfo"` + Instructions string `json:"instructions,omitempty"` +} + +type toolInfo struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema json.RawMessage `json:"inputSchema,omitempty"` +} + +type toolsListResult struct { + Tools []toolInfo `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` +} + +type callToolParams struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` +} + +type callToolResult struct { + Content []contentBlock `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +type contentBlock struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// Client 是一个已连接 MCP server 的最小面（initialize / tools/list / tools/call）。 +type Client interface { + Initialize(ctx context.Context) (initializeResult, error) + ListTools(ctx context.Context) ([]toolInfo, error) + CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) + Close() error +} diff --git a/runtime/mcp/stdio.go b/runtime/mcp/stdio.go new file mode 100644 index 0000000..2e5bcda --- /dev/null +++ b/runtime/mcp/stdio.go @@ -0,0 +1,198 @@ +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os/exec" + "sync" + "sync/atomic" +) + +// stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。 +// 同一时刻只允许一个在途请求（串行锁），匹配简单且足以支撑智能体工具调用。 +type stdioClient struct { + cmd *exec.Cmd + stdin io.WriteCloser + stdout *bufio.Reader + mu sync.Mutex + nextID int64 + closed bool +} + +// DialStdio 启动 server 子进程并完成 initialize 握手。 +func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) { + if err := cfg.Validate(); err != nil { + return nil, err + } + if cfg.Command == "" { + return nil, fmt.Errorf("mcp server %q: stdio transport requires command", cfg.Name) + } + cmd := exec.Command(cfg.Command, cfg.Args...) + cmd.Env = cfg.Env + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, err + } + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("mcp server %q: start %s: %w", cfg.Name, cfg.Command, err) + } + c := &stdioClient{cmd: cmd, stdin: stdin, stdout: bufio.NewReaderSize(stdout, 1<<20)} + if err := c.handshake(ctx); err != nil { + _ = c.Close() + return nil, err + } + return c, nil +} + +func (c *stdioClient) handshake(ctx context.Context) error { + var res initializeResult + if err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res); err != nil { + return err + } + // initialized 通知无需响应。 + return c.notify("notifications/initialized") +} + +func (c *stdioClient) Initialize(ctx context.Context) (initializeResult, error) { + // Initialize 已在 Dial 时完成；这里补一次协议层调用以支持探测场景。 + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return res, err +} + +func (c *stdioClient) ListTools(ctx context.Context) ([]toolInfo, error) { + var out []toolInfo + cursor := "" + for { + params := map[string]any{} + if cursor != "" { + params["cursor"] = cursor + } + var res toolsListResult + if err := c.roundTrip(ctx, "tools/list", params, &res); err != nil { + return nil, err + } + out = append(out, res.Tools...) + if res.NextCursor == "" || len(out) > maxToolsPerServer*2 { + return out, nil + } + cursor = res.NextCursor + } +} + +func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) { + var res callToolResult + if len(args) == 0 { + args = json.RawMessage(`{}`) + } + err := c.roundTrip(ctx, "tools/call", callToolParams{Name: name, Arguments: args}, &res) + return res, err +} + +func (c *stdioClient) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil + } + c.closed = true + _ = c.stdin.Close() + if c.cmd.Process != nil { + _ = c.cmd.Process.Kill() + } + return c.cmd.Wait() +} + +func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return fmt.Errorf("mcp client closed") + } + id := atomic.AddInt64(&c.nextID, 1) + req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params} + body, err := json.Marshal(req) + if err != nil { + return err + } + body = append(body, '\n') + if _, err := c.stdin.Write(body); err != nil { + return fmt.Errorf("mcp write %s: %w", method, err) + } + line, err := c.readResponseLine(ctx) + if err != nil { + return fmt.Errorf("mcp read %s: %w", method, err) + } + var resp rpcResponse + if err := json.Unmarshal(line, &resp); err != nil { + return fmt.Errorf("mcp %s: bad response frame: %w", method, err) + } + if resp.ID == nil || *resp.ID != id { + return fmt.Errorf("mcp %s: response id mismatch", method) + } + if resp.Error != nil { + return resp.Error + } + if out != nil && len(resp.Result) > 0 { + return json.Unmarshal(resp.Result, out) + } + return nil +} + +// readResponseLine 逐行读取，跳过通知帧（无 id），带 ctx 取消感知。 +func (c *stdioClient) readResponseLine(ctx context.Context) ([]byte, error) { + type result struct { + line []byte + err error + } + done := make(chan result, 1) + go func() { + for { + line, err := c.stdout.ReadBytes('\n') + if err != nil { + done <- result{nil, err} + return + } + var probe rpcResponse + if json.Unmarshal(line, &probe) == nil && probe.ID != nil { + done <- result{line, nil} + return + } + // 通知或日志行：忽略。 + } + }() + select { + case r := <-done: + return r.line, r.err + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (c *stdioClient) notify(method string) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return fmt.Errorf("mcp client closed") + } + body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method}) + if err != nil { + return err + } + body = append(body, '\n') + _, err = c.stdin.Write(body) + return err +} diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index 795d250..17ad4c8 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -62,28 +62,56 @@ func (r *Registry) Decls() []Decl { } // Restrict 把注册表收窄到 allow 列出的工具（工具策略的唯一执行点之一： -// 调用方声明 allowlist，注册表执行收窄）。allow 中出现未注册的工具名时 -// 返回错误并列出全部未知名（fail-closed，不静默忽略配置错误）。 +// 调用方声明 allowlist，注册表执行收窄）。条目支持两种形式： +// +// 精确名（read_file）或命名空间通配（mcp__github__*，匹配该前缀全部工具）。 +// +// 任一条目未命中任何已注册工具时返回错误（fail-closed，不静默忽略配置错误）。 // 空列表不收窄；可重复调用（在已收窄集合上继续求交）。 func (r *Registry) Restrict(allow []string) error { if len(allow) == 0 { return nil } - keep := make(map[string]bool, len(allow)) - var unknown []string - for _, name := range allow { - if _, ok := r.tools[name]; !ok { - unknown = append(unknown, name) + exact := make(map[string]bool, len(allow)) + prefixes := make([]string, 0) + for _, entry := range allow { + if name, ok := strings.CutSuffix(entry, "*"); ok { + prefixes = append(prefixes, name) continue } - keep[name] = true + exact[entry] = true + } + var unknown []string + for _, name := range allow { + if _, ok := strings.CutSuffix(name, "*"); ok { + continue + } + if _, ok := r.tools[name]; !ok { + unknown = append(unknown, name) + } + } + for _, p := range prefixes { + if !r.hasPrefix(p) { + unknown = append(unknown, p+"*") + } } if len(unknown) > 0 { return fmt.Errorf("tool allowlist names not registered: %s", strings.Join(unknown, ", ")) } + keep := func(name string) bool { + if exact[name] { + return true + } + for _, p := range prefixes { + if strings.HasPrefix(name, p) { + return true + } + } + return false + } var order []string for _, name := range r.order { - if keep[name] { + if keep(name) { order = append(order, name) } else { delete(r.tools, name) @@ -93,6 +121,15 @@ func (r *Registry) Restrict(allow []string) error { return nil } +func (r *Registry) hasPrefix(prefix string) bool { + for _, name := range r.order { + if strings.HasPrefix(name, prefix) { + return true + } + } + return false +} + // Decl is the schema view of a tool. type Decl struct { Name string