diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 714396db..e21aaa56 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -837,6 +837,46 @@ func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool { return false } +// CancelRunningMCPToolsForConversation cancels all currently running internal/external MCP executions +// owned by the conversation. It is used when a session ends or the user stops a task. +func (a *Agent) CancelRunningMCPToolsForConversation(conversationID, note string) int { + conversationID = strings.TrimSpace(conversationID) + if a == nil || conversationID == "" { + return 0 + } + note = strings.TrimSpace(note) + seen := make(map[string]struct{}) + cancelled := 0 + cancelIfConversationMatches := func(execID string, get func(string) (*mcp.ToolExecution, bool), cancel func(string, string) bool) { + execID = strings.TrimSpace(execID) + if execID == "" { + return + } + if _, ok := seen[execID]; ok { + return + } + seen[execID] = struct{}{} + exec, ok := get(execID) + if !ok || exec == nil || strings.TrimSpace(exec.ConversationID) != conversationID { + return + } + if cancel(execID, note) { + cancelled++ + } + } + if a.mcpServer != nil { + for execID := range a.mcpServer.ActiveRunningExecutionIDs() { + cancelIfConversationMatches(execID, a.mcpServer.GetExecution, a.mcpServer.CancelToolExecutionWithNote) + } + } + if a.externalMCPMgr != nil { + for execID := range a.externalMCPMgr.ActiveRunningExecutionIDs() { + cancelIfConversationMatches(execID, a.externalMCPMgr.GetExecution, a.externalMCPMgr.CancelToolExecutionWithNote) + } + } + return cancelled +} + // extractQuotedToolName 尝试从错误信息中提取被引用的工具名称 func extractQuotedToolName(errMsg string) string { start := strings.Index(errMsg, "\"") diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 1ebe6dff..6e1fa018 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -1,7 +1,10 @@ package agent import ( + "context" + "strings" "testing" + "time" "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/mcp" @@ -65,3 +68,51 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) { t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations) } } + +func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) { + ag := setupTestAgent(t) + ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1) + ag.mcpServer.RegisterTool(mcp.Tool{Name: "block", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + <-ctx.Done() + return nil, ctx.Err() + }) + + ctx1 := mcp.WithMCPConversationID(context.Background(), "conv-1") + result1, execID1, err := ag.mcpServer.CallTool(ctx1, "block", nil) + if err != nil { + t.Fatalf("CallTool conv-1: %v", err) + } + if result1 == nil || !result1.IsError || execID1 == "" { + t.Fatalf("expected bounded wait for conv-1, result=%#v id=%q", result1, execID1) + } + + ctx2 := mcp.WithMCPConversationID(context.Background(), "conv-2") + result2, execID2, err := ag.mcpServer.CallTool(ctx2, "block", nil) + if err != nil { + t.Fatalf("CallTool conv-2: %v", err) + } + if result2 == nil || !result2.IsError || execID2 == "" { + t.Fatalf("expected bounded wait for conv-2, result=%#v id=%q", result2, execID2) + } + + if got := ag.CancelRunningMCPToolsForConversation("conv-1", "session ended"); got != 1 { + t.Fatalf("cancelled count = %d, want 1", got) + } + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + exec1, _ := ag.mcpServer.GetExecution(execID1) + exec2, _ := ag.mcpServer.GetExecution(execID2) + if exec1 != nil && exec1.Status == mcp.ToolExecutionStatusCancelled { + if exec2 == nil || exec2.Status != mcp.ToolExecutionStatusRunning { + t.Fatalf("conv-2 execution should remain running, got %#v", exec2) + } + if !strings.Contains(exec1.Error, "session ended") && (exec1.Result == nil || !strings.Contains(mcp.ToolResultPlainText(exec1.Result), "session ended")) { + t.Fatalf("cancel note missing from conv-1 execution: %#v", exec1) + } + _ = ag.CancelRunningMCPToolsForConversation("conv-2", "") + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("conv-1 execution did not become cancelled") +}