From 48ba7cf823e6b5a9b33baa9d50fb4aaa46d5f768 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:05:49 +0800 Subject: [PATCH 01/13] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=20MCP=20Regist?= =?UTF-8?q?ry=20=E4=B8=8E=20Adapter=20=E5=9F=BA=E7=A1=80=E8=83=BD=E5=8A=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/adapter.go | 171 +++++++++++++++ internal/tools/mcp/adapter_test.go | 135 ++++++++++++ internal/tools/mcp/registry.go | 319 ++++++++++++++++++++++++++++ internal/tools/mcp/registry_test.go | 163 ++++++++++++++ 4 files changed, 788 insertions(+) create mode 100644 internal/tools/mcp/adapter.go create mode 100644 internal/tools/mcp/adapter_test.go create mode 100644 internal/tools/mcp/registry.go create mode 100644 internal/tools/mcp/registry_test.go diff --git a/internal/tools/mcp/adapter.go b/internal/tools/mcp/adapter.go new file mode 100644 index 00000000..e238e07e --- /dev/null +++ b/internal/tools/mcp/adapter.go @@ -0,0 +1,171 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + + "neo-code/internal/tools" +) + +const mcpToolNamePrefix = "mcp." + +// AdapterFactory 基于 registry 快照构造 MCP tool 适配器集合。 +type AdapterFactory struct { + registry *Registry +} + +// NewAdapterFactory 创建 MCP adapter 工厂。 +func NewAdapterFactory(registry *Registry) *AdapterFactory { + return &AdapterFactory{registry: registry} +} + +// BuildTools 将当前所有 MCP tool 快照转换为统一 tools.Tool 列表。 +func (f *AdapterFactory) BuildTools(ctx context.Context) ([]tools.Tool, 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([]tools.Tool, 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 适配为统一 tools.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 +} + +// Name 返回统一的 MCP tool 名称：mcp..。 +func (a *Adapter) Name() string { + return composeToolName(a.serverID, 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) +} + +// MicroCompactPolicy 返回 MCP tool 历史结果默认 micro compact 策略。 +func (a *Adapter) MicroCompactPolicy() tools.MicroCompactPolicy { + return tools.MicroCompactPolicyCompact +} + +// Execute 分发 MCP tool 调用并收敛为统一 ToolResult。 +func (a *Adapter) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) { + if a == nil || a.registry == nil { + err := errors.New("mcp: adapter is not initialized") + return tools.NewErrorResult("mcp", tools.NormalizeErrorReason("mcp", err), "", nil), err + } + if err := ctx.Err(); err != nil { + return tools.NewErrorResult(a.Name(), tools.NormalizeErrorReason(a.Name(), err), "", adapterMetadata(a.serverID, a.toolName)), err + } + + result, err := a.registry.Call(ctx, a.serverID, a.toolName, call.Arguments) + if err != nil { + errorResult := tools.NewErrorResult(a.Name(), tools.NormalizeErrorReason(a.Name(), err), "", adapterMetadata(a.serverID, a.toolName)) + errorResult.ToolCallID = call.ID + return errorResult, err + } + + metadata := adapterMetadata(a.serverID, a.toolName) + for key, value := range result.Metadata { + metadata[key] = value + } + + toolResult := tools.ToolResult{ + ToolCallID: call.ID, + Name: a.Name(), + Content: strings.TrimSpace(result.Content), + IsError: result.IsError, + Metadata: metadata, + } + if strings.TrimSpace(toolResult.Content) == "" { + toolResult.Content = "ok" + } + return tools.ApplyOutputLimit(toolResult, tools.DefaultOutputLimitBytes), nil +} + +// 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 +} + +// adapterMetadata 生成 MCP 调用结果的基础元信息。 +func adapterMetadata(serverID string, toolName string) map[string]any { + return map[string]any{ + "mcp_server_id": normalizeServerID(serverID), + "mcp_tool_name": strings.TrimSpace(toolName), + } +} diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go new file mode 100644 index 00000000..689075b3 --- /dev/null +++ b/internal/tools/mcp/adapter_test.go @@ -0,0 +1,135 @@ +package mcp + +import ( + "context" + "errors" + "testing" + + "neo-code/internal/tools" +) + +func TestAdapterFactoryBuildTools(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) + toolsList, err := factory.BuildTools(context.Background()) + if err != nil { + t.Fatalf("BuildTools() error = %v", err) + } + if len(toolsList) != 1 { + t.Fatalf("expected one adapter tool, got %d", len(toolsList)) + } + if toolsList[0].Name() != "mcp.docs.search" { + t.Fatalf("unexpected adapter tool name: %q", toolsList[0].Name()) + } +} + +func TestAdapterExecute(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.Execute(context.Background(), tools.ToolCallInput{ + ID: "tool-call-1", + Name: adapter.Name(), + Arguments: []byte(`{"q":"mcp"}`), + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if result.ToolCallID != "tool-call-1" { + t.Fatalf("expected tool call id tool-call-1, got %q", result.ToolCallID) + } + if result.Name != "mcp.docs.search" { + t.Fatalf("expected tool name mcp.docs.search, got %q", result.Name) + } + if result.Content != "result body" { + t.Fatalf("expected result content, got %q", result.Content) + } + if result.Metadata["mcp_server_id"] != "docs" || result.Metadata["mcp_tool_name"] != "search" { + t.Fatalf("unexpected metadata: %+v", result.Metadata) + } +} + +func TestAdapterExecuteErrorMapping(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) + } + + result, execErr := adapter.Execute(context.Background(), tools.ToolCallInput{ + ID: "tool-call-error", + Name: adapter.Name(), + Arguments: []byte(`{"q":"mcp"}`), + }) + if execErr == nil { + t.Fatalf("expected execute error") + } + if !result.IsError { + t.Fatalf("expected error result, got %+v", result) + } + if result.Metadata["mcp_server_id"] != "docs" { + t.Fatalf("unexpected metadata for error result: %+v", result.Metadata) + } +} diff --git a/internal/tools/mcp/registry.go b/internal/tools/mcp/registry.go new file mode 100644 index 00000000..026983ec --- /dev/null +++ b/internal/tools/mcp/registry.go @@ -0,0 +1,319 @@ +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] = value + } + return cloned +} diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go new file mode 100644 index 00000000..d7f5f26d --- /dev/null +++ b/internal/tools/mcp/registry_test.go @@ -0,0 +1,163 @@ +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") + } +} From e6fb96eb3eaff28d820aeebcc4a18cf954e5d456 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:06:03 +0800 Subject: [PATCH 02/13] =?UTF-8?q?feat:=20=E5=B0=86=20MCP=20=E5=B7=A5?= =?UTF-8?q?=E5=85=B7=E6=8E=A5=E5=85=A5=20Registry=20=E4=B8=BB=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=E9=93=BE=E5=B9=B6=E5=A2=9E=E5=8A=A0=20stdio=20?= =?UTF-8?q?=E5=AE=A2=E6=88=B7=E7=AB=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/adapter.go | 71 +--- internal/tools/mcp/adapter_test.go | 52 +-- internal/tools/mcp/stdio_client.go | 499 ++++++++++++++++++++++++ internal/tools/mcp/stdio_client_test.go | 139 +++++++ internal/tools/registry.go | 136 ++++++- internal/tools/registry_test.go | 105 +++++ 6 files changed, 903 insertions(+), 99 deletions(-) create mode 100644 internal/tools/mcp/stdio_client.go create mode 100644 internal/tools/mcp/stdio_client_test.go diff --git a/internal/tools/mcp/adapter.go b/internal/tools/mcp/adapter.go index e238e07e..d9ecb470 100644 --- a/internal/tools/mcp/adapter.go +++ b/internal/tools/mcp/adapter.go @@ -5,8 +5,6 @@ import ( "errors" "fmt" "strings" - - "neo-code/internal/tools" ) const mcpToolNamePrefix = "mcp." @@ -21,8 +19,8 @@ func NewAdapterFactory(registry *Registry) *AdapterFactory { return &AdapterFactory{registry: registry} } -// BuildTools 将当前所有 MCP tool 快照转换为统一 tools.Tool 列表。 -func (f *AdapterFactory) BuildTools(ctx context.Context) ([]tools.Tool, error) { +// 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") } @@ -35,7 +33,7 @@ func (f *AdapterFactory) BuildTools(ctx context.Context) ([]tools.Tool, error) { return nil, nil } - result := make([]tools.Tool, 0, len(snapshots)*2) + result := make([]*Adapter, 0, len(snapshots)*2) for _, snapshot := range snapshots { for _, descriptor := range snapshot.Tools { adapter, err := NewAdapter(f.registry, snapshot.ServerID, descriptor) @@ -48,7 +46,7 @@ func (f *AdapterFactory) BuildTools(ctx context.Context) ([]tools.Tool, error) { return result, nil } -// Adapter 将单个 MCP tool 适配为统一 tools.Tool 接口。 +// Adapter 将单个 MCP tool 适配为统一调用描述。 type Adapter struct { registry *Registry serverID string @@ -80,11 +78,21 @@ func NewAdapter(registry *Registry, serverID string, descriptor ToolDescriptor) }, nil } -// Name 返回统一的 MCP tool 名称：mcp..。 -func (a *Adapter) Name() string { +// 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) != "" { @@ -98,44 +106,15 @@ func (a *Adapter) Schema() map[string]any { return cloneSchema(a.schema) } -// MicroCompactPolicy 返回 MCP tool 历史结果默认 micro compact 策略。 -func (a *Adapter) MicroCompactPolicy() tools.MicroCompactPolicy { - return tools.MicroCompactPolicyCompact -} - -// Execute 分发 MCP tool 调用并收敛为统一 ToolResult。 -func (a *Adapter) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) { +// Call 分发 MCP tool 调用并返回统一结果。 +func (a *Adapter) Call(ctx context.Context, arguments []byte) (CallResult, error) { if a == nil || a.registry == nil { - err := errors.New("mcp: adapter is not initialized") - return tools.NewErrorResult("mcp", tools.NormalizeErrorReason("mcp", err), "", nil), err + return CallResult{}, errors.New("mcp: adapter is not initialized") } if err := ctx.Err(); err != nil { - return tools.NewErrorResult(a.Name(), tools.NormalizeErrorReason(a.Name(), err), "", adapterMetadata(a.serverID, a.toolName)), err - } - - result, err := a.registry.Call(ctx, a.serverID, a.toolName, call.Arguments) - if err != nil { - errorResult := tools.NewErrorResult(a.Name(), tools.NormalizeErrorReason(a.Name(), err), "", adapterMetadata(a.serverID, a.toolName)) - errorResult.ToolCallID = call.ID - return errorResult, err + return CallResult{}, err } - - metadata := adapterMetadata(a.serverID, a.toolName) - for key, value := range result.Metadata { - metadata[key] = value - } - - toolResult := tools.ToolResult{ - ToolCallID: call.ID, - Name: a.Name(), - Content: strings.TrimSpace(result.Content), - IsError: result.IsError, - Metadata: metadata, - } - if strings.TrimSpace(toolResult.Content) == "" { - toolResult.Content = "ok" - } - return tools.ApplyOutputLimit(toolResult, tools.DefaultOutputLimitBytes), nil + return a.registry.Call(ctx, a.serverID, a.toolName, arguments) } // composeToolName 组装统一的 MCP tool 名称，保持权限映射可预测。 @@ -161,11 +140,3 @@ func ensureObjectSchema(schema map[string]any) map[string]any { } return cloned } - -// adapterMetadata 生成 MCP 调用结果的基础元信息。 -func adapterMetadata(serverID string, toolName string) map[string]any { - return map[string]any{ - "mcp_server_id": normalizeServerID(serverID), - "mcp_tool_name": strings.TrimSpace(toolName), - } -} diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go index 689075b3..d5dd16f7 100644 --- a/internal/tools/mcp/adapter_test.go +++ b/internal/tools/mcp/adapter_test.go @@ -4,11 +4,9 @@ import ( "context" "errors" "testing" - - "neo-code/internal/tools" ) -func TestAdapterFactoryBuildTools(t *testing.T) { +func TestAdapterFactoryBuildAdapters(t *testing.T) { t.Parallel() registry := NewRegistry() @@ -29,19 +27,19 @@ func TestAdapterFactoryBuildTools(t *testing.T) { } factory := NewAdapterFactory(registry) - toolsList, err := factory.BuildTools(context.Background()) + adapters, err := factory.BuildAdapters(context.Background()) if err != nil { - t.Fatalf("BuildTools() error = %v", err) + t.Fatalf("BuildAdapters() error = %v", err) } - if len(toolsList) != 1 { - t.Fatalf("expected one adapter tool, got %d", len(toolsList)) + if len(adapters) != 1 { + t.Fatalf("expected one adapter, got %d", len(adapters)) } - if toolsList[0].Name() != "mcp.docs.search" { - t.Fatalf("unexpected adapter tool name: %q", toolsList[0].Name()) + if adapters[0].FullName() != "mcp.docs.search" { + t.Fatalf("unexpected adapter full name: %q", adapters[0].FullName()) } } -func TestAdapterExecute(t *testing.T) { +func TestAdapterCall(t *testing.T) { t.Parallel() registry := NewRegistry() @@ -72,29 +70,16 @@ func TestAdapterExecute(t *testing.T) { t.Fatalf("NewAdapter() error = %v", err) } - result, err := adapter.Execute(context.Background(), tools.ToolCallInput{ - ID: "tool-call-1", - Name: adapter.Name(), - Arguments: []byte(`{"q":"mcp"}`), - }) + result, err := adapter.Call(context.Background(), []byte(`{"q":"mcp"}`)) if err != nil { - t.Fatalf("Execute() error = %v", err) - } - if result.ToolCallID != "tool-call-1" { - t.Fatalf("expected tool call id tool-call-1, got %q", result.ToolCallID) - } - if result.Name != "mcp.docs.search" { - t.Fatalf("expected tool name mcp.docs.search, got %q", result.Name) + t.Fatalf("Call() error = %v", err) } if result.Content != "result body" { t.Fatalf("expected result content, got %q", result.Content) } - if result.Metadata["mcp_server_id"] != "docs" || result.Metadata["mcp_tool_name"] != "search" { - t.Fatalf("unexpected metadata: %+v", result.Metadata) - } } -func TestAdapterExecuteErrorMapping(t *testing.T) { +func TestAdapterCallError(t *testing.T) { t.Parallel() registry := NewRegistry() @@ -118,18 +103,7 @@ func TestAdapterExecuteErrorMapping(t *testing.T) { t.Fatalf("NewAdapter() error = %v", err) } - result, execErr := adapter.Execute(context.Background(), tools.ToolCallInput{ - ID: "tool-call-error", - Name: adapter.Name(), - Arguments: []byte(`{"q":"mcp"}`), - }) - if execErr == nil { - t.Fatalf("expected execute error") - } - if !result.IsError { - t.Fatalf("expected error result, got %+v", result) - } - if result.Metadata["mcp_server_id"] != "docs" { - t.Fatalf("unexpected metadata for error result: %+v", result.Metadata) + if _, err := adapter.Call(context.Background(), []byte(`{"q":"mcp"}`)); err == nil { + t.Fatalf("expected call error") } } diff --git a/internal/tools/mcp/stdio_client.go b/internal/tools/mcp/stdio_client.go new file mode 100644 index 00000000..7e92fe80 --- /dev/null +++ b/internal/tools/mcp/stdio_client.go @@ -0,0 +1,499 @@ +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 +) + +// 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 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 + 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 + shutdown bool +} + +// 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 +} + +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) +} + +func (c *StdIOClient) call(ctx context.Context, method string, params any) (json.RawMessage, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := c.ensureStarted(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 + c.mu.Unlock() + + 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) + } + if err := writeFramedMessage(stdin, requestPayload); err != nil { + c.removePending(requestID) + return nil, fmt.Errorf("mcp: send request: %w", err) + } + + select { + case <-ctx.Done(): + c.removePending(requestID) + return nil, ctx.Err() + case reply := <-replyCh: + return reply.result, reply.err + } +} + +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.backoff = c.cfg.RestartBackoff + c.retryAt = time.Time{} + + go c.readLoop() + go c.waitLoop() + go io.Copy(io.Discard, stderr) + return nil +} + +func (c *StdIOClient) readLoop() { + for { + message, err := readFramedMessage(c.reader) + if err != nil { + c.markExited(fmt.Errorf("mcp: read response: %w", err)) + return + } + + 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} + } +} + +func (c *StdIOClient) waitLoop() { + err := c.cmd.Wait() + c.markExited(fmt.Errorf("mcp: stdio process exited: %w", err)) +} + +func (c *StdIOClient) markExited(err error) { + c.mu.Lock() + defer c.mu.Unlock() + + if !c.started { + return + } + c.started = false + 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() +} + +func (c *StdIOClient) removePending(requestID string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.pending, requestID) +} + +func (c *StdIOClient) failAllPendingLocked(err error) { + for requestID, replyCh := range c.pending { + replyCh <- rpcReply{err: err} + delete(c.pending, requestID) + } +} + +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 + } +} + +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 +} + +func readFramedMessage(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 == "" { + break + } + + 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, fmt.Errorf("mcp: invalid content-length %q", rawLength) + } + contentLength = length + } + } + if contentLength < 0 { + return nil, errors.New("mcp: missing content-length header") + } + + payload := make([]byte, contentLength) + if _, err := io.ReadFull(reader, payload); err != nil { + return nil, err + } + return payload, nil +} + +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"] = 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..5235e454 --- /dev/null +++ b/internal/tools/mcp/stdio_client_test.go @@ -0,0 +1,139 @@ +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "strings" + "testing" + "time" +) + +func TestStdIOClientListToolsAndCallTool(t *testing.T) { + t.Parallel() + + client := newTestStdIOClient(t) + 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 TestStdIOClientHealthCheck(t *testing.T) { + t.Parallel() + + client := newTestStdIOClient(t) + defer func() { _ = client.Close() }() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := client.HealthCheck(ctx); err != nil { + t.Fatalf("HealthCheck() error = %v", err) + } +} + +func newTestStdIOClient(t *testing.T) *StdIOClient { + t.Helper() + + client, err := NewStdIOClient(StdioClientConfig{ + Command: os.Args[0], + Args: []string{"-test.run=TestHelperProcessMCPStdioServer", "--"}, + Env: []string{"GO_WANT_MCP_STDIO_HELPER=1"}, + StartTimeout: 3 * time.Second, + CallTimeout: 3 * time.Second, + }) + if err != nil { + t.Fatalf("NewStdIOClient() error = %v", err) + } + return client +} + +func TestHelperProcessMCPStdioServer(t *testing.T) { + if os.Getenv("GO_WANT_MCP_STDIO_HELPER") != "1" { + return + } + + reader := bufio.NewReader(os.Stdin) + for { + 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 "tools/list": + 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": + 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) + } + if err := writeFramedMessage(os.Stdout, rawResponse); err != nil { + os.Exit(5) + } + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 90f8a485..0b6b7d72 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,12 +107,42 @@ 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, provider.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 { + 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, err), "") return ToolResult{ ToolCallID: input.ID, @@ -107,15 +151,30 @@ func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult IsError: true, }, err } - - 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 result.Content == "" { + result.Content = "ok" + } + result = ApplyOutputLimit(result, DefaultOutputLimitBytes) + 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 + return result, callErr } return result, nil } @@ -124,3 +183,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..83632adb 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,107 @@ 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) + } +} From 7b5db03addbf3103788f6ea6a6d540ca462f33b2 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:06:17 +0800 Subject: [PATCH 03/13] =?UTF-8?q?feat:=20=E5=BC=BA=E5=8C=96=20MCP=20?= =?UTF-8?q?=E8=B0=83=E7=94=A8=E9=93=BE=E9=94=99=E8=AF=AF=E5=A4=84=E7=90=86?= =?UTF-8?q?=E4=B8=8E=E5=AE=89=E5=85=A8=E8=BE=B9=E7=95=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/registry.go | 22 +++++++++++++++++++++- internal/tools/mcp/stdio_client.go | 28 ++++++++++++++++++++++------ internal/tools/registry.go | 13 +++++++------ 3 files changed, 50 insertions(+), 13 deletions(-) diff --git a/internal/tools/mcp/registry.go b/internal/tools/mcp/registry.go index 026983ec..490b5142 100644 --- a/internal/tools/mcp/registry.go +++ b/internal/tools/mcp/registry.go @@ -313,7 +313,27 @@ func cloneSchema(schema map[string]any) map[string]any { } cloned := make(map[string]any, len(schema)) for key, value := range schema { - cloned[key] = value + 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/stdio_client.go b/internal/tools/mcp/stdio_client.go index 7e92fe80..7b9196f9 100644 --- a/internal/tools/mcp/stdio_client.go +++ b/internal/tools/mcp/stdio_client.go @@ -22,6 +22,7 @@ const ( defaultStdioCallTimeout = 15 * time.Second defaultStdioRestartBackoff = 1 * time.Second maxStdioRestartBackoff = 30 * time.Second + maxStdioFrameBytes = 8 * 1024 * 1024 ) // StdioClientConfig 描述 MCP stdio 客户端的启动与调用参数。 @@ -64,6 +65,7 @@ type StdIOClient struct { cfg StdioClientConfig idSeed uint64 mu sync.Mutex + writeMu sync.Mutex cmd *exec.Cmd stdin io.WriteCloser stdout io.ReadCloser @@ -217,6 +219,10 @@ func (c *StdIOClient) call(ctx context.Context, method string, params any) (json c.pending[requestID] = replyCh stdin := c.stdin 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", @@ -228,9 +234,12 @@ func (c *StdIOClient) call(ctx context.Context, method string, params any) (json c.removePending(requestID) return nil, fmt.Errorf("mcp: marshal request: %w", err) } - if err := writeFramedMessage(stdin, requestPayload); err != nil { + c.writeMu.Lock() + writeErr := writeFramedMessage(stdin, requestPayload) + c.writeMu.Unlock() + if writeErr != nil { c.removePending(requestID) - return nil, fmt.Errorf("mcp: send request: %w", err) + return nil, fmt.Errorf("mcp: send request: %w", writeErr) } select { @@ -301,7 +310,7 @@ func (c *StdIOClient) ensureStarted(ctx context.Context) error { c.retryAt = time.Time{} go c.readLoop() - go c.waitLoop() + go c.waitLoop(command) go io.Copy(io.Discard, stderr) return nil } @@ -342,8 +351,12 @@ func (c *StdIOClient) readLoop() { } } -func (c *StdIOClient) waitLoop() { - err := c.cmd.Wait() +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)) } @@ -427,6 +440,9 @@ func readFramedMessage(reader *bufio.Reader) ([]byte, error) { 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 { @@ -489,7 +505,7 @@ func decodeCallResult(raw json.RawMessage) CallResult { } metadata[key] = value } - metadata["raw_result"] = bytes.TrimSpace(raw) + metadata["raw_result"] = string(bytes.TrimSpace(raw)) return CallResult{ Content: content, diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 0b6b7d72..c03bfcb6 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -143,13 +143,13 @@ func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult adapter, resolveErr := r.resolveMCPAdapter(ctx, input.Name) if resolveErr != nil { - content := FormatError(input.Name, NormalizeErrorReason(input.Name, err), "") + content := FormatError(input.Name, NormalizeErrorReason(input.Name, resolveErr), "") return ToolResult{ ToolCallID: input.ID, Name: input.Name, Content: content, IsError: true, - }, err + }, resolveErr } callResult, callErr := adapter.Call(ctx, input.Arguments) result := ToolResult{ @@ -165,17 +165,18 @@ func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult for key, value := range callResult.Metadata { result.Metadata[key] = value } - if result.Content == "" { - result.Content = "ok" - } - result = ApplyOutputLimit(result, DefaultOutputLimitBytes) if callErr != nil { result.IsError = true if strings.TrimSpace(result.Content) == "" { result.Content = FormatError(result.Name, NormalizeErrorReason(result.Name, callErr), "") } + result = ApplyOutputLimit(result, DefaultOutputLimitBytes) return result, callErr } + if result.Content == "" { + result.Content = "ok" + } + result = ApplyOutputLimit(result, DefaultOutputLimitBytes) return result, nil } From 4bc19abef7f5b287ecf918dc4854d58b040a7f6e Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:06:27 +0800 Subject: [PATCH 04/13] =?UTF-8?q?feat:=20=E8=A1=A5=E9=BD=90=20MCP=20?= =?UTF-8?q?=E5=B9=B6=E5=8F=91=E4=B8=8E=E9=94=99=E8=AF=AF=E8=B7=AF=E5=BE=84?= =?UTF-8?q?=E5=9B=9E=E5=BD=92=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/registry_test.go | 47 ++++++++++++++++++ internal/tools/mcp/stdio_client_test.go | 46 +++++++++++++++++ internal/tools/registry_test.go | 65 +++++++++++++++++++++++++ 3 files changed, 158 insertions(+) diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go index d7f5f26d..b0b9c99d 100644 --- a/internal/tools/mcp/registry_test.go +++ b/internal/tools/mcp/registry_test.go @@ -161,3 +161,50 @@ func TestRegistryConcurrentSnapshotAndRefresh(t *testing.T) { 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"]) + } +} diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index 5235e454..2dea5e16 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -8,6 +8,7 @@ import ( "io" "os" "strings" + "sync" "testing" "time" ) @@ -48,6 +49,51 @@ func TestStdIOClientHealthCheck(t *testing.T) { } } +func TestStdIOClientConcurrentCallTool(t *testing.T) { + t.Parallel() + + client := newTestStdIOClient(t) + 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 newTestStdIOClient(t *testing.T) *StdIOClient { t.Helper() diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index 83632adb..3cc882d8 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -347,3 +347,68 @@ func TestRegistryExecuteDispatchesToMCPAdapter(t *testing.T) { 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) + } +} From 76f1b4862e99a1ccf99d37b2e8bccff4a7ab52dd Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:38:02 +0800 Subject: [PATCH 05/13] =?UTF-8?q?feat:=20=E5=AE=8C=E5=96=84=20MCP=20?= =?UTF-8?q?=E6=96=B0=E8=83=BD=E5=8A=9B=E6=B5=8B=E8=AF=95=E8=A6=86=E7=9B=96?= =?UTF-8?q?=E4=B8=8E=E8=BE=B9=E7=95=8C=E7=94=A8=E4=BE=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/adapter_test.go | 39 +++++++ internal/tools/mcp/registry_test.go | 36 +++++++ internal/tools/mcp/stdio_client_test.go | 131 ++++++++++++++++++++++++ 3 files changed, 206 insertions(+) diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go index d5dd16f7..f51f9191 100644 --- a/internal/tools/mcp/adapter_test.go +++ b/internal/tools/mcp/adapter_test.go @@ -107,3 +107,42 @@ func TestAdapterCallError(t *testing.T) { 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"]) + } +} diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go index b0b9c99d..49f2c654 100644 --- a/internal/tools/mcp/registry_test.go +++ b/internal/tools/mcp/registry_test.go @@ -208,3 +208,39 @@ func TestRegistrySnapshotSchemaIsDeepCloned(t *testing.T) { 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") + } +} diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index 2dea5e16..530677ad 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "io" "os" @@ -94,6 +95,136 @@ func TestReadFramedMessageRejectsOversizedPayload(t *testing.T) { } } +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") { + 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 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 newTestStdIOClient(t *testing.T) *StdIOClient { t.Helper() From 8486c26f3af141054b4474e47a7daa8c4d4d5184 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 12:57:42 +0800 Subject: [PATCH 06/13] =?UTF-8?q?feat:=20=E6=8F=90=E5=8D=87=20#176/#177=20?= =?UTF-8?q?=E5=B7=AE=E5=BC=82=E8=A6=86=E7=9B=96=E7=8E=87=E5=B9=B6=E8=A1=A5?= =?UTF-8?q?=E5=85=85=E5=88=86=E6=94=AF=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/adapter_test.go | 47 +++++++++++++++++++++++++ internal/tools/mcp/stdio_client_test.go | 45 +++++++++++++++++++++++ internal/tools/registry_test.go | 36 +++++++++++++++++++ 3 files changed, 128 insertions(+) diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go index f51f9191..8f3adc8f 100644 --- a/internal/tools/mcp/adapter_test.go +++ b/internal/tools/mcp/adapter_test.go @@ -146,3 +146,50 @@ func TestAdapterAccessorsAndSchemaClone(t *testing.T) { 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") + } +} diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index 530677ad..3a352d3a 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -14,6 +14,12 @@ import ( "time" ) +type errWriter struct{} + +func (errWriter) Write(p []byte) (int, error) { + return 0, errors.New("write failed") +} + func TestStdIOClientListToolsAndCallTool(t *testing.T) { t.Parallel() @@ -225,6 +231,45 @@ func TestFailAllPendingLocked(t *testing.T) { } } +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 { t.Helper() diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index 3cc882d8..a26843e4 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -412,3 +412,39 @@ func TestRegistryExecuteMCPCallErrorDoesNotReturnOK(t *testing.T) { 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) + } +} From 31c61df9d468811b7ee184040d25ed0434a73421 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 15:50:06 +0800 Subject: [PATCH 07/13] =?UTF-8?q?feat:=20=E6=8E=A5=E5=85=A5=E9=85=8D?= =?UTF-8?q?=E7=BD=AE=E9=A9=B1=E5=8A=A8=20MCP=20=E6=B3=A8=E5=86=8C=E5=B9=B6?= =?UTF-8?q?=E8=A1=A5=E5=85=85=E4=BD=BF=E7=94=A8=E6=96=87=E6=A1=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/guides/mcp-configuration.md | 62 ++++++++++++ internal/app/bootstrap.go | 16 ++- internal/app/bootstrap_test.go | 162 ++++++++++++++++++++++++++++++- internal/app/mcp_bootstrap.go | 154 +++++++++++++++++++++++++++++ internal/config/config_test.go | 69 +++++++++++++ internal/config/model.go | 137 ++++++++++++++++++++++++++ 6 files changed, 596 insertions(+), 4 deletions(-) create mode 100644 docs/guides/mcp-configuration.md create mode 100644 internal/app/mcp_bootstrap.go 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..67e7ba2a 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -10,9 +10,11 @@ import ( "path/filepath" "strings" "testing" + "time" "neo-code/internal/config" "neo-code/internal/tools" + "neo-code/internal/tools/mcp" ) func TestNewProgram(t *testing.T) { @@ -84,7 +86,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 +115,145 @@ 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 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 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 +396,19 @@ func disableBuiltinProviderAPIKeys(t *testing.T) { t.Setenv(config.OpenLLDefaultAPIKeyEnv, "") t.Setenv(config.QiniuDefaultAPIKeyEnv, "") } + +type stubMCPServerClient struct { + tools []mcp.ToolDescriptor +} + +func (s *stubMCPServerClient) ListTools(ctx context.Context) ([]mcp.ToolDescriptor, error) { + 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 +} 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 { From a07b88deac1105a3170913a8b3986422e078537a Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 16:55:57 +0800 Subject: [PATCH 08/13] =?UTF-8?q?feat:=20=E8=A1=A5=E9=BD=90=20MCP=20stdio?= =?UTF-8?q?=20initialize=20=E6=8F=A1=E6=89=8B=E4=B8=8E=E9=87=8D=E8=BF=9E?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E6=B5=81=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/stdio_client.go | 171 +++++++++++++++++++++--- internal/tools/mcp/stdio_client_test.go | 76 ++++++++++- 2 files changed, 229 insertions(+), 18 deletions(-) diff --git a/internal/tools/mcp/stdio_client.go b/internal/tools/mcp/stdio_client.go index 7b9196f9..71d28d3c 100644 --- a/internal/tools/mcp/stdio_client.go +++ b/internal/tools/mcp/stdio_client.go @@ -23,6 +23,9 @@ const ( defaultStdioRestartBackoff = 1 * time.Second maxStdioRestartBackoff = 30 * time.Second maxStdioFrameBytes = 8 * 1024 * 1024 + defaultMCPProtocolVersion = "2024-11-05" + defaultMCPClientName = "neocode" + defaultMCPClientVersion = "0.1.0" ) // StdioClientConfig 描述 MCP stdio 客户端的启动与调用参数。 @@ -43,6 +46,12 @@ type jsonRPCRequest struct { 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"` @@ -62,21 +71,24 @@ type rpcReply struct { // 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 - shutdown bool + 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{} + shutdown bool } // NewStdIOClient 创建 stdio MCP client。 @@ -201,11 +213,21 @@ func (c *StdIOClient) callContext(ctx context.Context) (context.Context, context } 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) { if err := ctx.Err(); err != nil { return nil, err } - if err := c.ensureStarted(ctx); 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) @@ -251,6 +273,49 @@ func (c *StdIOClient) call(ctx context.Context, method string, params any) (json } } +// sendNotification 发送无需响应的 RPC 通知；skipEnsure=true 用于初始化流程。 +func (c *StdIOClient) sendNotification(ctx context.Context, method string, params any, skipEnsure bool) 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 + 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 := writeFramedMessage(stdin, payload) + c.writeMu.Unlock() + if writeErr != nil { + return fmt.Errorf("mcp: send notification: %w", writeErr) + } + return nil +} + func (c *StdIOClient) ensureStarted(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() @@ -306,6 +371,9 @@ func (c *StdIOClient) ensureStarted(ctx context.Context) error { c.exited = make(chan struct{}) c.exitErr = nil c.started = true + c.initialized = false + c.initializing = false + c.initDone = nil c.backoff = c.cfg.RestartBackoff c.retryAt = time.Time{} @@ -315,6 +383,69 @@ func (c *StdIOClient) ensureStarted(ctx context.Context) error { 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 初始化握手：initialize -> notifications/initialized。 +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, + }, + } + if _, err := c.callRequest(ctx, "initialize", params, true); err != nil { + return fmt.Errorf("mcp: initialize session: %w", err) + } + if err := c.sendNotification(ctx, "notifications/initialized", map[string]any{}, true); err != nil { + return fmt.Errorf("mcp: notify initialized: %w", err) + } + return nil +} + func (c *StdIOClient) readLoop() { for { message, err := readFramedMessage(c.reader) @@ -368,6 +499,12 @@ func (c *StdIOClient) markExited(err error) { return } c.started = false + c.initialized = false + if c.initializing && c.initDone != nil { + close(c.initDone) + } + c.initializing = false + c.initDone = nil c.exitErr = err if c.exited != nil { close(c.exited) diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index 3a352d3a..a54757e6 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -276,7 +276,7 @@ func newTestStdIOClient(t *testing.T) *StdIOClient { client, err := NewStdIOClient(StdioClientConfig{ Command: os.Args[0], Args: []string{"-test.run=TestHelperProcessMCPStdioServer", "--"}, - Env: []string{"GO_WANT_MCP_STDIO_HELPER=1"}, + Env: []string{"GO_WANT_MCP_STDIO_HELPER=1", "GO_MCP_STDIO_REQUIRE_INITIALIZE=1"}, StartTimeout: 3 * time.Second, CallTimeout: 3 * time.Second, }) @@ -286,11 +286,36 @@ func newTestStdIOClient(t *testing.T) *StdIOClient { 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"}, + 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" + initialized := !requireInitialize + reader := bufio.NewReader(os.Stdin) for { payload, err := readFramedMessage(reader) @@ -311,7 +336,45 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { 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, @@ -329,6 +392,17 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { }, } 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{ From 48530c6fb21d411a4cd1f5dbf24589c2167b852f Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 17:49:22 +0800 Subject: [PATCH 09/13] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BAMCP=20stdio?= =?UTF-8?q?=E6=8F=A1=E6=89=8B=E5=85=BC=E5=AE=B9=E5=B9=B6=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E8=A6=86=E7=9B=96=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/mcp/stdio_client.go | 247 +++++++++++++++++++++--- internal/tools/mcp/stdio_client_test.go | 132 ++++++++++++- 2 files changed, 347 insertions(+), 32 deletions(-) diff --git a/internal/tools/mcp/stdio_client.go b/internal/tools/mcp/stdio_client.go index 71d28d3c..b6808d51 100644 --- a/internal/tools/mcp/stdio_client.go +++ b/internal/tools/mcp/stdio_client.go @@ -23,6 +23,7 @@ const ( 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" @@ -88,9 +89,18 @@ type StdIOClient struct { 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) == "" { @@ -132,7 +142,7 @@ func (c *StdIOClient) Close() error { return nil } -// ListTools 调用 MCP `tools/list` 获取工具清单。 +// ListTools 调用 MCP tools/list 获取工具清单。 func (c *StdIOClient) ListTools(ctx context.Context) ([]ToolDescriptor, error) { callCtx, cancel := c.callContext(ctx) defer cancel() @@ -169,7 +179,7 @@ func (c *StdIOClient) ListTools(ctx context.Context) ([]ToolDescriptor, error) { return result, nil } -// CallTool 调用 MCP `tools/call` 并收敛返回值。 +// CallTool 调用 MCP tools/call 并收敛返回值。 func (c *StdIOClient) CallTool(ctx context.Context, toolName string, arguments []byte) (CallResult, error) { trimmedToolName := strings.TrimSpace(toolName) if trimmedToolName == "" { @@ -196,12 +206,13 @@ func (c *StdIOClient) CallTool(ctx context.Context, toolName string, arguments [ return decodeCallResult(raw), nil } -// HealthCheck 通过一次短超时 `tools/list` 验证连接可用性。 +// 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 { @@ -212,12 +223,24 @@ func (c *StdIOClient) callContext(ctx context.Context) (context.Context, context 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 用于初始化阶段避免递归。 +// 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 } @@ -240,6 +263,7 @@ func (c *StdIOClient) callRequest(ctx context.Context, method string, params any } c.pending[requestID] = replyCh stdin := c.stdin + selectedProtocol := c.resolveWriteProtocolLocked(override) c.mu.Unlock() if stdin == nil { c.removePending(requestID) @@ -257,7 +281,7 @@ func (c *StdIOClient) callRequest(ctx context.Context, method string, params any return nil, fmt.Errorf("mcp: marshal request: %w", err) } c.writeMu.Lock() - writeErr := writeFramedMessage(stdin, requestPayload) + writeErr := writeMessageWithProtocol(stdin, requestPayload, selectedProtocol) c.writeMu.Unlock() if writeErr != nil { c.removePending(requestID) @@ -273,8 +297,19 @@ func (c *StdIOClient) callRequest(ctx context.Context, method string, params any } } -// sendNotification 发送无需响应的 RPC 通知；skipEnsure=true 用于初始化流程。 +// 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 } @@ -293,6 +328,7 @@ func (c *StdIOClient) sendNotification(ctx context.Context, method string, param 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") @@ -308,7 +344,7 @@ func (c *StdIOClient) sendNotification(ctx context.Context, method string, param } c.writeMu.Lock() - writeErr := writeFramedMessage(stdin, payload) + writeErr := writeMessageWithProtocol(stdin, payload, selectedProtocol) c.writeMu.Unlock() if writeErr != nil { return fmt.Errorf("mcp: send notification: %w", writeErr) @@ -316,6 +352,7 @@ func (c *StdIOClient) sendNotification(ctx context.Context, method string, param return nil } +// ensureStarted 确保 stdio 子进程已启动并处于可读写状态。 func (c *StdIOClient) ensureStarted(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() @@ -374,6 +411,7 @@ func (c *StdIOClient) ensureStarted(ctx context.Context) error { c.initialized = false c.initializing = false c.initDone = nil + c.protocol = stdioProtocolUnknown c.backoff = c.cfg.RestartBackoff c.retryAt = time.Time{} @@ -427,7 +465,7 @@ func (c *StdIOClient) ensureInitialized(ctx context.Context) error { } } -// performInitialize 执行标准 MCP 初始化握手：initialize -> notifications/initialized。 +// performInitialize 执行 MCP 握手并自动兼容 line/framed 两种 stdio 线协议。 func (c *StdIOClient) performInitialize(ctx context.Context) error { params := map[string]any{ "protocolVersion": defaultMCPProtocolVersion, @@ -437,23 +475,62 @@ func (c *StdIOClient) performInitialize(ctx context.Context) error { "version": defaultMCPClientVersion, }, } - if _, err := c.callRequest(ctx, "initialize", params, true); err != nil { - return fmt.Errorf("mcp: initialize session: %w", err) + + protocols := []stdioProtocol{stdioProtocolLine, stdioProtocolFramed} + c.mu.Lock() + if c.protocol == stdioProtocolLine || c.protocol == stdioProtocolFramed { + protocols = []stdioProtocol{c.protocol} } - if err := c.sendNotification(ctx, "notifications/initialized", map[string]any{}, true); err != nil { - return fmt.Errorf("mcp: notify initialized: %w", err) + 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 nil + + return fmt.Errorf("mcp: initialize session: %s", strings.Join(errs, "; ")) } +// readLoop 持续消费 MCP server 响应并分发给对应 pending 请求。 func (c *StdIOClient) readLoop() { for { - message, err := readFramedMessage(c.reader) + 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 @@ -482,6 +559,7 @@ func (c *StdIOClient) readLoop() { } } +// waitLoop 等待子进程退出并触发统一下线处理。 func (c *StdIOClient) waitLoop(command *exec.Cmd) { if command == nil { c.markExited(errors.New("mcp: stdio process is nil")) @@ -491,6 +569,7 @@ func (c *StdIOClient) waitLoop(command *exec.Cmd) { c.markExited(fmt.Errorf("mcp: stdio process exited: %w", err)) } +// markExited 将客户端状态原子切换为已下线并唤醒所有等待请求。 func (c *StdIOClient) markExited(err error) { c.mu.Lock() defer c.mu.Unlock() @@ -505,6 +584,7 @@ func (c *StdIOClient) markExited(err error) { } c.initializing = false c.initDone = nil + c.protocol = stdioProtocolUnknown c.exitErr = err if c.exited != nil { close(c.exited) @@ -517,12 +597,25 @@ func (c *StdIOClient) markExited(err error) { 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} @@ -530,6 +623,7 @@ func (c *StdIOClient) failAllPendingLocked(err error) { } } +// bumpBackoffLocked 按指数退避策略更新下次可重启时间，调用方需持有 c.mu。 func (c *StdIOClient) bumpBackoffLocked() { if c.backoff <= 0 { c.backoff = c.cfg.RestartBackoff @@ -541,6 +635,27 @@ func (c *StdIOClient) bumpBackoffLocked() { } } +// 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 { @@ -552,28 +667,85 @@ func writeFramedMessage(writer io.Writer, payload []byte) error { return nil } -func readFramedMessage(reader *bufio.Reader) ([]byte, error) { - contentLength := -1 +// 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, err + return nil, stdioProtocolUnknown, err } + trimmed := strings.TrimSpace(line) if trimmed == "" { - 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, fmt.Errorf("mcp: invalid content-length %q", rawLength) + 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)) } - contentLength = length + 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") } @@ -588,6 +760,35 @@ func readFramedMessage(reader *bufio.Reader) ([]byte, error) { 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 { diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index a54757e6..ccfd50d5 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "os" + "strconv" "strings" "sync" "testing" @@ -23,7 +24,7 @@ func (errWriter) Write(p []byte) (int, error) { func TestStdIOClientListToolsAndCallTool(t *testing.T) { t.Parallel() - client := newTestStdIOClient(t) + client := newTestStdIOClientWithMode(t, "framed") defer func() { _ = client.Close() }() toolsList, err := client.ListTools(context.Background()) @@ -43,13 +44,50 @@ func TestStdIOClientListToolsAndCallTool(t *testing.T) { } } +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 := newTestStdIOClient(t) + client := newTestStdIOClientWithMode(t, "framed") defer func() { _ = client.Close() }() - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() if err := client.HealthCheck(ctx); err != nil { t.Fatalf("HealthCheck() error = %v", err) @@ -59,7 +97,7 @@ func TestStdIOClientHealthCheck(t *testing.T) { func TestStdIOClientConcurrentCallTool(t *testing.T) { t.Parallel() - client := newTestStdIOClient(t) + client := newTestStdIOClientWithMode(t, "framed") defer func() { _ = client.Close() }() const workers = 16 @@ -101,6 +139,40 @@ func TestReadFramedMessageRejectsOversizedPayload(t *testing.T) { } } +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() @@ -176,7 +248,7 @@ 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") { + 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) } @@ -186,6 +258,21 @@ func TestReadFramedMessageHeaderErrors(t *testing.T) { } } +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() @@ -271,12 +358,20 @@ func TestStdIOClientBumpBackoffClamp(t *testing.T) { } 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"}, + 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, }) @@ -292,7 +387,7 @@ func TestStdIOClientInitializeFailure(t *testing.T) { 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"}, + 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, }) @@ -314,11 +409,24 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { 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 { - payload, err := readFramedMessage(reader) + 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) @@ -428,7 +536,13 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { if err != nil { os.Exit(4) } - if err := writeFramedMessage(os.Stdout, rawResponse); err != nil { + switch wireMode { + case "line": + err = writeLineMessage(os.Stdout, rawResponse) + default: + err = writeFramedMessage(os.Stdout, rawResponse) + } + if err != nil { os.Exit(5) } } From 5e099a583918894d38f07c81b72e0445178c1c50 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:25:35 +0800 Subject: [PATCH 10/13] =?UTF-8?q?test:=20=E8=A1=A5=E9=BD=90MCP=E6=8E=A5?= =?UTF-8?q?=E5=85=A5=E8=B7=AF=E5=BE=84=E8=A6=86=E7=9B=96=E7=94=A8=E4=BE=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/app/bootstrap_test.go | 270 +++++++++++++++++++++++- internal/tools/mcp/stdio_client_test.go | 66 ++++++ 2 files changed, 335 insertions(+), 1 deletion(-) diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 67e7ba2a..d706a1b2 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -1,13 +1,18 @@ package app import ( + "bufio" + "bytes" "context" "encoding/json" "errors" + "fmt" + "io" "net/http" "net/http/httptest" "os" "path/filepath" + "strconv" "strings" "testing" "time" @@ -159,6 +164,115 @@ func TestBuildMCPRegistryFromConfig(t *testing.T) { } } +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() @@ -398,10 +512,14 @@ func disableBuiltinProviderAPIKeys(t *testing.T) { } type stubMCPServerClient struct { - tools []mcp.ToolDescriptor + 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 } @@ -412,3 +530,153 @@ func (s *stubMCPServerClient) CallTool(ctx context.Context, toolName string, arg 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/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index ccfd50d5..01d8d20b 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -2,6 +2,7 @@ package mcp import ( "bufio" + "bytes" "context" "encoding/json" "errors" @@ -15,6 +16,12 @@ import ( "time" ) +type nopWriteCloser struct { + bytes.Buffer +} + +func (n *nopWriteCloser) Close() error { return nil } + type errWriter struct{} func (errWriter) Write(p []byte) (int, error) { @@ -94,6 +101,65 @@ func TestStdIOClientHealthCheck(t *testing.T) { } } +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() From e726eb338fec1f911167798f5768de0ead3b369d Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:48:22 +0800 Subject: [PATCH 11/13] =?UTF-8?q?fix:=20=E8=A7=84=E8=8C=83registry?= =?UTF-8?q?=E4=B8=AD=E7=9A=84provider=E5=BC=95=E7=94=A8=E4=BB=A5=E4=BF=AE?= =?UTF-8?q?=E5=A4=8DCI=E6=9E=84=E5=BB=BA=E5=BC=82=E5=B8=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/registry.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/registry.go b/internal/tools/registry.go index c03bfcb6..a392f629 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -114,7 +114,7 @@ func (r *Registry) ListAvailableSpecs(ctx context.Context, input SpecListInput) return nil, err } for _, adapter := range mcpAdapters { - specs = append(specs, provider.ToolSpec{ + specs = append(specs, providertypes.ToolSpec{ Name: adapter.FullName(), Description: adapter.Description(), Schema: adapter.Schema(), From 3cf3ea21f0b782b608914fd820302c223e61cf5c Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:51:42 +0800 Subject: [PATCH 12/13] =?UTF-8?q?chore:=20=E8=A7=A6=E5=8F=91CI=E9=87=8D?= =?UTF-8?q?=E8=B7=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From 704b018b8cbef0e7ee832712d99b5d4cd8434795 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Tue, 7 Apr 2026 20:43:01 +0800 Subject: [PATCH 13/13] =?UTF-8?q?test:=20=E6=8F=90=E5=8D=87MCP=E4=B8=8E?= =?UTF-8?q?=E5=BC=95=E5=AF=BC=E6=B5=81=E7=A8=8B=E8=BE=B9=E7=95=8C=E8=A6=86?= =?UTF-8?q?=E7=9B=96=E7=8E=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/app/bootstrap_test.go | 83 ++++++++++++++++++++++++++ internal/tools/mcp/adapter_test.go | 34 +++++++++++ internal/tools/mcp/registry_test.go | 90 +++++++++++++++++++++++++++++ 3 files changed, 207 insertions(+) diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index d706a1b2..a81903c0 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -351,6 +351,89 @@ func TestResolveMCPServerEnvAndWorkdir(t *testing.T) { } } +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() diff --git a/internal/tools/mcp/adapter_test.go b/internal/tools/mcp/adapter_test.go index 8f3adc8f..2b38691f 100644 --- a/internal/tools/mcp/adapter_test.go +++ b/internal/tools/mcp/adapter_test.go @@ -39,6 +39,19 @@ func TestAdapterFactoryBuildAdapters(t *testing.T) { } } +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() @@ -193,3 +206,24 @@ func TestAdapterCallBoundary(t *testing.T) { 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_test.go b/internal/tools/mcp/registry_test.go index 49f2c654..30c5cd3f 100644 --- a/internal/tools/mcp/registry_test.go +++ b/internal/tools/mcp/registry_test.go @@ -244,3 +244,93 @@ func TestRegistrySetServerStatusValidation(t *testing.T) { 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"]) + } +}