diff --git a/internal/mcp/run_context.go b/internal/mcp/run_context.go index 48dac642..d858c2ce 100644 --- a/internal/mcp/run_context.go +++ b/internal/mcp/run_context.go @@ -11,7 +11,16 @@ type ToolRunRegistry interface { UnregisterRunningTool(conversationID, executionID string) } +// EinoExecuteRunRegistry 登记进行中的 Eino filesystem execute,供「中断并继续」终止 amass 等长命令。 +type EinoExecuteRunRegistry interface { + RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc) + UnregisterActiveEinoExecute(conversationID string) + AbortActiveEinoExecute(conversationID, note string) bool + TakeEinoExecuteAbortNote(conversationID string) string +} + type toolRunRegistryCtxKey struct{} +type einoExecuteRunRegistryCtxKey struct{} type mcpConversationIDCtxKey struct{} // WithToolRunRegistry 将登记器注入 ctx(Eino / 原生 Agent 任务 ctx)。 @@ -31,6 +40,23 @@ func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry { return v } +// WithEinoExecuteRunRegistry 将 Eino execute 取消登记器注入 ctx。 +func WithEinoExecuteRunRegistry(ctx context.Context, reg EinoExecuteRunRegistry) context.Context { + if ctx == nil || reg == nil { + return ctx + } + return context.WithValue(ctx, einoExecuteRunRegistryCtxKey{}, reg) +} + +// EinoExecuteRunRegistryFromContext 取出 Eino execute 登记器(无则 nil)。 +func EinoExecuteRunRegistryFromContext(ctx context.Context) EinoExecuteRunRegistry { + if ctx == nil { + return nil + } + v, _ := ctx.Value(einoExecuteRunRegistryCtxKey{}).(EinoExecuteRunRegistry) + return v +} + // WithMCPConversationID 将对话 ID 注入 ctx,供 CallTool 内与 executionId 关联。 func WithMCPConversationID(ctx context.Context, conversationID string) context.Context { if ctx == nil {