Add files via upload

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, "\"")

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")
+}