diff --git a/docs/guides/mcp-configuration.md b/docs/guides/mcp-configuration.md new file mode 100644 index 00000000..64f941f6 --- /dev/null +++ b/docs/guides/mcp-configuration.md @@ -0,0 +1,62 @@ +# MCP 配置指南（stdio） + +本文档说明如何在 NeoCode 中通过配置注册 MCP server，并验证 `mcp..` 能力是否可用。 + +## 配置位置 + +在 `~/.neocode/config.yaml` 中添加 `tools.mcp.servers`： + +```yaml +tools: + mcp: + servers: + - id: docs + enabled: true + source: stdio + version: v1 + stdio: + command: node + args: + - ./mcp-server.js + workdir: ./mcp + start_timeout_sec: 8 + call_timeout_sec: 20 + restart_backoff_sec: 1 + env: + - name: MCP_TOKEN + value_env: MCP_TOKEN +``` + +## 字段说明 + +- `id`：server 稳定标识，用于工具命名空间（`mcp..`）。 +- `enabled`：是否启用该 server；仅 `true` 的 server 会在启动时注册。 +- `source`：传输类型，当前仅支持 `stdio`。 +- `version`：可选版本字段，用于可观测和后续策略命中。 +- `stdio.command`：启动命令（必填，启用时）。 +- `stdio.args`：启动参数列表。 +- `stdio.workdir`：子进程工作目录，支持相对路径（相对主 `workdir` 解析）。 +- `stdio.start_timeout_sec` / `call_timeout_sec` / `restart_backoff_sec`：可选秒级超时与重试参数。 +- `env`：传给 MCP 子进程的环境变量列表。 + - 每项必须配置 `value` 或 `value_env` 其中之一。 + - 推荐使用 `value_env` 引用系统环境变量，避免在 YAML 中写明文敏感信息。 + +## 启动行为 + +- 启动阶段会注册所有 `enabled: true` 的 server。 +- 注册后会执行一次 `tools/list` 初始化工具快照。 +- 若启用的 server 注册失败，启动会报错并中止（fail-fast）。 + +## 功能测试建议 + +1. 启动应用后让 Agent 列出工具： + - `请先列出你当前可用工具的完整名称。` +2. 检查是否存在 `mcp.docs.`。 +3. 发起一次明确调用： + - `请调用 mcp.docs.search，参数 {\"query\":\"hello\"}，并返回工具结果。` + +若返回 `tool not found`，优先检查： +- `enabled` 是否为 `true`； +- `stdio.command` 是否可执行； +- `env.value_env` 对应环境变量是否存在； +- MCP server 是否支持 `tools/list`。 diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index dd0b8436..d57e8b25 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -56,7 +56,10 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { cfg := manager.Get() - toolRegistry := buildToolRegistry(cfg) + toolRegistry, err := buildToolRegistry(cfg) + if err != nil { + return nil, err + } toolManager, err := buildToolManager(toolRegistry) if err != nil { return nil, err @@ -82,7 +85,7 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { ), nil } -func buildToolRegistry(cfg config.Config) *tools.Registry { +func buildToolRegistry(cfg config.Config) (*tools.Registry, error) { toolRegistry := tools.NewRegistry() toolRegistry.Register(filesystem.New(cfg.Workdir)) toolRegistry.Register(filesystem.NewWrite(cfg.Workdir)) @@ -95,7 +98,14 @@ func buildToolRegistry(cfg config.Config) *tools.Registry { MaxResponseBytes: cfg.Tools.WebFetch.MaxResponseBytes, SupportedContentTypes: cfg.Tools.WebFetch.SupportedContentTypes, })) - return toolRegistry + mcpRegistry, err := buildMCPRegistry(cfg) + if err != nil { + return nil, err + } + if mcpRegistry != nil { + toolRegistry.SetMCPRegistry(mcpRegistry) + } + return toolRegistry, nil } func buildToolManager(registry *tools.Registry) (tools.Manager, error) { diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 3abbf11a..a81903c0 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -1,18 +1,25 @@ package app import ( + "bufio" + "bytes" "context" "encoding/json" "errors" + "fmt" + "io" "net/http" "net/http/httptest" "os" "path/filepath" + "strconv" "strings" "testing" + "time" "neo-code/internal/config" "neo-code/internal/tools" + "neo-code/internal/tools/mcp" ) func TestNewProgram(t *testing.T) { @@ -84,7 +91,10 @@ func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { cfg.Workdir = t.TempDir() cfg.Tools.WebFetch.MaxResponseBytes = 4 - registry := buildToolRegistry(cfg) + registry, err := buildToolRegistry(cfg) + if err != nil { + t.Fatalf("buildToolRegistry() error = %v", err) + } tool, err := registry.Get("webfetch") if err != nil { t.Fatalf("registry.Get(webfetch) error = %v", err) @@ -110,6 +120,337 @@ func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { } } +func TestBuildMCPRegistryFromConfig(t *testing.T) { + t.Parallel() + + stubClient := &stubMCPServerClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + } + + cfg := config.Default().Clone() + cfg.Workdir = t.TempDir() + cfg.Tools.MCP.Servers = []config.MCPServerConfig{ + { + ID: "docs", + Enabled: true, + Source: "stdio", + Stdio: config.MCPStdioConfig{ + Command: "mock", + }, + }, + } + + originalRegister := registerMCPStdioServer + t.Cleanup(func() { registerMCPStdioServer = originalRegister }) + registerMCPStdioServer = func(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + if err := registry.RegisterServer(server.ID, "stdio", server.Version, stubClient); err != nil { + return err + } + return registry.RefreshServerTools(context.Background(), server.ID) + } + + registry, err := buildMCPRegistry(cfg) + if err != nil { + t.Fatalf("buildMCPRegistry() error = %v", err) + } + if registry == nil { + t.Fatalf("expected non-nil mcp registry") + } + snapshots := registry.Snapshot() + if len(snapshots) != 1 || snapshots[0].ServerID != "docs" { + t.Fatalf("unexpected snapshots: %+v", snapshots) + } +} + +func TestBuildMCPRegistryUnsupportedSource(t *testing.T) { + t.Parallel() + + cfg := config.Default().Clone() + cfg.Workdir = t.TempDir() + cfg.Tools.MCP.Servers = []config.MCPServerConfig{ + { + ID: "docs", + Enabled: true, + Source: "sse", + Stdio: config.MCPStdioConfig{ + Command: "mock", + }, + }, + } + + registry, err := buildMCPRegistry(cfg) + if err == nil { + t.Fatalf("expected unsupported source error") + } + if registry != nil { + t.Fatalf("expected nil registry when source unsupported") + } + if !strings.Contains(strings.ToLower(err.Error()), "unsupported mcp source") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDefaultRegisterMCPStdioServerSuccess(t *testing.T) { + t.Parallel() + + registry := mcp.NewRegistry() + cfg := config.Default().Clone() + wd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + cfg.Workdir = wd + cfg.ToolTimeoutSec = 9 + + server := config.MCPServerConfig{ + ID: "docs", + Enabled: true, + Source: "stdio", + Version: "v1", + Stdio: config.MCPStdioConfig{ + Command: os.Args[0], + Args: []string{"-test.run=TestHelperProcessAppMCPStdioServer", "--"}, + Workdir: "", + StartTimeoutSec: 3, + CallTimeoutSec: 3, + }, + Env: []config.MCPEnvVarConfig{ + {Name: "MODE", Value: "test"}, + {Name: "GO_WANT_APP_MCP_STDIO_HELPER", Value: "1"}, + }, + } + t.Cleanup(func() { _ = registry.UnregisterServer("docs") }) + + if err := defaultRegisterMCPStdioServer(registry, cfg, server); err != nil { + t.Fatalf("defaultRegisterMCPStdioServer() error = %v", err) + } + + snapshots := registry.Snapshot() + if len(snapshots) != 1 || snapshots[0].ServerID != "docs" { + t.Fatalf("unexpected snapshots: %+v", snapshots) + } + if len(snapshots[0].Tools) != 1 || snapshots[0].Tools[0].Name != "search" { + t.Fatalf("unexpected tools snapshot: %+v", snapshots[0].Tools) + } +} + +func TestDefaultRegisterMCPStdioServerRefreshFailure(t *testing.T) { + t.Parallel() + + registry := mcp.NewRegistry() + cfg := config.Default().Clone() + wd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + cfg.Workdir = wd + + server := config.MCPServerConfig{ + ID: "broken", + Enabled: true, + Source: "stdio", + Stdio: config.MCPStdioConfig{ + Command: os.Args[0], + Args: []string{"-test.run=TestHelperProcessAppMCPStdioServer", "--"}, + StartTimeoutSec: 3, + CallTimeoutSec: 3, + }, + Env: []config.MCPEnvVarConfig{ + {Name: "GO_WANT_APP_MCP_STDIO_HELPER", Value: "1"}, + {Name: "GO_APP_MCP_STDIO_LIST_FAIL", Value: "1"}, + }, + } + t.Cleanup(func() { _ = registry.UnregisterServer("broken") }) + + err = defaultRegisterMCPStdioServer(registry, cfg, server) + if err == nil { + t.Fatalf("expected refresh failure") + } + if !strings.Contains(strings.ToLower(err.Error()), "list tools failed") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestBuildToolRegistryIncludesMCPFromConfig(t *testing.T) { + t.Parallel() + + cfg := config.Default().Clone() + cfg.Workdir = t.TempDir() + cfg.Tools.MCP.Servers = []config.MCPServerConfig{ + { + ID: "docs", + Enabled: true, + Source: "stdio", + Stdio: config.MCPStdioConfig{ + Command: "mock", + }, + }, + } + + originalRegister := registerMCPStdioServer + t.Cleanup(func() { registerMCPStdioServer = originalRegister }) + registerMCPStdioServer = func(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + client := &stubMCPServerClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + } + if err := registry.RegisterServer(server.ID, "stdio", server.Version, client); err != nil { + return err + } + return registry.RefreshServerTools(context.Background(), server.ID) + } + + registry, err := buildToolRegistry(cfg) + if err != nil { + t.Fatalf("buildToolRegistry() error = %v", err) + } + specs, err := registry.ListAvailableSpecs(context.Background(), tools.SpecListInput{}) + if err != nil { + t.Fatalf("ListAvailableSpecs() error = %v", err) + } + found := false + for _, spec := range specs { + if spec.Name == "mcp.docs.search" { + found = true + break + } + } + if !found { + t.Fatalf("expected mcp.docs.search in specs, got %+v", specs) + } +} + +func TestResolveMCPServerEnvAndWorkdir(t *testing.T) { + t.Setenv("MCP_TOKEN", "secret") + env, err := resolveMCPServerEnv(config.MCPServerConfig{ + Env: []config.MCPEnvVarConfig{ + {Name: "TOKEN", ValueEnv: "MCP_TOKEN"}, + {Name: "MODE", Value: "test"}, + }, + }) + if err != nil { + t.Fatalf("resolveMCPServerEnv() error = %v", err) + } + joined := strings.Join(env, ",") + if !strings.Contains(joined, "TOKEN=secret") || !strings.Contains(joined, "MODE=test") { + t.Fatalf("unexpected env result: %+v", env) + } + + base := t.TempDir() + relative := resolveMCPServerWorkdir(base, "tools/mcp") + if !strings.HasSuffix(filepath.ToSlash(relative), "tools/mcp") { + t.Fatalf("unexpected relative workdir: %q", relative) + } + absoluteTarget := filepath.Join(t.TempDir(), "absolute") + absolute := resolveMCPServerWorkdir(base, absoluteTarget) + if absolute != filepath.Clean(absoluteTarget) { + t.Fatalf("unexpected absolute workdir: %q", absolute) + } +} + +func TestResolveMCPServerEnvValidationErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + server config.MCPServerConfig + }{ + { + name: "empty name", + server: config.MCPServerConfig{ + Env: []config.MCPEnvVarConfig{{Name: " ", Value: "x"}}, + }, + }, + { + name: "both value and value_env", + server: config.MCPServerConfig{ + Env: []config.MCPEnvVarConfig{{Name: "A", Value: "x", ValueEnv: "B"}}, + }, + }, + { + name: "missing value and value_env", + server: config.MCPServerConfig{ + Env: []config.MCPEnvVarConfig{{Name: "A"}}, + }, + }, + { + name: "value_env unresolved", + server: config.MCPServerConfig{ + Env: []config.MCPEnvVarConfig{{Name: "A", ValueEnv: "MISSING_ENV_FOR_TEST"}}, + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if _, err := resolveMCPServerEnv(tt.server); err == nil { + t.Fatalf("expected validation error") + } + }) + } +} + +func TestBuildMCPRegistryNoEnabledServerReturnsNil(t *testing.T) { + t.Parallel() + + cfg := config.Default().Clone() + cfg.Workdir = t.TempDir() + cfg.Tools.MCP.Servers = []config.MCPServerConfig{ + {ID: "docs", Enabled: false, Source: "stdio"}, + } + + registry, err := buildMCPRegistry(cfg) + if err != nil { + t.Fatalf("buildMCPRegistry() error = %v", err) + } + if registry != nil { + t.Fatalf("expected nil registry when no enabled server") + } +} + +func TestBuildMCPRegistryRegisterError(t *testing.T) { + t.Parallel() + + cfg := config.Default().Clone() + cfg.Workdir = t.TempDir() + cfg.Tools.MCP.Servers = []config.MCPServerConfig{ + {ID: "docs", Enabled: true, Source: "stdio"}, + } + + originalRegister := registerMCPStdioServer + t.Cleanup(func() { registerMCPStdioServer = originalRegister }) + registerMCPStdioServer = func(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + return errors.New("register failed") + } + + _, err := buildMCPRegistry(cfg) + if err == nil || !strings.Contains(err.Error(), "register failed") { + t.Fatalf("expected wrapped register error, got %v", err) + } +} + +func TestInitialMCPRefreshTimeoutAndDurationConversion(t *testing.T) { + t.Parallel() + + cfg := config.Default().Clone() + cfg.ToolTimeoutSec = 1 + timeout := initialMCPRefreshTimeout(cfg) + if timeout < 5*time.Second { + t.Fatalf("expected minimum timeout >= 5s, got %v", timeout) + } + if durationFromSeconds(0) != 0 { + t.Fatalf("expected zero duration for non-positive input") + } + if durationFromSeconds(2) != 2*time.Second { + t.Fatalf("expected 2s duration") + } +} + func TestBuildToolManagerWrapsRegistry(t *testing.T) { t.Parallel() @@ -252,3 +593,173 @@ func disableBuiltinProviderAPIKeys(t *testing.T) { t.Setenv(config.OpenLLDefaultAPIKeyEnv, "") t.Setenv(config.QiniuDefaultAPIKeyEnv, "") } + +type stubMCPServerClient struct { + tools []mcp.ToolDescriptor + listErr error +} + +func (s *stubMCPServerClient) ListTools(ctx context.Context) ([]mcp.ToolDescriptor, error) { + if s.listErr != nil { + return nil, s.listErr + } + return append([]mcp.ToolDescriptor(nil), s.tools...), nil +} + +func (s *stubMCPServerClient) CallTool(ctx context.Context, toolName string, arguments []byte) (mcp.CallResult, error) { + return mcp.CallResult{Content: "ok"}, nil +} + +func (s *stubMCPServerClient) HealthCheck(ctx context.Context) error { + return nil +} + +func TestHelperProcessAppMCPStdioServer(t *testing.T) { + if os.Getenv("GO_WANT_APP_MCP_STDIO_HELPER") != "1" { + return + } + + listFail := os.Getenv("GO_APP_MCP_STDIO_LIST_FAIL") == "1" + initialized := false + reader := bufio.NewReader(os.Stdin) + + for { + payload, err := readFramedForAppTest(reader) + if err != nil { + if errors.Is(err, os.ErrClosed) || strings.Contains(strings.ToLower(err.Error()), "eof") { + os.Exit(0) + } + os.Exit(2) + } + + var request map[string]any + if err := json.Unmarshal(payload, &request); err != nil { + os.Exit(3) + } + + method, _ := request["method"].(string) + requestID, _ := request["id"].(string) + var response any + + switch method { + case "initialize": + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{}, + "serverInfo": map[string]any{ + "name": "app-helper", + "version": "1.0.0", + }, + }, + } + case "notifications/initialized": + initialized = true + continue + case "tools/list": + if listFail { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32001, + "message": "list tools failed", + }, + } + break + } + if !initialized { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32002, + "message": "server not initialized", + }, + } + break + } + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": map[string]any{ + "tools": []map[string]any{ + { + "name": "search", + "description": "search docs", + "inputSchema": map[string]any{ + "type": "object", + "properties": map[string]any{"query": map[string]any{"type": "string"}}, + }, + }, + }, + }, + } + default: + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32601, + "message": "method not found", + }, + } + } + + rawResponse, err := json.Marshal(response) + if err != nil { + os.Exit(4) + } + if err := writeFramedForAppTest(os.Stdout, rawResponse); err != nil { + os.Exit(5) + } + } +} + +func readFramedForAppTest(reader *bufio.Reader) ([]byte, error) { + contentLength := -1 + for { + line, err := reader.ReadString('\n') + if err != nil { + return nil, err + } + trimmed := strings.TrimSpace(line) + if trimmed == "" { + if contentLength >= 0 { + break + } + continue + } + lower := strings.ToLower(trimmed) + if strings.HasPrefix(lower, "content-length:") { + rawLength := strings.TrimSpace(trimmed[len("content-length:"):]) + length, convErr := strconv.Atoi(rawLength) + if convErr != nil { + return nil, convErr + } + contentLength = length + continue + } + } + if contentLength < 0 { + return nil, errors.New("missing content-length") + } + payload := make([]byte, contentLength) + if _, err := io.ReadFull(reader, payload); err != nil { + return nil, err + } + return payload, nil +} + +func writeFramedForAppTest(writer io.Writer, payload []byte) error { + header := fmt.Sprintf("Content-Length: %d\r\n\r\n", len(payload)) + if _, err := io.WriteString(writer, header); err != nil { + return err + } + if _, err := writer.Write(bytes.TrimSpace(payload)); err != nil { + return err + } + return nil +} diff --git a/internal/app/mcp_bootstrap.go b/internal/app/mcp_bootstrap.go new file mode 100644 index 00000000..27982e14 --- /dev/null +++ b/internal/app/mcp_bootstrap.go @@ -0,0 +1,154 @@ +package app + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "neo-code/internal/config" + "neo-code/internal/tools/mcp" +) + +var newMCPStdioClient = mcp.NewStdIOClient +var registerMCPStdioServer = defaultRegisterMCPStdioServer + +// buildMCPRegistry 按配置构建并初始化 MCP registry；若无启用 server 则返回 nil。 +func buildMCPRegistry(cfg config.Config) (*mcp.Registry, error) { + if len(cfg.Tools.MCP.Servers) == 0 { + return nil, nil + } + + registry := mcp.NewRegistry() + enabledCount := 0 + for index := range cfg.Tools.MCP.Servers { + server := cfg.Tools.MCP.Servers[index] + if !server.Enabled { + continue + } + enabledCount++ + + switch strings.ToLower(strings.TrimSpace(server.Source)) { + case "", "stdio": + if err := registerMCPStdioServer(registry, cfg, server); err != nil { + return nil, fmt.Errorf("app: register mcp server %q: %w", strings.TrimSpace(server.ID), err) + } + default: + return nil, fmt.Errorf("app: unsupported mcp source %q", server.Source) + } + } + + if enabledCount == 0 { + return nil, nil + } + return registry, nil +} + +// defaultRegisterMCPStdioServer 创建 stdio client 并完成 server 注册与 tools 快照初始化。 +func defaultRegisterMCPStdioServer(registry *mcp.Registry, cfg config.Config, server config.MCPServerConfig) error { + env, err := resolveMCPServerEnv(server) + if err != nil { + return err + } + + workdir := resolveMCPServerWorkdir(cfg.Workdir, server.Stdio.Workdir) + client, err := newMCPStdioClient(mcp.StdioClientConfig{ + Command: strings.TrimSpace(server.Stdio.Command), + Args: append([]string(nil), server.Stdio.Args...), + Env: env, + Workdir: workdir, + StartTimeout: durationFromSeconds(server.Stdio.StartTimeoutSec), + CallTimeout: durationFromSeconds(server.Stdio.CallTimeoutSec), + RestartBackoff: durationFromSeconds(server.Stdio.RestartBackoffSec), + }) + if err != nil { + return err + } + + serverID := strings.TrimSpace(server.ID) + source := strings.ToLower(strings.TrimSpace(server.Source)) + if source == "" { + source = "stdio" + } + if err := registry.RegisterServer(serverID, source, strings.TrimSpace(server.Version), client); err != nil { + return err + } + + refreshCtx, cancel := context.WithTimeout(context.Background(), initialMCPRefreshTimeout(cfg)) + defer cancel() + if err := registry.RefreshServerTools(refreshCtx, serverID); err != nil { + return err + } + return nil +} + +// resolveMCPServerEnv 将配置中的 env 绑定解析为子进程环境变量。 +func resolveMCPServerEnv(server config.MCPServerConfig) ([]string, error) { + if len(server.Env) == 0 { + return nil, nil + } + result := make([]string, 0, len(server.Env)) + for index, item := range server.Env { + name := strings.TrimSpace(item.Name) + if name == "" { + return nil, fmt.Errorf("env[%d].name is empty", index) + } + + value := strings.TrimSpace(item.Value) + valueEnv := strings.TrimSpace(item.ValueEnv) + switch { + case value != "" && valueEnv != "": + return nil, fmt.Errorf("env[%d] must set either value or value_env", index) + case value != "": + result = append(result, name+"="+value) + case valueEnv != "": + resolved := strings.TrimSpace(os.Getenv(valueEnv)) + if resolved == "" { + return nil, fmt.Errorf("env[%d] value_env %q is empty", index, valueEnv) + } + result = append(result, name+"="+resolved) + default: + return nil, fmt.Errorf("env[%d] must set one of value/value_env", index) + } + } + return result, nil +} + +// resolveMCPServerWorkdir 解析 MCP server 子进程工作目录，支持相对路径。 +func resolveMCPServerWorkdir(baseWorkdir string, override string) string { + trimmedOverride := strings.TrimSpace(override) + if trimmedOverride == "" { + return strings.TrimSpace(baseWorkdir) + } + if filepath.IsAbs(trimmedOverride) { + return filepath.Clean(trimmedOverride) + } + + trimmedBase := strings.TrimSpace(baseWorkdir) + if trimmedBase == "" { + return filepath.Clean(trimmedOverride) + } + return filepath.Clean(filepath.Join(trimmedBase, trimmedOverride)) +} + +// initialMCPRefreshTimeout 计算启动阶段首轮 tools 刷新的超时时间。 +func initialMCPRefreshTimeout(cfg config.Config) time.Duration { + timeout := time.Duration(cfg.ToolTimeoutSec) * time.Second + if timeout <= 0 { + timeout = 20 * time.Second + } + if timeout < 5*time.Second { + timeout = 5 * time.Second + } + return timeout +} + +// durationFromSeconds 将秒级配置转换为 duration；非正值返回 0 以启用 client 默认值。 +func durationFromSeconds(seconds int) time.Duration { + if seconds <= 0 { + return 0 + } + return time.Duration(seconds) * time.Second +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index a2504ac2..29b11b52 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -471,6 +471,47 @@ func TestConfigValidateFailures(t *testing.T) { }(), expectErr: "duplicate provider endpoint", }, + { + name: "invalid mcp duplicate server id", + config: func() *Config { + cfg := validConfig.Clone() + cfg.Tools.MCP.Servers = []MCPServerConfig{ + {ID: "docs", Enabled: true, Stdio: MCPStdioConfig{Command: "cmd-1"}}, + {ID: "docs", Enabled: true, Stdio: MCPStdioConfig{Command: "cmd-2"}}, + } + return &cfg + }(), + expectErr: "duplicate servers", + }, + { + name: "invalid mcp source", + config: func() *Config { + cfg := validConfig.Clone() + cfg.Tools.MCP.Servers = []MCPServerConfig{ + {ID: "docs", Enabled: true, Source: "sse", Stdio: MCPStdioConfig{Command: "cmd"}}, + } + return &cfg + }(), + expectErr: "not supported", + }, + { + name: "invalid mcp env binding", + config: func() *Config { + cfg := validConfig.Clone() + cfg.Tools.MCP.Servers = []MCPServerConfig{ + { + ID: "docs", + Enabled: true, + Stdio: MCPStdioConfig{Command: "cmd"}, + Env: []MCPEnvVarConfig{ + {Name: "TOKEN", Value: "a", ValueEnv: "TOKEN_ENV"}, + }, + }, + } + return &cfg + }(), + expectErr: "exactly one of value/value_env", + }, } for _, tt := range tests { @@ -485,6 +526,34 @@ func TestConfigValidateFailures(t *testing.T) { } } +func TestMCPConfigApplyDefaultsAndClone(t *testing.T) { + t.Parallel() + + cfg := MCPConfig{ + Servers: []MCPServerConfig{ + { + ID: " Docs ", + Enabled: true, + Source: "", + Stdio: MCPStdioConfig{ + Command: "mock", + Args: []string{"a"}, + }, + }, + }, + } + cfg.ApplyDefaults(defaultMCPConfig()) + if cfg.Servers[0].Source != "stdio" { + t.Fatalf("expected default source stdio, got %q", cfg.Servers[0].Source) + } + + cloned := cfg.Clone() + cloned.Servers[0].Stdio.Args[0] = "b" + if cfg.Servers[0].Stdio.Args[0] == "b" { + t.Fatalf("expected MCP clone to be independent") + } +} + func TestProviderConfigValidateFailures(t *testing.T) { t.Parallel() diff --git a/internal/config/model.go b/internal/config/model.go index 75992d1c..3296dc20 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -60,6 +60,7 @@ type ResolvedProviderConfig struct { type ToolsConfig struct { WebFetch WebFetchConfig `yaml:"webfetch,omitempty"` + MCP MCPConfig `yaml:"mcp,omitempty"` } type ContextConfig struct { @@ -78,6 +79,34 @@ type WebFetchConfig struct { SupportedContentTypes []string `yaml:"supported_content_types,omitempty"` } +type MCPConfig struct { + Servers []MCPServerConfig `yaml:"servers,omitempty"` +} + +type MCPServerConfig struct { + ID string `yaml:"id"` + Enabled bool `yaml:"enabled,omitempty"` + Source string `yaml:"source,omitempty"` + Version string `yaml:"version,omitempty"` + Stdio MCPStdioConfig `yaml:"stdio,omitempty"` + Env []MCPEnvVarConfig `yaml:"env,omitempty"` +} + +type MCPStdioConfig struct { + Command string `yaml:"command,omitempty"` + Args []string `yaml:"args,omitempty"` + Workdir string `yaml:"workdir,omitempty"` + StartTimeoutSec int `yaml:"start_timeout_sec,omitempty"` + CallTimeoutSec int `yaml:"call_timeout_sec,omitempty"` + RestartBackoffSec int `yaml:"restart_backoff_sec,omitempty"` +} + +type MCPEnvVarConfig struct { + Name string `yaml:"name"` + Value string `yaml:"value,omitempty"` + ValueEnv string `yaml:"value_env,omitempty"` +} + func DefaultWebFetchSupportedContentTypes() []string { return append([]string(nil), defaultWebFetchSupportedContentTypes...) } @@ -91,6 +120,7 @@ func Default() *Config { Context: defaultContextConfig(), Tools: ToolsConfig{ WebFetch: defaultWebFetchConfig(), + MCP: defaultMCPConfig(), }, } } @@ -330,6 +360,13 @@ func defaultWebFetchConfig() WebFetchConfig { } } +// defaultMCPConfig 返回 MCP 工具接入配置的默认值（默认无 server）。 +func defaultMCPConfig() MCPConfig { + return MCPConfig{ + Servers: nil, + } +} + // defaultContextConfig 返回上下文压缩相关配置的默认值。 func defaultContextConfig() ContextConfig { return ContextConfig{ @@ -349,6 +386,7 @@ func defaultCompactConfig() CompactConfig { func (c ToolsConfig) Clone() ToolsConfig { return ToolsConfig{ WebFetch: c.WebFetch.Clone(), + MCP: c.MCP.Clone(), } } @@ -365,6 +403,7 @@ func (c *ToolsConfig) ApplyDefaults(defaults ToolsConfig) { } c.WebFetch.ApplyDefaults(defaults.WebFetch) + c.MCP.ApplyDefaults(defaults.MCP) } // ApplyDefaults 为上下文配置补齐缺省的 compact 参数。 @@ -380,6 +419,93 @@ func (c ToolsConfig) Validate() error { if err := c.WebFetch.Validate(); err != nil { return fmt.Errorf("webfetch: %w", err) } + if err := c.MCP.Validate(); err != nil { + return fmt.Errorf("mcp: %w", err) + } + return nil +} + +// Clone 返回 MCP 配置的独立副本，避免引用共享造成并发污染。 +func (c MCPConfig) Clone() MCPConfig { + if len(c.Servers) == 0 { + return MCPConfig{} + } + cloned := make([]MCPServerConfig, 0, len(c.Servers)) + for _, server := range c.Servers { + cloned = append(cloned, server.Clone()) + } + return MCPConfig{ + Servers: cloned, + } +} + +// Clone 返回单个 MCP server 配置的独立副本。 +func (c MCPServerConfig) Clone() MCPServerConfig { + cloned := c + cloned.Stdio.Args = append([]string(nil), c.Stdio.Args...) + if len(c.Env) > 0 { + cloned.Env = make([]MCPEnvVarConfig, 0, len(c.Env)) + cloned.Env = append(cloned.Env, c.Env...) + } else { + cloned.Env = nil + } + return cloned +} + +// ApplyDefaults 为 MCP 配置补齐缺省字段，保证运行时行为可预测。 +func (c *MCPConfig) ApplyDefaults(defaults MCPConfig) { + if c == nil { + return + } + if len(c.Servers) == 0 { + c.Servers = defaults.Clone().Servers + } + for index := range c.Servers { + c.Servers[index].ApplyDefaults() + } +} + +// Validate 校验 MCP server 列表与字段合法性，防止启动后失败。 +func (c MCPConfig) Validate() error { + if len(c.Servers) == 0 { + return nil + } + seen := make(map[string]struct{}, len(c.Servers)) + for index, server := range c.Servers { + normalizedID := strings.ToLower(strings.TrimSpace(server.ID)) + if normalizedID == "" { + return fmt.Errorf("servers[%d].id is empty", index) + } + if _, exists := seen[normalizedID]; exists { + return fmt.Errorf("duplicate servers[%d].id %q", index, server.ID) + } + seen[normalizedID] = struct{}{} + + source := strings.ToLower(strings.TrimSpace(server.Source)) + if source == "" { + source = "stdio" + } + if source != "stdio" { + return fmt.Errorf("servers[%d].source %q is not supported", index, server.Source) + } + if !server.Enabled { + continue + } + + if strings.TrimSpace(server.Stdio.Command) == "" { + return fmt.Errorf("servers[%d].stdio.command is empty", index) + } + for envIndex, env := range server.Env { + if strings.TrimSpace(env.Name) == "" { + return fmt.Errorf("servers[%d].env[%d].name is empty", index, envIndex) + } + hasValue := strings.TrimSpace(env.Value) != "" + hasValueEnv := strings.TrimSpace(env.ValueEnv) != "" + if hasValue == hasValueEnv { + return fmt.Errorf("servers[%d].env[%d] must set exactly one of value/value_env", index, envIndex) + } + } + } return nil } @@ -413,6 +539,17 @@ func (c *WebFetchConfig) ApplyDefaults(defaults WebFetchConfig) { c.SupportedContentTypes = normalizeContentTypes(c.SupportedContentTypes, defaults.SupportedContentTypes) } +// ApplyDefaults 为 MCP server stdio 配置补齐默认值。 +func (c *MCPServerConfig) ApplyDefaults() { + if c == nil { + return + } + c.Source = strings.ToLower(strings.TrimSpace(c.Source)) + if c.Source == "" { + c.Source = "stdio" + } +} + // ApplyDefaults 为 compact 配置填充缺省策略和阈值。 func (c *CompactConfig) ApplyDefaults(defaults CompactConfig) { if c == nil { diff --git a/internal/tools/mcp/adapter.go b/internal/tools/mcp/adapter.go new file mode 100644 index 00000000..d9ecb470 --- /dev/null +++ b/internal/tools/mcp/adapter.go @@ -0,0 +1,142 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" +) + +const mcpToolNamePrefix = "mcp." + +// AdapterFactory 基于 registry 快照构造 MCP tool 适配器集合。 +type AdapterFactory struct { + registry *Registry +} + +// NewAdapterFactory 创建 MCP adapter 工厂。 +func NewAdapterFactory(registry *Registry) *AdapterFactory { + return &AdapterFactory{registry: registry} +} + +// BuildAdapters 将当前所有 MCP tool 快照转换为 Adapter 列表。 +func (f *AdapterFactory) BuildAdapters(ctx context.Context) ([]*Adapter, error) { + if f == nil || f.registry == nil { + return nil, errors.New("mcp: adapter factory registry is nil") + } + if err := ctx.Err(); err != nil { + return nil, err + } + + snapshots := f.registry.Snapshot() + if len(snapshots) == 0 { + return nil, nil + } + + result := make([]*Adapter, 0, len(snapshots)*2) + for _, snapshot := range snapshots { + for _, descriptor := range snapshot.Tools { + adapter, err := NewAdapter(f.registry, snapshot.ServerID, descriptor) + if err != nil { + return nil, err + } + result = append(result, adapter) + } + } + return result, nil +} + +// Adapter 将单个 MCP tool 适配为统一调用描述。 +type Adapter struct { + registry *Registry + serverID string + toolName string + description string + schema map[string]any +} + +// NewAdapter 创建指定 server/tool 的 MCP 适配器。 +func NewAdapter(registry *Registry, serverID string, descriptor ToolDescriptor) (*Adapter, error) { + if registry == nil { + return nil, errors.New("mcp: registry is nil") + } + normalizedServerID := normalizeServerID(serverID) + if normalizedServerID == "" { + return nil, errors.New("mcp: server id is empty") + } + normalizedToolName := strings.TrimSpace(descriptor.Name) + if normalizedToolName == "" { + return nil, errors.New("mcp: descriptor tool name is empty") + } + + return &Adapter{ + registry: registry, + serverID: normalizedServerID, + toolName: normalizedToolName, + description: strings.TrimSpace(descriptor.Description), + schema: ensureObjectSchema(descriptor.InputSchema), + }, nil +} + +// FullName 返回统一的 MCP tool 名称：mcp..。 +func (a *Adapter) FullName() string { + return composeToolName(a.serverID, a.toolName) +} + +// ServerID 返回 MCP server 标识。 +func (a *Adapter) ServerID() string { + return a.serverID +} + +// ToolName 返回 MCP tool 原始名称。 +func (a *Adapter) ToolName() string { + return a.toolName +} + +// Description 返回工具描述，不存在时回退到稳定默认文案。 +func (a *Adapter) Description() string { + if strings.TrimSpace(a.description) != "" { + return a.description + } + return fmt.Sprintf("MCP tool %s from server %s", a.toolName, a.serverID) +} + +// Schema 返回 MCP 工具输入 schema 的标准对象结构。 +func (a *Adapter) Schema() map[string]any { + return cloneSchema(a.schema) +} + +// Call 分发 MCP tool 调用并返回统一结果。 +func (a *Adapter) Call(ctx context.Context, arguments []byte) (CallResult, error) { + if a == nil || a.registry == nil { + return CallResult{}, errors.New("mcp: adapter is not initialized") + } + if err := ctx.Err(); err != nil { + return CallResult{}, err + } + return a.registry.Call(ctx, a.serverID, a.toolName, arguments) +} + +// composeToolName 组装统一的 MCP tool 名称，保持权限映射可预测。 +func composeToolName(serverID string, toolName string) string { + return mcpToolNamePrefix + normalizeServerID(serverID) + "." + strings.TrimSpace(toolName) +} + +// ensureObjectSchema 确保 schema 至少是 object，避免上层 provider 解析异常。 +func ensureObjectSchema(schema map[string]any) map[string]any { + cloned := cloneSchema(schema) + if len(cloned) == 0 { + return map[string]any{ + "type": "object", + "properties": map[string]any{}, + } + } + + if strings.TrimSpace(fmt.Sprintf("%v", cloned["type"])) == "" { + cloned["type"] = "object" + } + if _, ok := cloned["properties"]; !ok { + cloned["properties"] = map[string]any{} + } + return cloned +} diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go new file mode 100644 index 00000000..2b38691f --- /dev/null +++ b/internal/tools/mcp/adapter_test.go @@ -0,0 +1,229 @@ +package mcp + +import ( + "context" + "errors" + "testing" +) + +func TestAdapterFactoryBuildAdapters(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + { + Name: "search", + Description: "search docs", + InputSchema: map[string]any{"type": "object"}, + }, + }, + } + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh tools: %v", err) + } + + factory := NewAdapterFactory(registry) + adapters, err := factory.BuildAdapters(context.Background()) + if err != nil { + t.Fatalf("BuildAdapters() error = %v", err) + } + if len(adapters) != 1 { + t.Fatalf("expected one adapter, got %d", len(adapters)) + } + if adapters[0].FullName() != "mcp.docs.search" { + t.Fatalf("unexpected adapter full name: %q", adapters[0].FullName()) + } +} + +func TestAdapterFactoryBuildAdaptersEmptySnapshot(t *testing.T) { + t.Parallel() + + factory := NewAdapterFactory(NewRegistry()) + adapters, err := factory.BuildAdapters(context.Background()) + if err != nil { + t.Fatalf("BuildAdapters() error = %v", err) + } + if len(adapters) != 0 { + t.Fatalf("expected empty adapters, got %d", len(adapters)) + } +} + +func TestAdapterCall(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + {Name: "search", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: CallResult{ + Content: "result body", + Metadata: map[string]any{ + "latency_ms": 20, + }, + }, + } + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh tools: %v", err) + } + + adapter, err := NewAdapter(registry, "docs", ToolDescriptor{ + Name: "search", + Description: "search docs", + InputSchema: map[string]any{"type": "object"}, + }) + if err != nil { + t.Fatalf("NewAdapter() error = %v", err) + } + + result, err := adapter.Call(context.Background(), []byte(`{"q":"mcp"}`)) + if err != nil { + t.Fatalf("Call() error = %v", err) + } + if result.Content != "result body" { + t.Fatalf("expected result content, got %q", result.Content) + } +} + +func TestAdapterCallError(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + {Name: "search", InputSchema: map[string]any{"type": "object"}}, + }, + callErr: errors.New("transport timeout"), + } + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh tools: %v", err) + } + + adapter, err := NewAdapter(registry, "docs", ToolDescriptor{ + Name: "search", + }) + if err != nil { + t.Fatalf("NewAdapter() error = %v", err) + } + + if _, err := adapter.Call(context.Background(), []byte(`{"q":"mcp"}`)); err == nil { + t.Fatalf("expected call error") + } +} + +func TestAdapterAccessorsAndSchemaClone(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + adapter, err := NewAdapter(registry, "Docs", ToolDescriptor{ + Name: "search", + Description: "", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "q": map[string]any{"type": "string"}, + }, + }, + }) + if err != nil { + t.Fatalf("NewAdapter() error = %v", err) + } + + if adapter.ServerID() != "docs" { + t.Fatalf("expected normalized server id docs, got %q", adapter.ServerID()) + } + if adapter.ToolName() != "search" { + t.Fatalf("expected tool name search, got %q", adapter.ToolName()) + } + if adapter.Description() == "" { + t.Fatalf("expected non-empty fallback description") + } + + schema1 := adapter.Schema() + schema2 := adapter.Schema() + props1, _ := schema1["properties"].(map[string]any) + props1["q"] = map[string]any{"type": "number"} + props2, _ := schema2["properties"].(map[string]any) + query2, _ := props2["q"].(map[string]any) + if query2["type"] != "string" { + t.Fatalf("expected schema clone not mutated, got %v", query2["type"]) + } +} + +func TestAdapterBuildAndCreateErrors(t *testing.T) { + t.Parallel() + + factory := NewAdapterFactory(nil) + if _, err := factory.BuildAdapters(context.Background()); err == nil { + t.Fatalf("expected nil registry error") + } + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + registry := NewRegistry() + factory = NewAdapterFactory(registry) + if _, err := factory.BuildAdapters(canceledCtx); err == nil { + t.Fatalf("expected canceled context error") + } + + if _, err := NewAdapter(nil, "docs", ToolDescriptor{Name: "search"}); err == nil { + t.Fatalf("expected nil registry error") + } + if _, err := NewAdapter(registry, " ", ToolDescriptor{Name: "search"}); err == nil { + t.Fatalf("expected empty server id error") + } + if _, err := NewAdapter(registry, "docs", ToolDescriptor{Name: " "}); err == nil { + t.Fatalf("expected empty tool name error") + } +} + +func TestAdapterCallBoundary(t *testing.T) { + t.Parallel() + + var nilAdapter *Adapter + if _, err := nilAdapter.Call(context.Background(), nil); err == nil { + t.Fatalf("expected nil adapter error") + } + + registry := NewRegistry() + adapter, err := NewAdapter(registry, "docs", ToolDescriptor{Name: "search"}) + if err != nil { + t.Fatalf("NewAdapter() error = %v", err) + } + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := adapter.Call(canceledCtx, nil); err == nil { + t.Fatalf("expected context canceled error") + } +} + +func TestAdapterEnsureObjectSchemaDefaults(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + adapter, err := NewAdapter(registry, "docs", ToolDescriptor{ + Name: "search", + Description: "search docs", + InputSchema: map[string]any{}, + }) + if err != nil { + t.Fatalf("NewAdapter() error = %v", err) + } + schema := adapter.Schema() + if schema["type"] != "object" { + t.Fatalf("expected object type, got %v", schema["type"]) + } + if _, ok := schema["properties"].(map[string]any); !ok { + t.Fatalf("expected properties object, got %+v", schema["properties"]) + } +} diff --git a/internal/tools/mcp/registry.go b/internal/tools/mcp/registry.go new file mode 100644 index 00000000..490b5142 --- /dev/null +++ b/internal/tools/mcp/registry.go @@ -0,0 +1,339 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "sync" + "time" +) + +// ServerStatus 描述 MCP server 在 registry 中的生命周期状态。 +type ServerStatus string + +const ( + // ServerStatusConnecting 表示 server 已注册但仍在连接或初始化阶段。 + ServerStatusConnecting ServerStatus = "connecting" + // ServerStatusReady 表示 server 已可用，支持正常工具调用。 + ServerStatusReady ServerStatus = "ready" + // ServerStatusDegraded 表示 server 可部分服务，但存在健康或调用异常。 + ServerStatusDegraded ServerStatus = "degraded" + // ServerStatusOffline 表示 server 当前不可用。 + ServerStatusOffline ServerStatus = "offline" +) + +// ToolDescriptor 描述 MCP tool 的稳定元信息与输入 schema。 +type ToolDescriptor struct { + Name string + Description string + InputSchema map[string]any +} + +// ServerSnapshot 描述 registry 对外暴露的 server 只读快照。 +type ServerSnapshot struct { + ServerID string + Source string + Version string + Status ServerStatus + UpdatedAt time.Time + Tools []ToolDescriptor +} + +// CallResult 收敛 MCP tool 调用后的统一结果语义。 +type CallResult struct { + Content string + IsError bool + Metadata map[string]any +} + +// ServerClient 描述 registry 与具体 MCP server 交互所需的最小能力。 +type ServerClient interface { + ListTools(ctx context.Context) ([]ToolDescriptor, error) + CallTool(ctx context.Context, toolName string, arguments []byte) (CallResult, error) + HealthCheck(ctx context.Context) error +} + +type serverEntry struct { + snapshot ServerSnapshot + client ServerClient +} + +// Registry 维护 MCP server 注册、快照读取和工具调用分发。 +type Registry struct { + mu sync.RWMutex + servers map[string]*serverEntry +} + +// NewRegistry 创建线程安全的 MCP registry 实例。 +func NewRegistry() *Registry { + return &Registry{ + servers: make(map[string]*serverEntry), + } +} + +// RegisterServer 注册一个 MCP server，并初始化其生命周期状态。 +func (r *Registry) RegisterServer(serverID string, source string, version string, client ServerClient) error { + if r == nil { + return errors.New("mcp: registry is nil") + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return errors.New("mcp: server id is empty") + } + if client == nil { + return errors.New("mcp: server client is nil") + } + + r.mu.Lock() + defer r.mu.Unlock() + + if _, exists := r.servers[normalizedID]; exists { + return fmt.Errorf("mcp: server %q already exists", normalizedID) + } + r.servers[normalizedID] = &serverEntry{ + snapshot: ServerSnapshot{ + ServerID: normalizedID, + Source: strings.TrimSpace(source), + Version: strings.TrimSpace(version), + Status: ServerStatusConnecting, + UpdatedAt: time.Now(), + }, + client: client, + } + return nil +} + +// UnregisterServer 注销一个 MCP server，返回是否实际删除。 +func (r *Registry) UnregisterServer(serverID string) bool { + if r == nil { + return false + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return false + } + + r.mu.Lock() + defer r.mu.Unlock() + if _, exists := r.servers[normalizedID]; !exists { + return false + } + delete(r.servers, normalizedID) + return true +} + +// SetServerStatus 更新指定 server 的生命周期状态。 +func (r *Registry) SetServerStatus(serverID string, status ServerStatus) error { + if r == nil { + return errors.New("mcp: registry is nil") + } + if !isValidStatus(status) { + return fmt.Errorf("mcp: unsupported server status %q", status) + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return errors.New("mcp: server id is empty") + } + + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.servers[normalizedID] + if !ok { + return fmt.Errorf("mcp: server %q not found", normalizedID) + } + entry.snapshot.Status = status + entry.snapshot.UpdatedAt = time.Now() + return nil +} + +// RefreshServerTools 从 server 拉取工具清单并刷新快照。 +func (r *Registry) RefreshServerTools(ctx context.Context, serverID string) error { + if r == nil { + return errors.New("mcp: registry is nil") + } + if err := ctx.Err(); err != nil { + return err + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return errors.New("mcp: server id is empty") + } + + r.mu.RLock() + entry, ok := r.servers[normalizedID] + r.mu.RUnlock() + if !ok { + return fmt.Errorf("mcp: server %q not found", normalizedID) + } + + tools, err := entry.client.ListTools(ctx) + if err != nil { + _ = r.SetServerStatus(normalizedID, ServerStatusDegraded) + return fmt.Errorf("mcp: list tools for server %q: %w", normalizedID, err) + } + + r.mu.Lock() + defer r.mu.Unlock() + current, exists := r.servers[normalizedID] + if !exists { + return fmt.Errorf("mcp: server %q not found", normalizedID) + } + current.snapshot.Tools = cloneToolDescriptors(tools) + current.snapshot.Status = ServerStatusReady + current.snapshot.UpdatedAt = time.Now() + return nil +} + +// HealthCheck 触发指定 server 的健康探测并同步状态。 +func (r *Registry) HealthCheck(ctx context.Context, serverID string) error { + if r == nil { + return errors.New("mcp: registry is nil") + } + if err := ctx.Err(); err != nil { + return err + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return errors.New("mcp: server id is empty") + } + + r.mu.RLock() + entry, ok := r.servers[normalizedID] + r.mu.RUnlock() + if !ok { + return fmt.Errorf("mcp: server %q not found", normalizedID) + } + + if err := entry.client.HealthCheck(ctx); err != nil { + _ = r.SetServerStatus(normalizedID, ServerStatusOffline) + return fmt.Errorf("mcp: health check failed for server %q: %w", normalizedID, err) + } + return r.SetServerStatus(normalizedID, ServerStatusReady) +} + +// Call 通过 registry 分发指定 server/tool 的调用请求。 +func (r *Registry) Call(ctx context.Context, serverID string, toolName string, arguments []byte) (CallResult, error) { + if r == nil { + return CallResult{}, errors.New("mcp: registry is nil") + } + if err := ctx.Err(); err != nil { + return CallResult{}, err + } + normalizedID := normalizeServerID(serverID) + if normalizedID == "" { + return CallResult{}, errors.New("mcp: server id is empty") + } + trimmedToolName := strings.TrimSpace(toolName) + if trimmedToolName == "" { + return CallResult{}, errors.New("mcp: tool name is empty") + } + + r.mu.RLock() + entry, ok := r.servers[normalizedID] + r.mu.RUnlock() + if !ok { + return CallResult{}, fmt.Errorf("mcp: server %q not found", normalizedID) + } + + result, err := entry.client.CallTool(ctx, trimmedToolName, arguments) + if err != nil { + _ = r.SetServerStatus(normalizedID, ServerStatusDegraded) + return CallResult{}, fmt.Errorf("mcp: call %s on %s failed: %w", trimmedToolName, normalizedID, err) + } + return result, nil +} + +// Snapshot 返回当前 registry 的不可变 server 快照集合。 +func (r *Registry) Snapshot() []ServerSnapshot { + if r == nil { + return nil + } + + r.mu.RLock() + defer r.mu.RUnlock() + if len(r.servers) == 0 { + return nil + } + + keys := make([]string, 0, len(r.servers)) + for serverID := range r.servers { + keys = append(keys, serverID) + } + sort.Strings(keys) + + result := make([]ServerSnapshot, 0, len(keys)) + for _, serverID := range keys { + entry := r.servers[serverID] + snapshot := entry.snapshot + snapshot.Tools = cloneToolDescriptors(snapshot.Tools) + result = append(result, snapshot) + } + return result +} + +// normalizeServerID 统一规范化 server id 以保证匹配稳定性。 +func normalizeServerID(serverID string) string { + return strings.ToLower(strings.TrimSpace(serverID)) +} + +// isValidStatus 校验 server 状态是否属于已定义集合。 +func isValidStatus(status ServerStatus) bool { + switch status { + case ServerStatusConnecting, ServerStatusReady, ServerStatusDegraded, ServerStatusOffline: + return true + default: + return false + } +} + +// cloneToolDescriptors 深拷贝工具描述，避免快照被外部引用污染。 +func cloneToolDescriptors(input []ToolDescriptor) []ToolDescriptor { + if len(input) == 0 { + return nil + } + + result := make([]ToolDescriptor, 0, len(input)) + for _, descriptor := range input { + cloned := ToolDescriptor{ + Name: strings.TrimSpace(descriptor.Name), + Description: strings.TrimSpace(descriptor.Description), + InputSchema: cloneSchema(descriptor.InputSchema), + } + result = append(result, cloned) + } + return result +} + +// cloneSchema 深拷贝 schema 顶层 map，满足当前工具定义的只读需求。 +func cloneSchema(schema map[string]any) map[string]any { + if len(schema) == 0 { + return nil + } + cloned := make(map[string]any, len(schema)) + for key, value := range schema { + cloned[key] = cloneAny(value) + } + return cloned +} + +// cloneAny 递归复制 schema/metadata 中的 map 与 slice，避免跨层共享引用。 +func cloneAny(value any) any { + switch typed := value.(type) { + case map[string]any: + cloned := make(map[string]any, len(typed)) + for key, item := range typed { + cloned[key] = cloneAny(item) + } + return cloned + case []any: + cloned := make([]any, len(typed)) + for i, item := range typed { + cloned[i] = cloneAny(item) + } + return cloned + default: + return value + } +} diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go new file mode 100644 index 00000000..30c5cd3f --- /dev/null +++ b/internal/tools/mcp/registry_test.go @@ -0,0 +1,336 @@ +package mcp + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +type stubServerClient struct { + mu sync.Mutex + tools []ToolDescriptor + callResult CallResult + listErr error + callErr error + healthErr error + lastToolName string + lastArguments []byte +} + +func (s *stubServerClient) ListTools(ctx context.Context) ([]ToolDescriptor, error) { + if s.listErr != nil { + return nil, s.listErr + } + return cloneToolDescriptors(s.tools), nil +} + +func (s *stubServerClient) CallTool(ctx context.Context, toolName string, arguments []byte) (CallResult, error) { + s.mu.Lock() + s.lastToolName = toolName + s.lastArguments = append([]byte(nil), arguments...) + s.mu.Unlock() + if s.callErr != nil { + return CallResult{}, s.callErr + } + result := s.callResult + result.Metadata = cloneSchema(result.Metadata) + return result, nil +} + +func (s *stubServerClient) HealthCheck(ctx context.Context) error { + return s.healthErr +} + +func TestRegistryRegisterRefreshSnapshotCall(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + { + Name: "search", + Description: "search docs", + InputSchema: map[string]any{"type": "object"}, + }, + }, + callResult: CallResult{ + Content: "ok", + Metadata: map[string]any{ + "latency_ms": 18, + }, + }, + } + + if err := registry.RegisterServer("Docs", "stdio", "v1", client); err != nil { + t.Fatalf("RegisterServer() error = %v", err) + } + if err := registry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("RefreshServerTools() error = %v", err) + } + + snapshots := registry.Snapshot() + if len(snapshots) != 1 { + t.Fatalf("expected one snapshot, got %d", len(snapshots)) + } + snapshot := snapshots[0] + if snapshot.ServerID != "docs" || snapshot.Status != ServerStatusReady { + t.Fatalf("unexpected snapshot: %+v", snapshot) + } + if len(snapshot.Tools) != 1 || snapshot.Tools[0].Name != "search" { + t.Fatalf("unexpected tools in snapshot: %+v", snapshot.Tools) + } + + result, err := registry.Call(context.Background(), "docs", "search", []byte(`{"q":"mcp"}`)) + if err != nil { + t.Fatalf("Call() error = %v", err) + } + if result.Content != "ok" { + t.Fatalf("expected call content ok, got %q", result.Content) + } +} + +func TestRegistryStatusTransitions(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + {Name: "search", InputSchema: map[string]any{"type": "object"}}, + }, + } + if err := registry.RegisterServer("server-1", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + + client.healthErr = errors.New("offline") + if err := registry.HealthCheck(context.Background(), "server-1"); err == nil { + t.Fatalf("expected health check failure") + } + if snapshots := registry.Snapshot(); snapshots[0].Status != ServerStatusOffline { + t.Fatalf("expected offline status, got %+v", snapshots[0].Status) + } + + client.healthErr = nil + if err := registry.HealthCheck(context.Background(), "server-1"); err != nil { + t.Fatalf("unexpected health check error: %v", err) + } + if snapshots := registry.Snapshot(); snapshots[0].Status != ServerStatusReady { + t.Fatalf("expected ready status, got %+v", snapshots[0].Status) + } +} + +func TestRegistryConcurrentSnapshotAndRefresh(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + {Name: "search", InputSchema: map[string]any{"type": "object"}}, + }, + } + if err := registry.RegisterServer("server-1", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = registry.RefreshServerTools(context.Background(), "server-1") + }() + } + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = registry.Snapshot() + }() + } + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatalf("concurrent registry operations timed out") + } +} + +func TestRegistrySnapshotSchemaIsDeepCloned(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + { + Name: "search", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]any{ + "type": "string", + }, + }, + }, + }, + }, + } + if err := registry.RegisterServer("server-1", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.RefreshServerTools(context.Background(), "server-1"); err != nil { + t.Fatalf("refresh tools: %v", err) + } + + first := registry.Snapshot() + properties, ok := first[0].Tools[0].InputSchema["properties"].(map[string]any) + if !ok { + t.Fatalf("expected properties map") + } + properties["query"] = map[string]any{"type": "number"} + + second := registry.Snapshot() + secondProperties, ok := second[0].Tools[0].InputSchema["properties"].(map[string]any) + if !ok { + t.Fatalf("expected properties map in second snapshot") + } + query, ok := secondProperties["query"].(map[string]any) + if !ok { + t.Fatalf("expected query schema map") + } + if query["type"] != "string" { + t.Fatalf("expected deep cloned schema type string, got %v", query["type"]) + } +} + +func TestRegistryRegisterAndUnregisterBoundaries(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{} + + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.RegisterServer("docs", "stdio", "v1", client); err == nil { + t.Fatalf("expected duplicate register error") + } + if !registry.UnregisterServer("docs") { + t.Fatalf("expected unregister success") + } + if registry.UnregisterServer("docs") { + t.Fatalf("expected unregister miss to be false") + } +} + +func TestRegistrySetServerStatusValidation(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{} + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + if err := registry.SetServerStatus("docs", ServerStatus("unknown")); err == nil { + t.Fatalf("expected invalid status error") + } + if err := registry.SetServerStatus("missing", ServerStatusReady); err == nil { + t.Fatalf("expected missing server error") + } +} + +func TestRegistryNilAndValidationBoundaries(t *testing.T) { + t.Parallel() + + var nilRegistry *Registry + if err := nilRegistry.RegisterServer("docs", "stdio", "v1", &stubServerClient{}); err == nil { + t.Fatalf("expected nil registry error") + } + if nilRegistry.UnregisterServer("docs") { + t.Fatalf("nil registry should return false on unregister") + } + if err := nilRegistry.SetServerStatus("docs", ServerStatusReady); err == nil { + t.Fatalf("expected nil registry error for set status") + } + if err := nilRegistry.RefreshServerTools(context.Background(), "docs"); err == nil { + t.Fatalf("expected nil registry error for refresh") + } + if err := nilRegistry.HealthCheck(context.Background(), "docs"); err == nil { + t.Fatalf("expected nil registry error for health check") + } + if _, err := nilRegistry.Call(context.Background(), "docs", "search", nil); err == nil { + t.Fatalf("expected nil registry error for call") + } + if snapshots := nilRegistry.Snapshot(); snapshots != nil { + t.Fatalf("expected nil snapshots from nil registry") + } +} + +func TestRegistryRefreshHealthCallValidation(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{} + if err := registry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register server: %v", err) + } + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + if err := registry.RefreshServerTools(canceledCtx, "docs"); err == nil { + t.Fatalf("expected canceled refresh error") + } + if err := registry.HealthCheck(canceledCtx, "docs"); err == nil { + t.Fatalf("expected canceled health check error") + } + if _, err := registry.Call(canceledCtx, "docs", "search", nil); err == nil { + t.Fatalf("expected canceled call error") + } + + if err := registry.RefreshServerTools(context.Background(), " "); err == nil { + t.Fatalf("expected empty server id error") + } + if err := registry.HealthCheck(context.Background(), " "); err == nil { + t.Fatalf("expected empty server id error") + } + if _, err := registry.Call(context.Background(), " ", "search", nil); err == nil { + t.Fatalf("expected empty server id error") + } + if _, err := registry.Call(context.Background(), "docs", " ", nil); err == nil { + t.Fatalf("expected empty tool name error") + } +} + +func TestRegistryCloneAnyCoversSlicesAndMaps(t *testing.T) { + t.Parallel() + + source := map[string]any{ + "items": []any{ + map[string]any{"name": "a"}, + []any{"nested"}, + }, + } + cloned := cloneSchema(source) + + items, ok := cloned["items"].([]any) + if !ok { + t.Fatalf("expected []any clone") + } + nestedMap, ok := items[0].(map[string]any) + if !ok { + t.Fatalf("expected nested map clone") + } + nestedMap["name"] = "changed" + + originalItems := source["items"].([]any) + originalMap := originalItems[0].(map[string]any) + if originalMap["name"] != "a" { + t.Fatalf("expected deep cloned map, got %v", originalMap["name"]) + } +} diff --git a/internal/tools/mcp/stdio_client.go b/internal/tools/mcp/stdio_client.go new file mode 100644 index 00000000..b6808d51 --- /dev/null +++ b/internal/tools/mcp/stdio_client.go @@ -0,0 +1,853 @@ +package mcp + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + defaultStdioStartTimeout = 5 * time.Second + defaultStdioCallTimeout = 15 * time.Second + defaultStdioRestartBackoff = 1 * time.Second + maxStdioRestartBackoff = 30 * time.Second + maxStdioFrameBytes = 8 * 1024 * 1024 + maxStdioLineBytes = 8 * 1024 * 1024 + defaultMCPProtocolVersion = "2024-11-05" + defaultMCPClientName = "neocode" + defaultMCPClientVersion = "0.1.0" +) + +// StdioClientConfig 描述 MCP stdio 客户端的启动与调用参数。 +type StdioClientConfig struct { + Command string + Args []string + Env []string + Workdir string + StartTimeout time.Duration + CallTimeout time.Duration + RestartBackoff time.Duration +} + +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID string `json:"id"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + +type jsonRPCNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID string `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` +} + +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type rpcReply struct { + result json.RawMessage + err error +} + +// StdIOClient 通过 stdio 子进程与 MCP server 进行 JSON-RPC 通信。 +type StdIOClient struct { + cfg StdioClientConfig + idSeed uint64 + mu sync.Mutex + writeMu sync.Mutex + cmd *exec.Cmd + stdin io.WriteCloser + stdout io.ReadCloser + reader *bufio.Reader + pending map[string]chan rpcReply + exited chan struct{} + exitErr error + backoff time.Duration + retryAt time.Time + started bool + initialized bool + initializing bool + initDone chan struct{} + protocol stdioProtocol + shutdown bool +} + +type stdioProtocol string + +const ( + stdioProtocolUnknown stdioProtocol = "" + stdioProtocolLine stdioProtocol = "line" + stdioProtocolFramed stdioProtocol = "framed" +) + +// NewStdIOClient 创建 stdio MCP client。 +func NewStdIOClient(cfg StdioClientConfig) (*StdIOClient, error) { + if strings.TrimSpace(cfg.Command) == "" { + return nil, errors.New("mcp: stdio command is empty") + } + if cfg.StartTimeout <= 0 { + cfg.StartTimeout = defaultStdioStartTimeout + } + if cfg.CallTimeout <= 0 { + cfg.CallTimeout = defaultStdioCallTimeout + } + if cfg.RestartBackoff <= 0 { + cfg.RestartBackoff = defaultStdioRestartBackoff + } + + return &StdIOClient{ + cfg: cfg, + pending: make(map[string]chan rpcReply), + backoff: cfg.RestartBackoff, + }, nil +} + +// Close 关闭 stdio 子进程并释放资源。 +func (c *StdIOClient) Close() error { + if c == nil { + return nil + } + + c.mu.Lock() + defer c.mu.Unlock() + c.shutdown = true + if c.stdin != nil { + _ = c.stdin.Close() + } + if c.cmd != nil && c.cmd.Process != nil { + _ = c.cmd.Process.Kill() + } + c.failAllPendingLocked(errors.New("mcp: stdio client closed")) + return nil +} + +// ListTools 调用 MCP tools/list 获取工具清单。 +func (c *StdIOClient) ListTools(ctx context.Context) ([]ToolDescriptor, error) { + callCtx, cancel := c.callContext(ctx) + defer cancel() + + raw, err := c.call(callCtx, "tools/list", map[string]any{}) + if err != nil { + return nil, err + } + + var payload struct { + Tools []struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema map[string]any `json:"inputSchema"` + InputSchema2 map[string]any `json:"input_schema"` + } `json:"tools"` + } + if err := json.Unmarshal(raw, &payload); err != nil { + return nil, fmt.Errorf("mcp: decode tools/list result: %w", err) + } + + result := make([]ToolDescriptor, 0, len(payload.Tools)) + for _, item := range payload.Tools { + schema := item.InputSchema + if len(schema) == 0 { + schema = item.InputSchema2 + } + result = append(result, ToolDescriptor{ + Name: strings.TrimSpace(item.Name), + Description: strings.TrimSpace(item.Description), + InputSchema: ensureObjectSchema(schema), + }) + } + return result, nil +} + +// CallTool 调用 MCP tools/call 并收敛返回值。 +func (c *StdIOClient) CallTool(ctx context.Context, toolName string, arguments []byte) (CallResult, error) { + trimmedToolName := strings.TrimSpace(toolName) + if trimmedToolName == "" { + return CallResult{}, errors.New("mcp: tool name is empty") + } + + callCtx, cancel := c.callContext(ctx) + defer cancel() + + var args any = map[string]any{} + if len(arguments) > 0 { + if err := json.Unmarshal(arguments, &args); err != nil { + return CallResult{}, fmt.Errorf("mcp: decode tool arguments: %w", err) + } + } + + raw, err := c.call(callCtx, "tools/call", map[string]any{ + "name": trimmedToolName, + "arguments": args, + }) + if err != nil { + return CallResult{}, err + } + return decodeCallResult(raw), nil +} + +// HealthCheck 通过一次短超时 tools/list 验证连接可用性。 +func (c *StdIOClient) HealthCheck(ctx context.Context) error { + _, err := c.ListTools(ctx) + return err +} + +// callContext 基于配置与上游截止时间生成单次 RPC 调用上下文。 +func (c *StdIOClient) callContext(ctx context.Context) (context.Context, context.CancelFunc) { + timeout := c.cfg.CallTimeout + if deadline, ok := ctx.Deadline(); ok { + if remaining := time.Until(deadline); remaining > 0 && remaining < timeout { + timeout = remaining + } + } + return context.WithTimeout(ctx, timeout) +} + +// call 发送 RPC 请求并等待响应。 +func (c *StdIOClient) call(ctx context.Context, method string, params any) (json.RawMessage, error) { + return c.callRequest(ctx, method, params, false) +} + +// callRequest 在可用连接上发送 RPC 请求，skipEnsure=true 用于初始化阶段避免递归。 +func (c *StdIOClient) callRequest(ctx context.Context, method string, params any, skipEnsure bool) (json.RawMessage, error) { + return c.callRequestWithProtocol(ctx, method, params, skipEnsure, stdioProtocolUnknown) +} + +// callRequestWithProtocol 按指定协议发送 RPC 请求，用于 initialize 时的协议探测与回退。 +func (c *StdIOClient) callRequestWithProtocol( + ctx context.Context, + method string, + params any, + skipEnsure bool, + override stdioProtocol, +) (json.RawMessage, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if !skipEnsure { + if err := c.ensureStarted(ctx); err != nil { + return nil, err + } + if err := c.ensureInitialized(ctx); err != nil { + return nil, err + } + } + + requestID := "req-" + strconv.FormatUint(atomic.AddUint64(&c.idSeed, 1), 10) + replyCh := make(chan rpcReply, 1) + + c.mu.Lock() + if c.shutdown { + c.mu.Unlock() + return nil, errors.New("mcp: stdio client closed") + } + c.pending[requestID] = replyCh + stdin := c.stdin + selectedProtocol := c.resolveWriteProtocolLocked(override) + c.mu.Unlock() + if stdin == nil { + c.removePending(requestID) + return nil, errors.New("mcp: stdio client is not connected") + } + + requestPayload, err := json.Marshal(jsonRPCRequest{ + JSONRPC: "2.0", + ID: requestID, + Method: method, + Params: params, + }) + if err != nil { + c.removePending(requestID) + return nil, fmt.Errorf("mcp: marshal request: %w", err) + } + c.writeMu.Lock() + writeErr := writeMessageWithProtocol(stdin, requestPayload, selectedProtocol) + c.writeMu.Unlock() + if writeErr != nil { + c.removePending(requestID) + return nil, fmt.Errorf("mcp: send request: %w", writeErr) + } + + select { + case <-ctx.Done(): + c.removePending(requestID) + return nil, ctx.Err() + case reply := <-replyCh: + return reply.result, reply.err + } +} + +// sendNotification 发送无响应 RPC 通知，skipEnsure=true 用于初始化流程。 +func (c *StdIOClient) sendNotification(ctx context.Context, method string, params any, skipEnsure bool) error { + return c.sendNotificationWithProtocol(ctx, method, params, skipEnsure, stdioProtocolUnknown) +} + +// sendNotificationWithProtocol 按指定协议发送 RPC 通知，用于 initialize 完成后的 initialized 事件。 +func (c *StdIOClient) sendNotificationWithProtocol( + ctx context.Context, + method string, + params any, + skipEnsure bool, + override stdioProtocol, +) error { + if err := ctx.Err(); err != nil { + return err + } + if !skipEnsure { + if err := c.ensureStarted(ctx); err != nil { + return err + } + if err := c.ensureInitialized(ctx); err != nil { + return err + } + } + + c.mu.Lock() + if c.shutdown { + c.mu.Unlock() + return errors.New("mcp: stdio client closed") + } + stdin := c.stdin + selectedProtocol := c.resolveWriteProtocolLocked(override) + c.mu.Unlock() + if stdin == nil { + return errors.New("mcp: stdio client is not connected") + } + + payload, err := json.Marshal(jsonRPCNotification{ + JSONRPC: "2.0", + Method: method, + Params: params, + }) + if err != nil { + return fmt.Errorf("mcp: marshal notification: %w", err) + } + + c.writeMu.Lock() + writeErr := writeMessageWithProtocol(stdin, payload, selectedProtocol) + c.writeMu.Unlock() + if writeErr != nil { + return fmt.Errorf("mcp: send notification: %w", writeErr) + } + return nil +} + +// ensureStarted 确保 stdio 子进程已启动并处于可读写状态。 +func (c *StdIOClient) ensureStarted(ctx context.Context) error { + c.mu.Lock() + defer c.mu.Unlock() + + if c.shutdown { + return errors.New("mcp: stdio client closed") + } + if c.started { + return nil + } + if !c.retryAt.IsZero() && time.Now().Before(c.retryAt) { + return fmt.Errorf("mcp: stdio restart backoff in effect until %s", c.retryAt.Format(time.RFC3339)) + } + + startCtx, cancel := context.WithTimeout(ctx, c.cfg.StartTimeout) + defer cancel() + + command := exec.Command(c.cfg.Command, c.cfg.Args...) + command.Env = append(os.Environ(), c.cfg.Env...) + command.Dir = strings.TrimSpace(c.cfg.Workdir) + + stdin, err := command.StdinPipe() + if err != nil { + return fmt.Errorf("mcp: create stdin pipe: %w", err) + } + stdout, err := command.StdoutPipe() + if err != nil { + return fmt.Errorf("mcp: create stdout pipe: %w", err) + } + stderr, err := command.StderrPipe() + if err != nil { + return fmt.Errorf("mcp: create stderr pipe: %w", err) + } + + startErrCh := make(chan error, 1) + go func() { + startErrCh <- command.Start() + }() + select { + case <-startCtx.Done(): + return startCtx.Err() + case err := <-startErrCh: + if err != nil { + c.bumpBackoffLocked() + return fmt.Errorf("mcp: start stdio server: %w", err) + } + } + + c.cmd = command + c.stdin = stdin + c.stdout = stdout + c.reader = bufio.NewReader(stdout) + c.exited = make(chan struct{}) + c.exitErr = nil + c.started = true + c.initialized = false + c.initializing = false + c.initDone = nil + c.protocol = stdioProtocolUnknown + c.backoff = c.cfg.RestartBackoff + c.retryAt = time.Time{} + + go c.readLoop() + go c.waitLoop(command) + go io.Copy(io.Discard, stderr) + return nil +} + +// ensureInitialized 确保 MCP 会话完成 initialize/initialized 握手，并发调用共享结果。 +func (c *StdIOClient) ensureInitialized(ctx context.Context) error { + for { + c.mu.Lock() + if c.shutdown { + c.mu.Unlock() + return errors.New("mcp: stdio client closed") + } + if !c.started { + c.mu.Unlock() + return errors.New("mcp: stdio client is not started") + } + if c.initialized { + c.mu.Unlock() + return nil + } + if c.initializing { + wait := c.initDone + c.mu.Unlock() + select { + case <-ctx.Done(): + return ctx.Err() + case <-wait: + continue + } + } + c.initializing = true + c.initDone = make(chan struct{}) + done := c.initDone + c.mu.Unlock() + + initErr := c.performInitialize(ctx) + + c.mu.Lock() + if c.started && initErr == nil { + c.initialized = true + } + c.initializing = false + close(done) + c.mu.Unlock() + return initErr + } +} + +// performInitialize 执行 MCP 握手并自动兼容 line/framed 两种 stdio 线协议。 +func (c *StdIOClient) performInitialize(ctx context.Context) error { + params := map[string]any{ + "protocolVersion": defaultMCPProtocolVersion, + "capabilities": map[string]any{}, + "clientInfo": map[string]any{ + "name": defaultMCPClientName, + "version": defaultMCPClientVersion, + }, + } + + protocols := []stdioProtocol{stdioProtocolLine, stdioProtocolFramed} + c.mu.Lock() + if c.protocol == stdioProtocolLine || c.protocol == stdioProtocolFramed { + protocols = []stdioProtocol{c.protocol} + } + c.mu.Unlock() + + errs := make([]string, 0, len(protocols)) + for index, protocol := range protocols { + attemptCtx, cancel := initializeAttemptContext(ctx, len(protocols)-index) + _, initErr := c.callRequestWithProtocol(attemptCtx, "initialize", params, true, protocol) + cancel() + if initErr != nil { + errs = append(errs, fmt.Sprintf("%s=%v", protocol, initErr)) + continue + } + + notifyCtx, notifyCancel := initializeAttemptContext(ctx, len(protocols)-index) + notifyErr := c.sendNotificationWithProtocol( + notifyCtx, + "notifications/initialized", + map[string]any{}, + true, + protocol, + ) + notifyCancel() + if notifyErr != nil { + errs = append(errs, fmt.Sprintf("%s=%v", protocol, notifyErr)) + continue + } + + c.mu.Lock() + c.protocol = protocol + c.mu.Unlock() + return nil + } + + return fmt.Errorf("mcp: initialize session: %s", strings.Join(errs, "; ")) +} + +// readLoop 持续消费 MCP server 响应并分发给对应 pending 请求。 +func (c *StdIOClient) readLoop() { + for { + message, protocol, err := readRPCMessage(c.reader) + if err != nil { + c.markExited(fmt.Errorf("mcp: read response: %w", err)) + return + } + + c.mu.Lock() + if c.protocol == stdioProtocolUnknown && (protocol == stdioProtocolLine || protocol == stdioProtocolFramed) { + c.protocol = protocol + } + c.mu.Unlock() + + var response jsonRPCResponse + if err := json.Unmarshal(message, &response); err != nil { + continue + } + if strings.TrimSpace(response.ID) == "" { + continue + } + + c.mu.Lock() + replyCh, ok := c.pending[response.ID] + if ok { + delete(c.pending, response.ID) + } + c.mu.Unlock() + if !ok { + continue + } + + if response.Error != nil { + replyCh <- rpcReply{ + err: fmt.Errorf("mcp: rpc error %d: %s", response.Error.Code, strings.TrimSpace(response.Error.Message)), + } + continue + } + replyCh <- rpcReply{result: response.Result} + } +} + +// waitLoop 等待子进程退出并触发统一下线处理。 +func (c *StdIOClient) waitLoop(command *exec.Cmd) { + if command == nil { + c.markExited(errors.New("mcp: stdio process is nil")) + return + } + err := command.Wait() + c.markExited(fmt.Errorf("mcp: stdio process exited: %w", err)) +} + +// markExited 将客户端状态原子切换为已下线并唤醒所有等待请求。 +func (c *StdIOClient) markExited(err error) { + c.mu.Lock() + defer c.mu.Unlock() + + if !c.started { + return + } + c.started = false + c.initialized = false + if c.initializing && c.initDone != nil { + close(c.initDone) + } + c.initializing = false + c.initDone = nil + c.protocol = stdioProtocolUnknown + c.exitErr = err + if c.exited != nil { + close(c.exited) + } + c.stdin = nil + c.stdout = nil + c.reader = nil + c.cmd = nil + c.failAllPendingLocked(err) + c.bumpBackoffLocked() +} + +// resolveWriteProtocolLocked 根据调用方覆盖值与会话状态解析实际写入协议。 +func (c *StdIOClient) resolveWriteProtocolLocked(override stdioProtocol) stdioProtocol { + if override == stdioProtocolLine || override == stdioProtocolFramed { + return override + } + if c.protocol == stdioProtocolLine || c.protocol == stdioProtocolFramed { + return c.protocol + } + return stdioProtocolFramed +} + +// removePending 从挂起请求表中移除指定 requestID。 +func (c *StdIOClient) removePending(requestID string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.pending, requestID) +} + +// failAllPendingLocked 将所有挂起请求统一返回错误，调用方需持有 c.mu。 +func (c *StdIOClient) failAllPendingLocked(err error) { + for requestID, replyCh := range c.pending { + replyCh <- rpcReply{err: err} + delete(c.pending, requestID) + } +} + +// bumpBackoffLocked 按指数退避策略更新下次可重启时间，调用方需持有 c.mu。 +func (c *StdIOClient) bumpBackoffLocked() { + if c.backoff <= 0 { + c.backoff = c.cfg.RestartBackoff + } + c.retryAt = time.Now().Add(c.backoff) + c.backoff *= 2 + if c.backoff > maxStdioRestartBackoff { + c.backoff = maxStdioRestartBackoff + } +} + +// initializeAttemptContext 按剩余尝试次数切分超时预算，避免单个协议探测耗尽全部时间。 +func initializeAttemptContext(parent context.Context, remainingAttempts int) (context.Context, context.CancelFunc) { + if remainingAttempts <= 1 { + return context.WithCancel(parent) + } + deadline, ok := parent.Deadline() + if !ok { + return context.WithCancel(parent) + } + remaining := time.Until(deadline) + if remaining <= 0 { + return context.WithCancel(parent) + } + timeout := remaining / time.Duration(remainingAttempts) + if timeout < 200*time.Millisecond { + timeout = 200 * time.Millisecond + } + return context.WithTimeout(parent, timeout) +} + +// writeFramedMessage 以 Content-Length framed 格式写入 JSON-RPC 消息。 +func writeFramedMessage(writer io.Writer, payload []byte) error { + header := fmt.Sprintf("Content-Length: %d\r\n\r\n", len(payload)) + if _, err := io.WriteString(writer, header); err != nil { + return err + } + if _, err := writer.Write(payload); err != nil { + return err + } + return nil +} + +// writeLineMessage 以行分隔 JSON（NDJSON）格式写入 JSON-RPC 消息。 +func writeLineMessage(writer io.Writer, payload []byte) error { + if len(payload) == 0 { + return nil + } + if _, err := writer.Write(payload); err != nil { + return err + } + if !bytes.HasSuffix(payload, []byte("\n")) { + if _, err := io.WriteString(writer, "\n"); err != nil { + return err + } + } + return nil +} + +// writeMessageWithProtocol 按指定线协议写入消息；未知协议默认使用 framed。 +func writeMessageWithProtocol(writer io.Writer, payload []byte, protocol stdioProtocol) error { + switch protocol { + case stdioProtocolLine: + return writeLineMessage(writer, payload) + case stdioProtocolFramed, stdioProtocolUnknown: + return writeFramedMessage(writer, payload) + default: + return writeFramedMessage(writer, payload) + } +} + +// readRPCMessage 自动识别并读取 line/framed 两种 stdio 消息，忽略非协议日志行。 +func readRPCMessage(reader *bufio.Reader) ([]byte, stdioProtocol, error) { + for { + line, err := reader.ReadString('\n') + if err != nil { + return nil, stdioProtocolUnknown, err + } + + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + + lower := strings.ToLower(trimmed) + if strings.HasPrefix(lower, "content-length:") { + payload, framedErr := readFramedPayload(reader, trimmed) + if framedErr != nil { + return nil, stdioProtocolUnknown, framedErr + } + return payload, stdioProtocolFramed, nil + } + + if strings.HasPrefix(trimmed, "{") || strings.HasPrefix(trimmed, "[") { + if len(trimmed) > maxStdioLineBytes { + return nil, stdioProtocolUnknown, fmt.Errorf("mcp: line message too large: %d", len(trimmed)) + } + if !json.Valid([]byte(trimmed)) { + continue + } + return []byte(trimmed), stdioProtocolLine, nil + } + } +} + +// readFramedPayload 在已读取首个 Content-Length 头后，继续读取 header/body 并返回 payload。 +func readFramedPayload(reader *bufio.Reader, firstHeader string) ([]byte, error) { + contentLength, err := parseContentLength(firstHeader) + if err != nil { + return nil, err + } + + for { + line, readErr := reader.ReadString('\n') + if readErr != nil { + return nil, readErr + } + if strings.TrimSpace(line) == "" { + break + } + } + + if contentLength < 0 { + return nil, errors.New("mcp: missing content-length header") + } + if contentLength > maxStdioFrameBytes { + return nil, fmt.Errorf("mcp: content-length %d exceeds limit %d", contentLength, maxStdioFrameBytes) + } + + payload := make([]byte, contentLength) + if _, err := io.ReadFull(reader, payload); err != nil { + return nil, err + } + return payload, nil +} + +// parseContentLength 解析 Content-Length 头并返回消息体长度。 +func parseContentLength(header string) (int, error) { + trimmed := strings.TrimSpace(header) + lower := strings.ToLower(trimmed) + if !strings.HasPrefix(lower, "content-length:") { + return -1, errors.New("mcp: missing content-length header") + } + rawLength := strings.TrimSpace(trimmed[len("content-length:"):]) + contentLength, err := strconv.Atoi(rawLength) + if err != nil { + return -1, fmt.Errorf("mcp: invalid content-length %q", rawLength) + } + return contentLength, nil +} + +// readFramedMessage 仅读取 framed 消息；会跳过前置日志与 line 消息直到遇到 framed。 +func readFramedMessage(reader *bufio.Reader) ([]byte, error) { + for { + message, protocol, err := readRPCMessage(reader) + if err != nil { + return nil, err + } + if protocol == stdioProtocolFramed { + return message, nil + } + } +} + +// decodeCallResult 将 tools/call 结果统一收敛为 CallResult。 +func decodeCallResult(raw json.RawMessage) CallResult { + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + return CallResult{ + Content: strings.TrimSpace(string(raw)), + IsError: false, + Metadata: map[string]any{"raw_result": string(raw)}, + } + } + + content := "" + switch typed := payload["content"].(type) { + case string: + content = strings.TrimSpace(typed) + case []any: + lines := make([]string, 0, len(typed)) + for _, item := range typed { + switch value := item.(type) { + case map[string]any: + text, _ := value["text"].(string) + if strings.TrimSpace(text) != "" { + lines = append(lines, strings.TrimSpace(text)) + } + case string: + if strings.TrimSpace(value) != "" { + lines = append(lines, strings.TrimSpace(value)) + } + } + } + content = strings.Join(lines, "\n") + default: + if typed != nil { + content = strings.TrimSpace(fmt.Sprintf("%v", typed)) + } + } + if content == "" { + content = "ok" + } + + isError := false + if value, ok := payload["isError"].(bool); ok { + isError = value + } + if value, ok := payload["is_error"].(bool); ok { + isError = isError || value + } + + metadata := map[string]any{} + for key, value := range payload { + if key == "content" || key == "isError" || key == "is_error" { + continue + } + metadata[key] = value + } + metadata["raw_result"] = string(bytes.TrimSpace(raw)) + + return CallResult{ + Content: content, + IsError: isError, + Metadata: metadata, + } +} diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go new file mode 100644 index 00000000..01d8d20b --- /dev/null +++ b/internal/tools/mcp/stdio_client_test.go @@ -0,0 +1,615 @@ +package mcp + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "strconv" + "strings" + "sync" + "testing" + "time" +) + +type nopWriteCloser struct { + bytes.Buffer +} + +func (n *nopWriteCloser) Close() error { return nil } + +type errWriter struct{} + +func (errWriter) Write(p []byte) (int, error) { + return 0, errors.New("write failed") +} + +func TestStdIOClientListToolsAndCallTool(t *testing.T) { + t.Parallel() + + client := newTestStdIOClientWithMode(t, "framed") + defer func() { _ = client.Close() }() + + toolsList, err := client.ListTools(context.Background()) + if err != nil { + t.Fatalf("ListTools() error = %v", err) + } + if len(toolsList) != 1 || toolsList[0].Name != "search" { + t.Fatalf("unexpected tools list: %+v", toolsList) + } + + result, err := client.CallTool(context.Background(), "search", []byte(`{"query":"mcp"}`)) + if err != nil { + t.Fatalf("CallTool() error = %v", err) + } + if !strings.Contains(result.Content, "search") { + t.Fatalf("unexpected call result content: %q", result.Content) + } +} + +func TestStdIOClientLineProtocolInitializeAndCalls(t *testing.T) { + t.Parallel() + + client := newTestStdIOClientWithMode(t, "line") + defer func() { _ = client.Close() }() + + toolsList, err := client.ListTools(context.Background()) + if err != nil { + t.Fatalf("ListTools() with line protocol error = %v", err) + } + if len(toolsList) != 1 || toolsList[0].Name != "search" { + t.Fatalf("unexpected tools list: %+v", toolsList) + } + + result, err := client.CallTool(context.Background(), "search", []byte(`{"query":"mcp"}`)) + if err != nil { + t.Fatalf("CallTool() with line protocol error = %v", err) + } + if !strings.Contains(result.Content, "search") { + t.Fatalf("unexpected call result content: %q", result.Content) + } +} + +func TestStdIOClientInitializeFallbackToFramed(t *testing.T) { + t.Parallel() + + client := newTestStdIOClientWithMode(t, "framed") + defer func() { _ = client.Close() }() + + if _, err := client.ListTools(context.Background()); err != nil { + t.Fatalf("expected fallback initialize success, got %v", err) + } + if client.protocol != stdioProtocolFramed { + t.Fatalf("expected framed protocol selected, got %q", client.protocol) + } +} + +func TestStdIOClientHealthCheck(t *testing.T) { + t.Parallel() + + client := newTestStdIOClientWithMode(t, "framed") + defer func() { _ = client.Close() }() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + if err := client.HealthCheck(ctx); err != nil { + t.Fatalf("HealthCheck() error = %v", err) + } +} + +func TestStdIOClientSendNotificationFramedDefault(t *testing.T) { + t.Parallel() + + writer := &nopWriteCloser{} + client := &StdIOClient{ + pending: make(map[string]chan rpcReply), + stdin: writer, + started: true, + cfg: StdioClientConfig{ + CallTimeout: time.Second, + StartTimeout: time.Second, + RestartBackoff: time.Millisecond, + }, + } + + err := client.sendNotification(context.Background(), "notifications/initialized", map[string]any{}, true) + if err != nil { + t.Fatalf("sendNotification() error = %v", err) + } + if !strings.Contains(writer.String(), "Content-Length:") { + t.Fatalf("expected framed header, got: %q", writer.String()) + } +} + +func TestStdIOClientSendNotificationLineProtocol(t *testing.T) { + t.Parallel() + + writer := &nopWriteCloser{} + client := &StdIOClient{ + pending: make(map[string]chan rpcReply), + stdin: writer, + started: true, + protocol: stdioProtocolLine, + cfg: StdioClientConfig{ + CallTimeout: time.Second, + StartTimeout: time.Second, + RestartBackoff: time.Millisecond, + }, + } + + err := client.sendNotificationWithProtocol( + context.Background(), + "notifications/initialized", + map[string]any{}, + true, + stdioProtocolLine, + ) + if err != nil { + t.Fatalf("sendNotificationWithProtocol() error = %v", err) + } + raw := writer.String() + if strings.Contains(raw, "Content-Length:") { + t.Fatalf("expected line protocol write, got: %q", raw) + } + if !strings.HasSuffix(raw, "\n") { + t.Fatalf("expected newline-delimited payload, got: %q", raw) + } +} + +func TestStdIOClientConcurrentCallTool(t *testing.T) { + t.Parallel() + + client := newTestStdIOClientWithMode(t, "framed") + defer func() { _ = client.Close() }() + + const workers = 16 + var wg sync.WaitGroup + errCh := make(chan error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result, err := client.CallTool(context.Background(), "search", []byte(`{"query":"mcp"}`)) + if err != nil { + errCh <- err + return + } + if !strings.Contains(result.Content, "search") { + errCh <- fmt.Errorf("unexpected content: %q", result.Content) + } + }() + } + wg.Wait() + close(errCh) + for err := range errCh { + t.Fatalf("concurrent call failed: %v", err) + } +} + +func TestReadFramedMessageRejectsOversizedPayload(t *testing.T) { + t.Parallel() + + payload := strings.Repeat("x", 32) + raw := fmt.Sprintf("Content-Length: %d\r\n\r\n%s", maxStdioFrameBytes+1, payload) + reader := bufio.NewReader(strings.NewReader(raw)) + _, err := readFramedMessage(reader) + if err == nil { + t.Fatalf("expected oversized payload error") + } + if !strings.Contains(err.Error(), "exceeds limit") { + t.Fatalf("expected exceeds limit error, got %v", err) + } +} + +func TestReadRPCMessageLine(t *testing.T) { + t.Parallel() + + reader := bufio.NewReader(strings.NewReader("Starting...\n{\"jsonrpc\":\"2.0\",\"id\":\"1\",\"result\":{}}\n")) + payload, protocol, err := readRPCMessage(reader) + if err != nil { + t.Fatalf("readRPCMessage() error = %v", err) + } + if protocol != stdioProtocolLine { + t.Fatalf("expected line protocol, got %q", protocol) + } + if !strings.Contains(string(payload), `"jsonrpc":"2.0"`) { + t.Fatalf("unexpected payload: %s", payload) + } +} + +func TestReadRPCMessageFramed(t *testing.T) { + t.Parallel() + + body := `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + raw := "log\nContent-Length: " + strconv.Itoa(len(body)) + "\r\n\r\n" + body + reader := bufio.NewReader(strings.NewReader(raw)) + payload, protocol, err := readRPCMessage(reader) + if err != nil { + t.Fatalf("readRPCMessage() error = %v", err) + } + if protocol != stdioProtocolFramed { + t.Fatalf("expected framed protocol, got %q", protocol) + } + if string(payload) != body { + t.Fatalf("unexpected payload: %s", payload) + } +} + +func TestNewStdIOClientValidationAndDefaults(t *testing.T) { + t.Parallel() + + if _, err := NewStdIOClient(StdioClientConfig{}); err == nil { + t.Fatalf("expected empty command error") + } + client, err := NewStdIOClient(StdioClientConfig{Command: "cmd"}) + if err != nil { + t.Fatalf("NewStdIOClient() error = %v", err) + } + if client.cfg.StartTimeout <= 0 || client.cfg.CallTimeout <= 0 || client.cfg.RestartBackoff <= 0 { + t.Fatalf("expected default timeouts/backoff to be initialized") + } +} + +func TestStdIOClientCallToolInputValidation(t *testing.T) { + t.Parallel() + + client := &StdIOClient{} + if _, err := client.CallTool(context.Background(), "", nil); err == nil { + t.Fatalf("expected empty tool name error") + } + if _, err := client.CallTool(context.Background(), "search", []byte("{not-json")); err == nil { + t.Fatalf("expected invalid json arguments error") + } +} + +func TestStdIOClientCallRejectsClosedAndDisconnected(t *testing.T) { + t.Parallel() + + client := &StdIOClient{ + pending: make(map[string]chan rpcReply), + cfg: StdioClientConfig{ + CallTimeout: time.Second, + StartTimeout: time.Second, + RestartBackoff: time.Millisecond, + }, + } + client.shutdown = true + if _, err := client.call(context.Background(), "tools/list", map[string]any{}); err == nil { + t.Fatalf("expected closed error") + } + + client.shutdown = false + client.started = true + client.stdin = nil + if _, err := client.call(context.Background(), "tools/list", map[string]any{}); err == nil { + t.Fatalf("expected disconnected error") + } +} + +func TestStdIOClientEnsureStartedBackoff(t *testing.T) { + t.Parallel() + + client := &StdIOClient{ + cfg: StdioClientConfig{ + Command: "cmd", + StartTimeout: time.Second, + CallTimeout: time.Second, + RestartBackoff: time.Second, + }, + pending: make(map[string]chan rpcReply), + retryAt: time.Now().Add(2 * time.Second), + } + + err := client.ensureStarted(context.Background()) + if err == nil || !strings.Contains(err.Error(), "backoff") { + t.Fatalf("expected backoff error, got %v", err) + } +} + +func TestReadFramedMessageHeaderErrors(t *testing.T) { + t.Parallel() + + reader := bufio.NewReader(strings.NewReader("X-Test: 1\r\n\r\n{}")) + if _, err := readFramedMessage(reader); err == nil || !(strings.Contains(err.Error(), "missing content-length") || errors.Is(err, io.EOF)) { + t.Fatalf("expected missing content-length error, got %v", err) + } + + reader = bufio.NewReader(strings.NewReader("Content-Length: nope\r\n\r\n{}")) + if _, err := readFramedMessage(reader); err == nil || !strings.Contains(err.Error(), "invalid content-length") { + t.Fatalf("expected invalid content-length error, got %v", err) + } +} + +func TestReadFramedMessageIgnoresStdoutPreamble(t *testing.T) { + t.Parallel() + + body := `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + raw := "Starting Time MCP server...\nContent-Length: " + strconv.Itoa(len(body)) + "\r\n\r\n" + body + reader := bufio.NewReader(strings.NewReader(raw)) + payload, err := readFramedMessage(reader) + if err != nil { + t.Fatalf("readFramedMessage() error = %v", err) + } + if string(payload) != body { + t.Fatalf("unexpected payload: %s", string(payload)) + } +} + +func TestDecodeCallResultVariants(t *testing.T) { + t.Parallel() + + result := decodeCallResult(json.RawMessage(`{"content":" ok ","isError":true,"extra":1}`)) + if result.Content != "ok" || !result.IsError { + t.Fatalf("unexpected decode result: %+v", result) + } + if result.Metadata["extra"] != float64(1) { + t.Fatalf("expected metadata extra") + } + + result = decodeCallResult(json.RawMessage(`{"content":[{"text":"a"},"b"],"is_error":true}`)) + if result.Content != "a\nb" || !result.IsError { + t.Fatalf("unexpected list content decode: %+v", result) + } + + result = decodeCallResult(json.RawMessage(`{"content":{"nested":"x"}}`)) + if result.Content == "" { + t.Fatalf("expected fallback string content") + } + + result = decodeCallResult(json.RawMessage(`not-json`)) + if result.Content != "not-json" { + t.Fatalf("expected raw fallback content, got %q", result.Content) + } + if _, ok := result.Metadata["raw_result"]; !ok { + t.Fatalf("expected raw_result metadata") + } +} + +func TestFailAllPendingLocked(t *testing.T) { + t.Parallel() + + client := &StdIOClient{ + pending: map[string]chan rpcReply{ + "a": make(chan rpcReply, 1), + "b": make(chan rpcReply, 1), + }, + } + client.failAllPendingLocked(errors.New("closed")) + if len(client.pending) != 0 { + t.Fatalf("expected pending cleared") + } +} + +func TestWriteFramedMessageError(t *testing.T) { + t.Parallel() + + if err := writeFramedMessage(errWriter{}, []byte(`{}`)); err == nil { + t.Fatalf("expected write error") + } +} + +func TestStdIOClientWaitLoopNilCommand(t *testing.T) { + t.Parallel() + + client := &StdIOClient{ + pending: make(map[string]chan rpcReply), + started: true, + } + client.waitLoop(nil) + if client.started { + t.Fatalf("expected started=false after nil command waitLoop") + } +} + +func TestStdIOClientBumpBackoffClamp(t *testing.T) { + t.Parallel() + + client := &StdIOClient{ + cfg: StdioClientConfig{ + RestartBackoff: time.Second, + }, + backoff: maxStdioRestartBackoff, + } + client.bumpBackoffLocked() + if client.backoff != maxStdioRestartBackoff { + t.Fatalf("expected clamp to max backoff, got %v", client.backoff) + } + if client.retryAt.IsZero() { + t.Fatalf("expected retryAt assigned") + } +} + +func newTestStdIOClient(t *testing.T) *StdIOClient { + return newTestStdIOClientWithMode(t, "framed") +} + +func newTestStdIOClientWithMode(t *testing.T, wireMode string) *StdIOClient { + t.Helper() + + if strings.TrimSpace(wireMode) == "" { + wireMode = "framed" + } + + client, err := NewStdIOClient(StdioClientConfig{ + Command: os.Args[0], + Args: []string{"-test.run=TestHelperProcessMCPStdioServer", "--"}, + Env: []string{"GO_WANT_MCP_STDIO_HELPER=1", "GO_MCP_STDIO_REQUIRE_INITIALIZE=1", "GO_MCP_STDIO_WIRE=" + wireMode}, + StartTimeout: 3 * time.Second, + CallTimeout: 3 * time.Second, + }) + if err != nil { + t.Fatalf("NewStdIOClient() error = %v", err) + } + return client +} + +func TestStdIOClientInitializeFailure(t *testing.T) { + t.Parallel() + + client, err := NewStdIOClient(StdioClientConfig{ + Command: os.Args[0], + Args: []string{"-test.run=TestHelperProcessMCPStdioServer", "--"}, + Env: []string{"GO_WANT_MCP_STDIO_HELPER=1", "GO_MCP_STDIO_INIT_FAIL=1", "GO_MCP_STDIO_WIRE=framed"}, + StartTimeout: 3 * time.Second, + CallTimeout: 3 * time.Second, + }) + if err != nil { + t.Fatalf("NewStdIOClient() error = %v", err) + } + defer func() { _ = client.Close() }() + + _, callErr := client.ListTools(context.Background()) + if callErr == nil || !strings.Contains(callErr.Error(), "initialize session") { + t.Fatalf("expected initialize error, got %v", callErr) + } +} + +func TestHelperProcessMCPStdioServer(t *testing.T) { + if os.Getenv("GO_WANT_MCP_STDIO_HELPER") != "1" { + return + } + + requireInitialize := os.Getenv("GO_MCP_STDIO_REQUIRE_INITIALIZE") == "1" + initFail := os.Getenv("GO_MCP_STDIO_INIT_FAIL") == "1" + wireMode := strings.TrimSpace(os.Getenv("GO_MCP_STDIO_WIRE")) + if wireMode == "" { + wireMode = "framed" + } + initialized := !requireInitialize + + reader := bufio.NewReader(os.Stdin) + for { + var ( + payload []byte + err error + ) + switch wireMode { + case "line": + payload, _, err = readRPCMessage(reader) + default: + payload, err = readFramedMessage(reader) + } + if err != nil { + if err == io.EOF { + os.Exit(0) + } + os.Exit(2) + } + + var request map[string]any + if err := json.Unmarshal(payload, &request); err != nil { + os.Exit(3) + } + + method, _ := request["method"].(string) + requestID, _ := request["id"].(string) + + var response any + switch method { + case "initialize": + if initFail { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32600, + "message": "initialize rejected", + }, + } + break + } + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{}, + "serverInfo": map[string]any{ + "name": "test-helper", + "version": "1.0.0", + }, + }, + } + case "notifications/initialized": + initialized = true + continue + case "tools/list": + if !initialized { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32002, + "message": "server not initialized", + }, + } + break + } + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": map[string]any{ + "tools": []map[string]any{ + { + "name": "search", + "description": "search docs", + "inputSchema": map[string]any{ + "type": "object", + "properties": map[string]any{"query": map[string]any{"type": "string"}}, + }, + }, + }, + }, + } + case "tools/call": + if !initialized { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32002, + "message": "server not initialized", + }, + } + break + } + params, _ := request["params"].(map[string]any) + name, _ := params["name"].(string) + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": map[string]any{ + "content": fmt.Sprintf("ok:%s", name), + "isError": false, + }, + } + default: + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32601, + "message": "method not found", + }, + } + } + + rawResponse, err := json.Marshal(response) + if err != nil { + os.Exit(4) + } + switch wireMode { + case "line": + err = writeLineMessage(os.Stdout, rawResponse) + default: + err = writeFramedMessage(os.Stdout, rawResponse) + } + if err != nil { + os.Exit(5) + } + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 90f8a485..a392f629 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -8,11 +8,14 @@ import ( providertypes "neo-code/internal/provider/types" "neo-code/internal/security" + "neo-code/internal/tools/mcp" ) type Registry struct { tools map[string]Tool microCompactPolicies map[string]MicroCompactPolicy + mcpRegistry *mcp.Registry + mcpFactory *mcp.AdapterFactory } func NewRegistry() *Registry { @@ -22,6 +25,15 @@ func NewRegistry() *Registry { } } +// SetMCPRegistry 绑定 MCP registry，用于将远程工具纳入统一执行链。 +func (r *Registry) SetMCPRegistry(registry *mcp.Registry) { + if r == nil || registry == nil { + return + } + r.mcpRegistry = registry + r.mcpFactory = mcp.NewAdapterFactory(registry) +} + func (r *Registry) Register(tool Tool) { if tool == nil { return @@ -46,8 +58,10 @@ func (r *Registry) Get(name string) (Tool, error) { // Supports reports whether a tool is registered. func (r *Registry) Supports(name string) bool { - _, err := r.Get(name) - return err == nil + if _, err := r.Get(name); err == nil { + return true + } + return r.supportsMCPTool(name) } // MicroCompactPolicy 返回指定工具名的 micro compact 策略；未知工具按默认可压缩处理。 @@ -93,30 +107,76 @@ func (r *Registry) ListAvailableSpecs(ctx context.Context, input SpecListInput) if err := ctx.Err(); err != nil { return nil, err } - return r.GetSpecs(), nil + + specs := r.GetSpecs() + mcpAdapters, err := r.listMCPAdapters(ctx) + if err != nil { + return nil, err + } + for _, adapter := range mcpAdapters { + specs = append(specs, providertypes.ToolSpec{ + Name: adapter.FullName(), + Description: adapter.Description(), + Schema: adapter.Schema(), + }) + } + sort.Slice(specs, func(i, j int) bool { + return strings.ToLower(specs[i].Name) < strings.ToLower(specs[j].Name) + }) + return specs, nil } func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) { tool, err := r.Get(input.Name) - if err != nil { - content := FormatError(input.Name, NormalizeErrorReason(input.Name, err), "") + if err == nil { + result, execErr := tool.Execute(ctx, input) + result.ToolCallID = input.ID + if execErr != nil { + result.IsError = true + if strings.TrimSpace(result.Content) == "" { + result.Content = FormatError(result.Name, NormalizeErrorReason(result.Name, execErr), "") + } + return result, execErr + } + return result, nil + } + + adapter, resolveErr := r.resolveMCPAdapter(ctx, input.Name) + if resolveErr != nil { + content := FormatError(input.Name, NormalizeErrorReason(input.Name, resolveErr), "") return ToolResult{ ToolCallID: input.ID, Name: input.Name, Content: content, IsError: true, - }, err + }, resolveErr } - - result, execErr := tool.Execute(ctx, input) - result.ToolCallID = input.ID - if execErr != nil { + callResult, callErr := adapter.Call(ctx, input.Arguments) + result := ToolResult{ + ToolCallID: input.ID, + Name: adapter.FullName(), + Content: strings.TrimSpace(callResult.Content), + IsError: callResult.IsError, + Metadata: map[string]any{ + "mcp_server_id": adapter.ServerID(), + "mcp_tool_name": adapter.ToolName(), + }, + } + for key, value := range callResult.Metadata { + result.Metadata[key] = value + } + if callErr != nil { result.IsError = true if strings.TrimSpace(result.Content) == "" { - result.Content = FormatError(result.Name, NormalizeErrorReason(result.Name, execErr), "") + result.Content = FormatError(result.Name, NormalizeErrorReason(result.Name, callErr), "") } - return result, execErr + result = ApplyOutputLimit(result, DefaultOutputLimitBytes) + return result, callErr + } + if result.Content == "" { + result.Content = "ok" } + result = ApplyOutputLimit(result, DefaultOutputLimitBytes) return result, nil } @@ -124,3 +184,60 @@ func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult func (r *Registry) RememberSessionDecision(sessionID string, action security.Action, scope SessionPermissionScope) error { return errors.New("tools: session permission memory is unsupported by registry manager") } + +// supportsMCPTool 判断指定工具名是否可由当前 MCP 快照解析。 +func (r *Registry) supportsMCPTool(name string) bool { + if r == nil || r.mcpFactory == nil { + return false + } + lowerName := strings.ToLower(strings.TrimSpace(name)) + if !strings.HasPrefix(lowerName, "mcp.") { + return false + } + for _, snapshot := range r.mcpFactoryBuildSnapshot() { + for _, tool := range snapshot.Tools { + if strings.EqualFold(mcpToolFullName(snapshot.ServerID, tool.Name), lowerName) { + return true + } + } + } + return false +} + +// listMCPAdapters 返回 MCP 快照对应的 adapter 列表。 +func (r *Registry) listMCPAdapters(ctx context.Context) ([]*mcp.Adapter, error) { + if r == nil || r.mcpFactory == nil { + return nil, nil + } + return r.mcpFactory.BuildAdapters(ctx) +} + +// resolveMCPAdapter 按完整工具名解析并返回对应 adapter。 +func (r *Registry) resolveMCPAdapter(ctx context.Context, fullName string) (*mcp.Adapter, error) { + adapters, err := r.listMCPAdapters(ctx) + if err != nil { + return nil, err + } + lowerName := strings.ToLower(strings.TrimSpace(fullName)) + if !strings.HasPrefix(lowerName, "mcp.") { + return nil, errors.New("tool: not found") + } + for _, adapter := range adapters { + if strings.EqualFold(adapter.FullName(), lowerName) { + return adapter, nil + } + } + return nil, errors.New("tool: not found") +} + +// mcpFactoryBuildSnapshot 读取 MCP registry 快照，用于无上下文快速检查。 +func (r *Registry) mcpFactoryBuildSnapshot() []mcp.ServerSnapshot { + if r == nil || r.mcpRegistry == nil { + return nil + } + return r.mcpRegistry.Snapshot() +} + +func mcpToolFullName(serverID string, toolName string) string { + return "mcp." + strings.ToLower(strings.TrimSpace(serverID)) + "." + strings.ToLower(strings.TrimSpace(toolName)) +} diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index 2ef5a423..a26843e4 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -7,6 +7,7 @@ import ( "testing" "neo-code/internal/security" + "neo-code/internal/tools/mcp" ) type stubTool struct { @@ -242,3 +243,208 @@ func TestRegistryRememberSessionDecisionUnsupported(t *testing.T) { t.Fatalf("expected unsupported error, got %v", err) } } + +type stubMCPClient struct { + tools []mcp.ToolDescriptor + callResult mcp.CallResult + callErr error +} + +func (s *stubMCPClient) ListTools(ctx context.Context) ([]mcp.ToolDescriptor, error) { + return s.tools, nil +} + +func (s *stubMCPClient) CallTool(ctx context.Context, toolName string, arguments []byte) (mcp.CallResult, error) { + if s.callErr != nil { + return mcp.CallResult{}, s.callErr + } + return s.callResult, nil +} + +func (s *stubMCPClient) HealthCheck(ctx context.Context) error { + return nil +} + +func TestRegistryListAvailableSpecsIncludesMCP(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + registry.Register(stubTool{name: "a_tool", description: "built-in", schema: map[string]any{"type": "object"}}) + + mcpRegistry := mcp.NewRegistry() + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + }); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + specs, err := registry.ListAvailableSpecs(context.Background(), SpecListInput{}) + if err != nil { + t.Fatalf("ListAvailableSpecs() error = %v", err) + } + if len(specs) != 2 { + t.Fatalf("expected 2 specs (built-in + mcp), got %d", len(specs)) + } + foundMCP := false + for _, spec := range specs { + if spec.Name == "mcp.docs.search" { + foundMCP = true + break + } + } + if !foundMCP { + t.Fatalf("expected mcp.docs.search in specs, got %+v", specs) + } +} + +func TestRegistryExecuteDispatchesToMCPAdapter(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + mcpRegistry := mcp.NewRegistry() + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: mcp.CallResult{ + Content: "mcp ok", + Metadata: map[string]any{ + "latency_ms": 12, + }, + }, + }); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + result, err := registry.Execute(context.Background(), ToolCallInput{ + ID: "mcp-call-1", + Name: "mcp.docs.search", + Arguments: []byte(`{"query":"neocode"}`), + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if result.ToolCallID != "mcp-call-1" { + t.Fatalf("expected tool call id mcp-call-1, got %q", result.ToolCallID) + } + if result.Name != "mcp.docs.search" { + t.Fatalf("expected mcp tool name, got %q", result.Name) + } + if !strings.Contains(result.Content, "mcp ok") { + t.Fatalf("expected mcp content, got %q", result.Content) + } + if result.Metadata["mcp_server_id"] != "docs" || result.Metadata["mcp_tool_name"] != "search" { + t.Fatalf("unexpected mcp metadata: %+v", result.Metadata) + } +} + +func TestRegistryExecuteMCPResolveErrorPropagates(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + mcpRegistry := mcp.NewRegistry() + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + }); err != nil { + t.Fatalf("register mcp server: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + result, err := registry.Execute(ctx, ToolCallInput{ + ID: "mcp-call-canceled", + Name: "mcp.docs.search", + }) + if err == nil || !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", err) + } + if !result.IsError { + t.Fatalf("expected error result") + } + if !strings.Contains(result.Content, context.Canceled.Error()) { + t.Fatalf("expected canceled content, got %q", result.Content) + } +} + +func TestRegistryExecuteMCPCallErrorDoesNotReturnOK(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + mcpRegistry := mcp.NewRegistry() + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + callErr: errors.New("mcp transport timeout"), + }); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + result, err := registry.Execute(context.Background(), ToolCallInput{ + ID: "mcp-call-error", + Name: "mcp.docs.search", + Arguments: []byte(`{"query":"neocode"}`), + }) + if err == nil { + t.Fatalf("expected mcp call error") + } + if !result.IsError { + t.Fatalf("expected IsError true") + } + if strings.TrimSpace(result.Content) == "" || strings.EqualFold(strings.TrimSpace(result.Content), "ok") { + t.Fatalf("expected non-ok error content, got %q", result.Content) + } +} + +func TestRegistrySupportsMCPToolAndHelpers(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + mcpRegistry := mcp.NewRegistry() + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "search", Description: "search docs", InputSchema: map[string]any{"type": "object"}}, + }, + }); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + if !registry.Supports("mcp.docs.search") { + t.Fatalf("expected supports mcp.docs.search") + } + if registry.Supports("mcp.docs.missing") { + t.Fatalf("did not expect supports mcp.docs.missing") + } + if registry.Supports("search") { + t.Fatalf("did not expect supports non-prefixed mcp name") + } + + snapshots := registry.mcpFactoryBuildSnapshot() + if len(snapshots) != 1 { + t.Fatalf("expected one snapshot, got %d", len(snapshots)) + } + if got := mcpToolFullName(" Docs ", " Search "); got != "mcp.docs.search" { + t.Fatalf("unexpected mcp full name: %q", got) + } +}