From 1130fdbfa44068c05591f96b0035e99bdb772316 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Fri, 8 May 2026 13:08:45 +0800
Subject: [PATCH] Add files via upload

---
 internal/mcp/external_manager.go |  2 +
 internal/mcp/run_context.go      | 77 ++++++++++++++++++++++++++++++++
 internal/mcp/server.go           |  2 +
 3 files changed, 81 insertions(+)
 create mode 100644 internal/mcp/run_context.go

diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go
index c67c25d9..036f243a 100644
--- a/internal/mcp/external_manager.go
+++ b/internal/mcp/external_manager.go
@@ -458,7 +458,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
 
 	execCtx, runCancel := context.WithCancel(ctx)
 	m.registerRunningCancel(executionID, runCancel)
+	notifyToolRunBegin(ctx, executionID)
 	defer func() {
+		notifyToolRunEnd(ctx, executionID)
 		runCancel()
 		m.unregisterRunningCancel(executionID)
 	}()
diff --git a/internal/mcp/run_context.go b/internal/mcp/run_context.go
new file mode 100644
index 00000000..48dac642
--- /dev/null
+++ b/internal/mcp/run_context.go
@@ -0,0 +1,77 @@
+package mcp
+
+import (
+	"context"
+	"strings"
+)
+
+// ToolRunRegistry 在工具开始/结束时登记当前 executionId，供对话页「仅终止当前工具」与监控页共用取消逻辑。
+type ToolRunRegistry interface {
+	RegisterRunningTool(conversationID, executionID string)
+	UnregisterRunningTool(conversationID, executionID string)
+}
+
+type toolRunRegistryCtxKey struct{}
+type mcpConversationIDCtxKey struct{}
+
+// WithToolRunRegistry 将登记器注入 ctx（Eino / 原生 Agent 任务 ctx）。
+func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
+	if ctx == nil || reg == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, toolRunRegistryCtxKey{}, reg)
+}
+
+// ToolRunRegistryFromContext 取出登记器（无则 nil）。
+func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry {
+	if ctx == nil {
+		return nil
+	}
+	v, _ := ctx.Value(toolRunRegistryCtxKey{}).(ToolRunRegistry)
+	return v
+}
+
+// WithMCPConversationID 将对话 ID 注入 ctx，供 CallTool 内与 executionId 关联。
+func WithMCPConversationID(ctx context.Context, conversationID string) context.Context {
+	if ctx == nil {
+		return nil
+	}
+	id := strings.TrimSpace(conversationID)
+	if id == "" {
+		return ctx
+	}
+	return context.WithValue(ctx, mcpConversationIDCtxKey{}, id)
+}
+
+// MCPConversationIDFromContext 读取对话 ID。
+func MCPConversationIDFromContext(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+	v, _ := ctx.Value(mcpConversationIDCtxKey{}).(string)
+	return v
+}
+
+func notifyToolRunBegin(ctx context.Context, executionID string) {
+	reg := ToolRunRegistryFromContext(ctx)
+	if reg == nil {
+		return
+	}
+	conv := MCPConversationIDFromContext(ctx)
+	if conv == "" || strings.TrimSpace(executionID) == "" {
+		return
+	}
+	reg.RegisterRunningTool(conv, executionID)
+}
+
+func notifyToolRunEnd(ctx context.Context, executionID string) {
+	reg := ToolRunRegistryFromContext(ctx)
+	if reg == nil {
+		return
+	}
+	conv := MCPConversationIDFromContext(ctx)
+	if conv == "" || strings.TrimSpace(executionID) == "" {
+		return
+	}
+	reg.UnregisterRunningTool(conv, executionID)
+}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
index 48b2325d..00170255 100644
--- a/internal/mcp/server.go
+++ b/internal/mcp/server.go
@@ -803,7 +803,9 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
 
 	execCtx, runCancel := context.WithCancel(ctx)
 	s.registerRunningCancel(executionID, runCancel)
+	notifyToolRunBegin(ctx, executionID)
 	defer func() {
+		notifyToolRunEnd(ctx, executionID)
 		runCancel()
 		s.unregisterRunningCancel(executionID)
 	}()