diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 955a617c..d7dea0ad 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -167,8 +167,6 @@ 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"}}, @@ -323,8 +321,6 @@ func TestDefaultRegisterMCPStdioServerRefreshFailure(t *testing.T) { } func TestBuildToolRegistryIncludesMCPFromConfig(t *testing.T) { - t.Parallel() - cfg := config.StaticDefaults().Clone() cfg.Workdir = t.TempDir() cfg.Tools.MCP.Servers = []config.MCPServerConfig{ @@ -373,8 +369,6 @@ func TestBuildToolRegistryIncludesMCPFromConfig(t *testing.T) { } func TestBuildToolRegistryAppliesMCPExposureConfig(t *testing.T) { - t.Parallel() - cfg := config.StaticDefaults().Clone() cfg.Workdir = t.TempDir() cfg.Tools.MCP.Servers = []config.MCPServerConfig{ @@ -551,8 +545,6 @@ func TestBuildMCPRegistryNoEnabledServerReturnsNil(t *testing.T) { } func TestBuildMCPRegistryRegisterError(t *testing.T) { - t.Parallel() - cfg := config.StaticDefaults().Clone() cfg.Workdir = t.TempDir() cfg.Tools.MCP.Servers = []config.MCPServerConfig{ @@ -572,8 +564,6 @@ func TestBuildMCPRegistryRegisterError(t *testing.T) { } func TestBuildMCPRegistryRollbackRegisteredServersOnFailure(t *testing.T) { - t.Parallel() - cfg := config.StaticDefaults().Clone() cfg.Workdir = t.TempDir() cfg.Tools.MCP.Servers = []config.MCPServerConfig{ diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index 06538080..0ad9103a 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -43,6 +43,7 @@ func TestNewRootCommandAllowsEmptyWorkdir(t *testing.T) { } cmd := NewRootCommand() + cmd.SetArgs([]string{}) if err := cmd.ExecuteContext(context.Background()); err != nil { t.Fatalf("ExecuteContext() error = %v", err) } @@ -61,6 +62,7 @@ func TestNewRootCommandReturnsLauncherError(t *testing.T) { } cmd := NewRootCommand() + cmd.SetArgs([]string{}) err := cmd.ExecuteContext(context.Background()) if !errors.Is(err, expected) { t.Fatalf("expected launcher error %v, got %v", expected, err) diff --git a/internal/runtime/permission_test.go b/internal/runtime/permission_test.go index 9895f4da..447924a4 100644 --- a/internal/runtime/permission_test.go +++ b/internal/runtime/permission_test.go @@ -11,8 +11,38 @@ import ( approvalflow "neo-code/internal/runtime/approval" "neo-code/internal/security" "neo-code/internal/tools" + "neo-code/internal/tools/mcp" ) +type runtimeStubMCPClient struct { + tools []mcp.ToolDescriptor + callResult mcp.CallResult + callErr error + callCount int +} + +func (s *runtimeStubMCPClient) ListTools(ctx context.Context) ([]mcp.ToolDescriptor, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + return append([]mcp.ToolDescriptor(nil), s.tools...), nil +} + +func (s *runtimeStubMCPClient) CallTool(ctx context.Context, toolName string, arguments []byte) (mcp.CallResult, error) { + if err := ctx.Err(); err != nil { + return mcp.CallResult{}, err + } + s.callCount++ + if s.callErr != nil { + return mcp.CallResult{}, s.callErr + } + return s.callResult, nil +} + +func (s *runtimeStubMCPClient) HealthCheck(ctx context.Context) error { + return ctx.Err() +} + func TestResolvePermissionValidation(t *testing.T) { t.Parallel() @@ -226,6 +256,366 @@ waitRequest: } } +func TestServiceRunMCPPermissionAllowFlow(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + mcpRegistry := mcp.NewRegistry() + mcpClient := &runtimeStubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "create_issue", Description: "create issue", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: mcp.CallResult{Content: "mcp create ok"}, + } + if err := mcpRegistry.RegisterServer("github", "stdio", "v1", mcpClient); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "github"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + engine, err := security.NewPolicyEngine(security.DecisionAllow, []security.PolicyRule{ + { + ID: "ask-github-create", + Priority: 720, + Decision: security.DecisionAsk, + Reason: "mcp create requires approval", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.github.create_issue"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + }) + if err != nil { + t.Fatalf("new policy engine: %v", err) + } + toolManager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + + scripted := &scriptedProvider{ + responses: []scriptedResponse{ + { + Message: providertypes.Message{ + Role: "assistant", + ToolCalls: []providertypes.ToolCall{ + {ID: "call-mcp-allow", Name: "mcp.github.create_issue", Arguments: `{"title":"hello"}`}, + }, + }, + FinishReason: "tool_calls", + }, + { + Message: providertypes.Message{Role: "assistant", Content: "done"}, + FinishReason: "stop", + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) + runErrCh := make(chan error, 1) + go func() { + runErrCh <- service.Run(context.Background(), UserInput{RunID: "run-mcp-permission-allow", Content: "create issue"}) + }() + + var requestPayload PermissionRequestPayload +waitRequest: + for { + select { + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting permission request") + case event := <-service.Events(): + if event.Type != EventPermissionRequest { + continue + } + payload, ok := event.Payload.(PermissionRequestPayload) + if !ok { + t.Fatalf("expected permission request payload, got %#v", event.Payload) + } + requestPayload = payload + break waitRequest + } + } + + if requestPayload.ToolName != "mcp.github.create_issue" || + requestPayload.ToolCategory != "mcp.github" || + requestPayload.RuleID != "ask-github-create" || + requestPayload.Reason != "mcp create requires approval" || + requestPayload.Decision != "ask" { + t.Fatalf("unexpected permission request payload: %+v", requestPayload) + } + if requestPayload.RememberScope != "" { + t.Fatalf("expected empty request remember scope, got %+v", requestPayload) + } + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: requestPayload.RequestID, + Decision: approvalflow.DecisionAllowSession, + }); err != nil { + t.Fatalf("ResolvePermission() error = %v", err) + } + if err := <-runErrCh; err != nil { + t.Fatalf("Run() error = %v", err) + } + if mcpClient.callCount != 1 { + t.Fatalf("expected MCP tool to execute once, got %d", mcpClient.callCount) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventPermissionResolved, EventToolResult, EventAgentDone}) + + foundResolved := false + for _, event := range events { + if event.Type != EventPermissionResolved { + continue + } + payload, ok := event.Payload.(PermissionResolvedPayload) + if !ok { + t.Fatalf("expected PermissionResolvedPayload, got %#v", event.Payload) + } + if payload.ToolName != "mcp.github.create_issue" || + payload.ToolCategory != "mcp.github" || + payload.RuleID != "ask-github-create" || + payload.Reason != "permission approved by user" || + payload.Decision != "allow" || + payload.ResolvedAs != "approved" || + payload.RememberScope != string(tools.SessionPermissionScopeAlways) { + t.Fatalf("unexpected permission resolved payload: %+v", payload) + } + foundResolved = true + } + if !foundResolved { + t.Fatalf("expected permission resolved event") + } +} + +func TestServiceRunMCPPermissionRejectFlow(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + mcpRegistry := mcp.NewRegistry() + mcpClient := &runtimeStubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "create_issue", Description: "create issue", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: mcp.CallResult{Content: "should-not-run"}, + } + if err := mcpRegistry.RegisterServer("github", "stdio", "v1", mcpClient); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "github"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + engine, err := security.NewPolicyEngine(security.DecisionAllow, []security.PolicyRule{ + { + ID: "ask-github-create", + Priority: 720, + Decision: security.DecisionAsk, + Reason: "mcp create requires approval", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.github.create_issue"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + }) + if err != nil { + t.Fatalf("new policy engine: %v", err) + } + toolManager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + + scripted := &scriptedProvider{ + responses: []scriptedResponse{ + { + Message: providertypes.Message{ + Role: "assistant", + ToolCalls: []providertypes.ToolCall{ + {ID: "call-mcp-reject", Name: "mcp.github.create_issue", Arguments: `{"title":"hello"}`}, + }, + }, + FinishReason: "tool_calls", + }, + { + Message: providertypes.Message{Role: "assistant", Content: "done"}, + FinishReason: "stop", + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) + runErrCh := make(chan error, 1) + go func() { + runErrCh <- service.Run(context.Background(), UserInput{RunID: "run-mcp-permission-reject", Content: "create issue"}) + }() + + var requestPayload PermissionRequestPayload +waitRequest: + for { + select { + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting permission request") + case event := <-service.Events(): + if event.Type != EventPermissionRequest { + continue + } + payload, ok := event.Payload.(PermissionRequestPayload) + if !ok { + t.Fatalf("expected permission request payload, got %#v", event.Payload) + } + requestPayload = payload + break waitRequest + } + } + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: requestPayload.RequestID, + Decision: approvalflow.DecisionReject, + }); err != nil { + t.Fatalf("ResolvePermission() error = %v", err) + } + if err := <-runErrCh; err != nil { + t.Fatalf("Run() error = %v", err) + } + if mcpClient.callCount != 0 { + t.Fatalf("expected rejected MCP tool not to execute, got %d", mcpClient.callCount) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventPermissionResolved, EventToolResult, EventAgentDone}) + + foundResolved := false + for _, event := range events { + if event.Type != EventPermissionResolved { + continue + } + payload, ok := event.Payload.(PermissionResolvedPayload) + if !ok { + t.Fatalf("expected PermissionResolvedPayload, got %#v", event.Payload) + } + if payload.ToolName != "mcp.github.create_issue" || + payload.RuleID != "ask-github-create" || + payload.Decision != "deny" || + payload.ResolvedAs != "rejected" || + payload.RememberScope != string(tools.SessionPermissionScopeReject) { + t.Fatalf("unexpected permission resolved payload: %+v", payload) + } + foundResolved = true + } + if !foundResolved { + t.Fatalf("expected permission resolved event") + } +} + +func TestServiceRunMCPPermissionHardDenyFlow(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + mcpRegistry := mcp.NewRegistry() + mcpClient := &runtimeStubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "create_issue", Description: "create issue", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: mcp.CallResult{Content: "should-not-run"}, + } + if err := mcpRegistry.RegisterServer("github", "stdio", "v1", mcpClient); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "github"); err != nil { + t.Fatalf("refresh mcp tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + engine, err := security.NewPolicyEngine(security.DecisionAllow, []security.PolicyRule{ + { + ID: "deny-github-server", + Priority: 830, + Decision: security.DecisionDeny, + Reason: "github mcp server denied", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ToolCategories: []string{"mcp.github"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + { + ID: "ask-github-create", + Priority: 720, + Decision: security.DecisionAsk, + Reason: "mcp create requires approval", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.github.create_issue"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + }) + if err != nil { + t.Fatalf("new policy engine: %v", err) + } + toolManager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + + scripted := &scriptedProvider{ + responses: []scriptedResponse{ + { + Message: providertypes.Message{ + Role: "assistant", + ToolCalls: []providertypes.ToolCall{ + {ID: "call-mcp-deny", Name: "mcp.github.create_issue", Arguments: `{"title":"hello"}`}, + }, + }, + FinishReason: "tool_calls", + }, + { + Message: providertypes.Message{Role: "assistant", Content: "done"}, + FinishReason: "stop", + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) + if err := service.Run(context.Background(), UserInput{RunID: "run-mcp-permission-deny", Content: "create issue"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + if mcpClient.callCount != 0 { + t.Fatalf("expected hard denied MCP tool not to execute, got %d", mcpClient.callCount) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventPermissionResolved, EventToolResult, EventAgentDone}) + assertNoEventType(t, events, EventPermissionRequest) + + foundResolved := false + for _, event := range events { + if event.Type != EventPermissionResolved { + continue + } + payload, ok := event.Payload.(PermissionResolvedPayload) + if !ok { + t.Fatalf("expected PermissionResolvedPayload, got %#v", event.Payload) + } + if payload.ToolName != "mcp.github.create_issue" || + payload.RuleID != "deny-github-server" || + payload.Reason != "github mcp server denied" || + payload.Decision != "deny" || + payload.ResolvedAs != "denied" || + payload.RememberScope != "" { + t.Fatalf("unexpected permission resolved payload: %+v", payload) + } + foundResolved = true + } + if !foundResolved { + t.Fatalf("expected permission resolved event") + } +} + func TestPermissionHelpers(t *testing.T) { t.Parallel() @@ -415,6 +805,7 @@ func TestExecuteToolCallWithPermissionDoesNotRecheckContextAfterSuccessfulEmit(t ) ctx, cancel := context.WithCancel(context.Background()) + defer cancel() service.events = make(chan RuntimeEvent, 1) result, execErr := service.executeToolCallWithPermission(ctx, permissionExecutionInput{ diff --git a/internal/tools/manager_test.go b/internal/tools/manager_test.go index 757794d4..0004b079 100644 --- a/internal/tools/manager_test.go +++ b/internal/tools/manager_test.go @@ -1203,6 +1203,108 @@ func TestDefaultManagerExecuteMCPServerDenyUsesTraceableRule(t *testing.T) { } } +func TestDefaultManagerExecuteMCPServerDenyPriorityOverridesToolRules(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + mcpRegistry := mcp.NewRegistry() + client := &stubMCPClient{ + tools: []mcp.ToolDescriptor{ + {Name: "create_issue", Description: "create"}, + {Name: "list_issues", Description: "list"}, + {Name: "search", Description: "search"}, + }, + callResult: mcp.CallResult{Content: "ok"}, + } + if err := mcpRegistry.RegisterServer("github", "stdio", "v1", client); err != nil { + t.Fatalf("register mcp server: %v", err) + } + if err := mcpRegistry.RegisterServer("docs", "stdio", "v1", client); err != nil { + t.Fatalf("register docs server: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "github"); err != nil { + t.Fatalf("refresh github tools: %v", err) + } + if err := mcpRegistry.RefreshServerTools(context.Background(), "docs"); err != nil { + t.Fatalf("refresh docs tools: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + engine, err := security.NewPolicyEngine(security.DecisionAllow, []security.PolicyRule{ + { + ID: "deny-github-server", + Priority: 830, + Decision: security.DecisionDeny, + Reason: "github server denied", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ToolCategories: []string{"mcp.github"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + { + ID: "allow-github-create", + Priority: 700, + Decision: security.DecisionAllow, + Reason: "github create allowed", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.github.create_issue"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + { + ID: "ask-github-list", + Priority: 720, + Decision: security.DecisionAsk, + Reason: "github list requires approval", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.github.list_issues"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + { + ID: "allow-docs-search", + Priority: 700, + Decision: security.DecisionAllow, + Reason: "docs search allowed", + ActionTypes: []security.ActionType{security.ActionTypeMCP}, + ResourcePatterns: []string{"mcp.docs.search"}, + TargetTypes: []security.TargetType{security.TargetTypeMCP}, + }, + }) + if err != nil { + t.Fatalf("new policy engine: %v", err) + } + + manager, err := NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + for _, input := range []ToolCallInput{ + {ID: "call-github-create", Name: "mcp.github.create_issue", Arguments: []byte(`{"title":"hello"}`), SessionID: "session-priority"}, + {ID: "call-github-list", Name: "mcp.github.list_issues", Arguments: []byte(`{"state":"open"}`), SessionID: "session-priority"}, + } { + _, execErr := manager.Execute(context.Background(), input) + var permissionErr *PermissionDecisionError + if !errors.As(execErr, &permissionErr) { + t.Fatalf("expected permission error for %s, got %v", input.Name, execErr) + } + if permissionErr.Decision() != "deny" || permissionErr.RuleID() != "deny-github-server" { + t.Fatalf("expected server-level deny for %s, got decision=%q rule=%q", input.Name, permissionErr.Decision(), permissionErr.RuleID()) + } + } + + result, execErr := manager.Execute(context.Background(), ToolCallInput{ + ID: "call-docs-search", + Name: "mcp.docs.search", + Arguments: []byte(`{"query":"neo-code"}`), + SessionID: "session-priority", + }) + if execErr != nil { + t.Fatalf("expected docs search allow, got %v", execErr) + } + if result.Content != "ok" { + t.Fatalf("expected docs search to execute, got %+v", result) + } +} + func TestNoopWorkspaceSandbox(t *testing.T) { t.Parallel() diff --git a/internal/tools/mcp/registry_test.go b/internal/tools/mcp/registry_test.go index 249e2505..13d34179 100644 --- a/internal/tools/mcp/registry_test.go +++ b/internal/tools/mcp/registry_test.go @@ -172,6 +172,72 @@ func TestRegistryConcurrentSnapshotAndRefresh(t *testing.T) { } } +func TestRegistryConcurrentRefreshAndCall(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + client := &stubServerClient{ + tools: []ToolDescriptor{ + {Name: "search", InputSchema: map[string]any{"type": "object"}}, + }, + callResult: CallResult{Content: "ok"}, + } + 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) + } + + var wg sync.WaitGroup + errCh := make(chan error, 32) + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 16; j++ { + if err := registry.RefreshServerTools(context.Background(), "server-1"); err != nil { + errCh <- err + return + } + } + }() + } + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 16; j++ { + result, err := registry.Call(context.Background(), "server-1", "search", []byte(`{"q":"neo"}`)) + if err != nil { + errCh <- err + return + } + if result.Content != "ok" { + errCh <- errors.New("unexpected call result") + return + } + } + }() + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatalf("concurrent refresh and call timed out") + } + close(errCh) + for err := range errCh { + t.Fatalf("concurrent registry operation failed: %v", err) + } +} + func TestRegistrySnapshotSchemaIsDeepCloned(t *testing.T) { t.Parallel() diff --git a/internal/tools/mcp/stdio_client_test.go b/internal/tools/mcp/stdio_client_test.go index 07163406..146bbd13 100644 --- a/internal/tools/mcp/stdio_client_test.go +++ b/internal/tools/mcp/stdio_client_test.go @@ -516,6 +516,56 @@ func TestStdIOClientInitializeFailure(t *testing.T) { } } +func TestStdIOClientListToolsTimeout(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_REQUIRE_INITIALIZE=1", "GO_MCP_STDIO_WIRE=framed", "GO_MCP_STDIO_LIST_DELAY_MS=300"}, + StartTimeout: 3 * time.Second, + CallTimeout: 100 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewStdIOClient() error = %v", err) + } + defer func() { _ = client.Close() }() + + _, callErr := client.ListTools(context.Background()) + if callErr == nil { + t.Fatal("expected timeout error, got nil") + } + errMsg := callErr.Error() + if !strings.Contains(errMsg, context.DeadlineExceeded.Error()) { + t.Fatalf("expected deadline exceeded error, got %v", callErr) + } + if strings.Contains(errMsg, "initialize session") || strings.Contains(errMsg, "tools/list") || errMsg == context.DeadlineExceeded.Error() { + return + } + t.Fatalf("expected initialize session or tools/list timeout path, got %v", callErr) +} + +func TestStdIOClientListToolsProtocolError(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_REQUIRE_INITIALIZE=1", "GO_MCP_STDIO_WIRE=framed", "GO_MCP_STDIO_LIST_MALFORMED=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(), "decode tools/list result") { + t.Fatalf("expected decode tools/list result error, got %v", callErr) + } +} + func TestHelperProcessMCPStdioServer(t *testing.T) { if os.Getenv("GO_WANT_MCP_STDIO_HELPER") != "1" { return @@ -523,6 +573,8 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { requireInitialize := os.Getenv("GO_MCP_STDIO_REQUIRE_INITIALIZE") == "1" initFail := os.Getenv("GO_MCP_STDIO_INIT_FAIL") == "1" + listDelayMS, _ := strconv.Atoi(strings.TrimSpace(os.Getenv("GO_MCP_STDIO_LIST_DELAY_MS"))) + listMalformed := os.Getenv("GO_MCP_STDIO_LIST_MALFORMED") == "1" wireMode := strings.TrimSpace(os.Getenv("GO_MCP_STDIO_WIRE")) if wireMode == "" { wireMode = "framed" @@ -597,6 +649,17 @@ func TestHelperProcessMCPStdioServer(t *testing.T) { } break } + if listDelayMS > 0 { + time.Sleep(time.Duration(listDelayMS) * time.Millisecond) + } + if listMalformed { + response = map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "result": "broken", + } + break + } response = map[string]any{ "jsonrpc": "2.0", "id": requestID, diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index cc3a2028..59ddb78f 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -745,3 +745,54 @@ func TestParseMCPToolFullName(t *testing.T) { t.Fatalf("expected parse failure for incomplete name") } } + +func TestRegistryMCPFaultDoesNotAffectBuiltInTools(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + registry.Register(stubTool{ + name: "builtin_tool", + description: "built in", + schema: map[string]any{"type": "object"}, + result: ToolResult{ + Name: "builtin_tool", + Content: "builtin ok", + }, + }) + + 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) + } + if err := mcpRegistry.SetServerStatus("docs", mcp.ServerStatusOffline); err != nil { + t.Fatalf("set server offline: %v", err) + } + registry.SetMCPRegistry(mcpRegistry) + + specs, err := registry.ListAvailableSpecs(context.Background(), SpecListInput{}) + if err != nil { + t.Fatalf("ListAvailableSpecs() error = %v", err) + } + if len(specs) != 1 || specs[0].Name != "builtin_tool" { + t.Fatalf("expected built-in tool only, got %+v", specs) + } + + result, err := registry.Execute(context.Background(), ToolCallInput{ + ID: "builtin-call", + Name: "builtin_tool", + }) + if err != nil { + t.Fatalf("builtin Execute() error = %v", err) + } + if result.Content != "builtin ok" { + t.Fatalf("expected builtin execution success, got %+v", result) + } +} diff --git a/internal/tui/core/app/update.go b/internal/tui/core/app/update.go index c117524f..d133279d 100644 --- a/internal/tui/core/app/update.go +++ b/internal/tui/core/app/update.go @@ -722,7 +722,7 @@ func (a *App) syncConfigState(cfg config.Config) { } } -// refreshRuntimeSourceSnapshot 浠?runtime 鏌ヨ context/token/tool 蹇収锛岀敤浜庝細璇濆垏鎹㈡垨鎭㈠鏃跺洖濉?UI銆 +// refreshRuntimeSourceSnapshot 从 runtime 查询 context/token/tool 快照，用于会话切换或恢复时回填 UI。 func (a *App) refreshRuntimeSourceSnapshot() { sessionID := strings.TrimSpace(a.state.ActiveSessionID) if sessionID != "" { @@ -1791,7 +1791,7 @@ func runResolvePermission( ) } -// runCompact 鍦ㄧ嫭绔嬪懡浠や腑瑙﹀彂 runtime compact锛屽苟鎶婄粨鏋滃洖浼犵粰 TUI銆 +// runCompact 在独立命令中触发 runtime compact，并把结果回传给 TUI。 func runCompact(runtime agentruntime.Runtime, sessionID string) tea.Cmd { return tuiservices.RunCompactCmd( runtime, @@ -1800,7 +1800,7 @@ func runCompact(runtime agentruntime.Runtime, sessionID string) tea.Cmd { ) } -// isBusy 缁熶竴鍒ゆ柇褰撳墠鐣岄潰鏄惁瀛樺湪杩涜涓殑 agent 鎴?compact 鎿嶄綔銆 +// isBusy 统一判断当前界面是否存在进行中的 agent 或 compact 操作。 func (a App) isBusy() bool { return tuiutils.IsBusy(a.state.IsAgentRunning, a.state.IsCompacting) }