From 8bde9d749a868af0dbfc946ba6e2a546ce0b0d51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 15 Aug 2026 01:47:48 +0800 Subject: [PATCH] Add files via upload --- internal/handler/agent.go | 30 ++++ .../handler/agent_progress_callback_test.go | 47 ++++++ internal/handler/config.go | 40 +++++ .../handler/config_eino_resilience_test.go | 64 ++++++++ internal/handler/conversation.go | 58 ++++++- .../conversation_process_details_test.go | 66 ++++++++ internal/handler/eino_single_agent.go | 14 ++ internal/handler/multi_agent.go | 14 ++ internal/handler/task_manager.go | 78 +++++++++- .../handler/task_manager_tool_cancel_test.go | 146 ++++++++++++++++++ 10 files changed, 551 insertions(+), 6 deletions(-) create mode 100644 internal/handler/config_eino_resilience_test.go diff --git a/internal/handler/agent.go b/internal/handler/agent.go index b1581d07..350aee0a 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -911,6 +911,32 @@ func (h *AgentHandler) publishProgressToTaskEventBus(conversationID, eventType, h.taskEventBus.Publish(conversationID, sseLine) } +func isInternalEinoDiagnosticProgress(eventType, message string, data interface{}) bool { + switch eventType { + case "model_output_rejected": + return true + case "progress": + msg := strings.TrimSpace(message) + if msg == "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。" || + msg == "Eino TurnLoop 已在安全点切换到用户补充后的下一轮。" || + msg == "已将用户补充推入 Eino TurnLoop，正在等待安全点切换…" { + return true + } + m, ok := data.(map[string]interface{}) + if !ok { + return false + } + switch strings.TrimSpace(fmt.Sprint(m["kind"])) { + case "turn_loop_takeover", "turn_loop_preempted": + return true + default: + return false + } + default: + return false + } +} + // enrichProgressEventData 为 SSE / taskEventBus 事件补齐 conversationId、messageId，便于前端懒加载过程详情。 func enrichProgressEventData(data interface{}, conversationID, assistantMessageID string) interface{} { if strings.TrimSpace(conversationID) == "" && strings.TrimSpace(assistantMessageID) == "" { @@ -1076,6 +1102,10 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun progressMu.Lock() defer progressMu.Unlock() + if isInternalEinoDiagnosticProgress(eventType, message, data) { + return + } + // 上游在重试/补偿时可能重复回调相同 tool_call/tool_result。 // 这里做幂等过滤，保证前端展示和 process_details 都以唯一事件为准。 if (eventType == "tool_call" || eventType == "tool_result") && data != nil { diff --git a/internal/handler/agent_progress_callback_test.go b/internal/handler/agent_progress_callback_test.go index 3cab9b58..8c2b0da5 100644 --- a/internal/handler/agent_progress_callback_test.go +++ b/internal/handler/agent_progress_callback_test.go @@ -80,6 +80,53 @@ func TestCreateProgressCallback_MirrorsWebStreamEvents(t *testing.T) { } } +func TestCreateProgressCallback_HidesInternalEinoDiagnostics(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + conv, err := db.CreateConversation("diag-hidden", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + bus := NewTaskEventBus() + h := &AgentHandler{logger: zap.NewNop(), db: db, taskEventBus: bus} + _, events := bus.Subscribe(conv.ID) + primaryCalls := 0 + cb := h.createProgressCallback( + context.Background(), nil, conv.ID, asst.ID, + func(string, string, interface{}) { primaryCalls++ }, + ) + + cb("model_output_rejected", "模型工具调用不完整或参数不安全，已阻止执行并要求重写。", map[string]interface{}{ + "reason": "invalid_tool_arguments_json", + }) + cb("progress", "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。", map[string]interface{}{ + "kind": "turn_loop_takeover", + }) + + if primaryCalls != 0 { + t.Fatalf("primary SSE calls = %d, want hidden diagnostics", primaryCalls) + } + select { + case payload := <-events: + t.Fatalf("unexpected mirrored diagnostic event: %s", string(payload)) + default: + } + details, err := db.GetProcessDetails(asst.ID) + if err != nil { + t.Fatalf("GetProcessDetails: %v", err) + } + if len(details) != 0 { + t.Fatalf("process details = %+v, want no diagnostics persisted", details) + } +} + func TestCreateProgressCallback_PersistsRunningResponseBeforeDone(t *testing.T) { tmp := t.TempDir() db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) diff --git a/internal/handler/config.go b/internal/handler/config.go index e3c26e5c..cde56cdd 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -356,6 +356,10 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) { LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(), LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(), LatestUserMessageTailRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective(), + ModelRetryMaxRetries: h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries, + ModelRetryMaxBackoffSec: h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec, + ModelFailoverChannels: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels...), + ModelFailoverMaxRetries: h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries, ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...), ToolSearchAlwaysVisibleEffectiveTools: mergeToolNameLists( h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools, @@ -1002,6 +1006,30 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) { } h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunes = v } + if req.MultiAgent.ModelRetryMaxRetries != nil { + v := *req.MultiAgent.ModelRetryMaxRetries + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries = v + } + if req.MultiAgent.ModelRetryMaxBackoffSec != nil { + v := *req.MultiAgent.ModelRetryMaxBackoffSec + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec = v + } + if req.MultiAgent.ModelFailoverChannels != nil { + h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels = dedupeTrimmedStringList(*req.MultiAgent.ModelFailoverChannels) + } + if req.MultiAgent.ModelFailoverMaxRetries != nil { + v := *req.MultiAgent.ModelFailoverMaxRetries + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries = v + } if req.MultiAgent.ToolSearchAlwaysVisibleTools != nil { h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools = dedupeToolNameList(*req.MultiAgent.ToolSearchAlwaysVisibleTools) } @@ -1015,6 +1043,10 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) { zap.Int("latest_user_message_max_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective()), zap.Int("latest_user_message_head_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective()), zap.Int("latest_user_message_tail_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective()), + zap.Int("model_retry_max_retries", h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries), + zap.Int("model_retry_max_backoff_sec", h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec), + zap.Int("model_failover_channels", len(h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels)), + zap.Int("model_failover_max_retries", h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries), zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)), ) } @@ -2191,10 +2223,18 @@ func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) { setIntInMap(mwNode, "latest_user_message_max_runes", cfg.EinoMiddleware.LatestUserMessageMaxRunesEffective()) setIntInMap(mwNode, "latest_user_message_head_runes", cfg.EinoMiddleware.LatestUserMessageHeadRunesEffective()) setIntInMap(mwNode, "latest_user_message_tail_runes", cfg.EinoMiddleware.LatestUserMessageTailRunesEffective()) + setIntInMap(mwNode, "model_retry_max_retries", cfg.EinoMiddleware.ModelRetryMaxRetries) + setIntInMap(mwNode, "model_retry_max_backoff_sec", cfg.EinoMiddleware.ModelRetryMaxBackoffSec) + setFlowStringSliceInMap(mwNode, "model_failover_channels", dedupeTrimmedStringList(cfg.EinoMiddleware.ModelFailoverChannels)) + setIntInMap(mwNode, "model_failover_max_retries", cfg.EinoMiddleware.ModelFailoverMaxRetries) setFlowStringSliceInMap(mwNode, "tool_search_always_visible_tools", dedupeToolNameList(cfg.EinoMiddleware.ToolSearchAlwaysVisibleTools)) } func dedupeToolNameList(in []string) []string { + return dedupeTrimmedStringList(in) +} + +func dedupeTrimmedStringList(in []string) []string { if len(in) == 0 { return []string{} } diff --git a/internal/handler/config_eino_resilience_test.go b/internal/handler/config_eino_resilience_test.go new file mode 100644 index 00000000..079d6acc --- /dev/null +++ b/internal/handler/config_eino_resilience_test.go @@ -0,0 +1,64 @@ +package handler + +import ( + "testing" + + "cyberstrike-ai/internal/config" + "gopkg.in/yaml.v3" +) + +func TestUpdateMultiAgentConfigWritesEinoModelResilience(t *testing.T) { + doc := &yaml.Node{ + Kind: yaml.DocumentNode, + Content: []*yaml.Node{{ + Kind: yaml.MappingNode, + Tag: "!!map", + }}, + } + + updateMultiAgentConfig(doc, config.MultiAgentConfig{ + Enabled: true, + RobotDefaultAgentMode: "deep", + PlanExecuteLoopMaxIterations: 3, + EinoMiddleware: config.MultiAgentEinoMiddlewareConfig{ + ModelRetryMaxRetries: 5, + ModelRetryMaxBackoffSec: 45, + ModelFailoverChannels: []string{"backup-openai", "backup-claude", "backup-openai"}, + ModelFailoverMaxRetries: 2, + }, + }) + + var got struct { + MultiAgent struct { + EinoMiddleware struct { + ModelRetryMaxRetries int `yaml:"model_retry_max_retries"` + ModelRetryMaxBackoffSec int `yaml:"model_retry_max_backoff_sec"` + ModelFailoverChannels []string `yaml:"model_failover_channels"` + ModelFailoverMaxRetries int `yaml:"model_failover_max_retries"` + } `yaml:"eino_middleware"` + } `yaml:"multi_agent"` + } + if err := doc.Decode(&got); err != nil { + t.Fatalf("decode config yaml: %v", err) + } + + mw := got.MultiAgent.EinoMiddleware + if mw.ModelRetryMaxRetries != 5 { + t.Fatalf("model_retry_max_retries = %d, want 5", mw.ModelRetryMaxRetries) + } + if mw.ModelRetryMaxBackoffSec != 45 { + t.Fatalf("model_retry_max_backoff_sec = %d, want 45", mw.ModelRetryMaxBackoffSec) + } + if mw.ModelFailoverMaxRetries != 2 { + t.Fatalf("model_failover_max_retries = %d, want 2", mw.ModelFailoverMaxRetries) + } + wantChannels := []string{"backup-openai", "backup-claude"} + if len(mw.ModelFailoverChannels) != len(wantChannels) { + t.Fatalf("model_failover_channels = %#v, want %#v", mw.ModelFailoverChannels, wantChannels) + } + for i, want := range wantChannels { + if mw.ModelFailoverChannels[i] != want { + t.Fatalf("model_failover_channels[%d] = %q, want %q", i, mw.ModelFailoverChannels[i], want) + } + } +} diff --git a/internal/handler/conversation.go b/internal/handler/conversation.go index 36c0515f..3cc5e4fb 100644 --- a/internal/handler/conversation.go +++ b/internal/handler/conversation.go @@ -2,6 +2,7 @@ package handler import ( "encoding/json" + "fmt" "net/http" "strconv" "strings" @@ -325,7 +326,7 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) { } details = database.DedupeConsecutiveProcessDetails(details) - out := processDetailsToJSON(h.logger, details, true) + out := processDetailsToJSON(h.logger, h.db, details, true) c.JSON(http.StatusOK, gin.H{ "processDetails": out, "total": len(out), @@ -374,7 +375,7 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) { return } details = database.DedupeConsecutiveProcessDetails(details) - out := processDetailsToJSON(h.logger, details, false) + out := processDetailsToJSON(h.logger, h.db, details, false) // A page may end between tool_call and tool_result. Return the full-history // execution summary so the UI can render terminal status without pretending // that an unloaded result is still running. @@ -409,7 +410,7 @@ func (h *ConversationHandler) GetProcessDetail(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"}) return } - out := processDetailsToJSON(h.logger, []database.ProcessDetail{*detail}, true) + out := processDetailsToJSON(h.logger, h.db, []database.ProcessDetail{*detail}, true) if len(out) == 0 { c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"}) return @@ -417,7 +418,7 @@ func (h *ConversationHandler) GetProcessDetail(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"processDetail": out[0]}) } -func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail, includeToolPayload bool) []map[string]interface{} { +func processDetailsToJSON(logger *zap.Logger, db *database.DB, details []database.ProcessDetail, includeToolPayload bool) []map[string]interface{} { out := make([]map[string]interface{}, 0, len(details)) for _, d := range details { var data interface{} @@ -426,6 +427,9 @@ func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail, logger.Warn("解析过程详情数据失败", zap.Error(err)) } } + if m, ok := data.(map[string]interface{}); ok { + enrichEmptyToolCallArgumentsFromExecution(logger, db, d, m) + } if !includeToolPayload { data = summarizeProcessDetailData(d.EventType, data) } @@ -442,6 +446,50 @@ func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail, return out } +func enrichEmptyToolCallArgumentsFromExecution(logger *zap.Logger, db *database.DB, detail database.ProcessDetail, data map[string]interface{}) { + if db == nil || detail.EventType != "tool_call" || !toolCallArgumentsEmpty(data) { + return + } + toolName := strings.TrimSpace(fmt.Sprint(data["toolName"])) + if toolName == "" || detail.ConversationID == "" || detail.CreatedAt.IsZero() { + return + } + execID, args, err := db.FindNearestToolExecutionArguments(detail.ConversationID, toolName, detail.CreatedAt, 5*time.Second) + if err != nil { + if logger != nil { + logger.Debug("未能从工具执行记录补全过程详情参数", + zap.Error(err), + zap.String("processDetailId", detail.ID), + zap.String("toolName", toolName)) + } + return + } + if len(args) == 0 { + return + } + data["argumentsObj"] = args + if b, err := json.Marshal(args); err == nil { + data["arguments"] = string(b) + } + if strings.TrimSpace(execID) != "" { + data["executionId"] = strings.TrimSpace(execID) + } +} + +func toolCallArgumentsEmpty(data map[string]interface{}) bool { + if data == nil { + return true + } + if args, ok := data["argumentsObj"].(map[string]interface{}); ok && len(args) > 0 { + return false + } + if raw, ok := data["arguments"]; ok { + s := strings.TrimSpace(fmt.Sprint(raw)) + return s == "" || s == "{}" || s == "null" + } + return true +} + func summarizeProcessDetailData(eventType string, data interface{}) interface{} { m, ok := data.(map[string]interface{}) if !ok || (eventType != "tool_call" && eventType != "tool_result") { @@ -452,7 +500,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{} "success": true, "isError": true, "executionId": true, "einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true, "agentFacing": true, - "status": true, "modelFacingIsError": true, "resultPreview": true, + "status": true, "modelFacingIsError": true, "resultPreview": true, } out := make(map[string]interface{}, len(allow)+1) for k, v := range m { diff --git a/internal/handler/conversation_process_details_test.go b/internal/handler/conversation_process_details_test.go index 8c3ba700..ff53898e 100644 --- a/internal/handler/conversation_process_details_test.go +++ b/internal/handler/conversation_process_details_test.go @@ -6,8 +6,10 @@ import ( "net/http/httptest" "path/filepath" "testing" + "time" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -73,3 +75,67 @@ func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testi } } } + +func TestProcessDetailsFullBackfillsEmptyToolCallArgumentsFromExecution(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-args.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("empty args", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "done", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "calling exec", map[string]interface{}{ + "toolName": "exec", "toolCallId": "call-empty", "arguments": "", "argumentsObj": nil, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: "exec-whoami", + ToolName: "exec", + Arguments: map[string]interface{}{"command": "whoami"}, + Status: "completed", + StartTime: time.Now(), + ConversationID: conversation.ID, + }); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?full=1", nil) + c.Params = gin.Params{{Key: "id", Value: message.ID}} + NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c) + if w.Code != 200 { + t.Fatalf("status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + ProcessDetails []map[string]interface{} `json:"processDetails"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(response.ProcessDetails) != 1 { + t.Fatalf("process details = %d, want 1", len(response.ProcessDetails)) + } + data, ok := response.ProcessDetails[0]["data"].(map[string]interface{}) + if !ok { + t.Fatalf("data = %#v", response.ProcessDetails[0]["data"]) + } + args, ok := data["argumentsObj"].(map[string]interface{}) + if !ok { + t.Fatalf("argumentsObj = %#v", data["argumentsObj"]) + } + if args["command"] != "whoami" { + t.Fatalf("command = %#v, want whoami", args["command"]) + } + if data["executionId"] != "exec-whoami" { + t.Fatalf("executionId = %#v, want exec-whoami", data["executionId"]) + } +} diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index 6d965813..8ec130f0 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -227,6 +227,12 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID) taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks) taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = multiagent.WithAgentRuntimeCancelRegistrar(taskCtxLoop, func(cancel func(error) bool) func() { + return h.tasks.BindAgentRuntimeCancel(conversationID, cancel) + }) + taskCtxLoop = multiagent.WithAgentTurnLoopInterruptRegistrar(taskCtxLoop, func(push func(string) bool) func() { + return h.tasks.BindAgentTurnLoopInterrupt(conversationID, push) + }) taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) { return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments) }) @@ -273,6 +279,14 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { } cause := context.Cause(baseCtx) + if cause == nil { + switch { + case errors.Is(runErr, multiagent.ErrInterruptContinue): + cause = multiagent.ErrInterruptContinue + case errors.Is(runErr, ErrTaskCancelled): + cause = ErrTaskCancelled + } + } if errors.Is(cause, multiagent.ErrInterruptContinue) { if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { h.persistEinoAgentTraceForResume(conversationID, result) diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index e1d2c7bb..d68f0c28 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -240,6 +240,12 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID) taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks) taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = multiagent.WithAgentRuntimeCancelRegistrar(taskCtxLoop, func(cancel func(error) bool) func() { + return h.tasks.BindAgentRuntimeCancel(conversationID, cancel) + }) + taskCtxLoop = multiagent.WithAgentTurnLoopInterruptRegistrar(taskCtxLoop, func(push func(string) bool) func() { + return h.tasks.BindAgentTurnLoopInterrupt(conversationID, push) + }) taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) { return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments) }) @@ -287,6 +293,14 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { } cause := context.Cause(baseCtx) + if cause == nil { + switch { + case errors.Is(runErr, multiagent.ErrInterruptContinue): + cause = multiagent.ErrInterruptContinue + case errors.Is(runErr, ErrTaskCancelled): + cause = ErrTaskCancelled + } + } if errors.Is(cause, multiagent.ErrInterruptContinue) { if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { h.persistEinoAgentTraceForResume(conversationID, result) diff --git a/internal/handler/task_manager.go b/internal/handler/task_manager.go index bdd045b5..1a55789e 100644 --- a/internal/handler/task_manager.go +++ b/internal/handler/task_manager.go @@ -46,6 +46,14 @@ type AgentTask struct { // hitlCognition 本轮运行中供 HITL/审计 Agent 读取的上下文（用户原话 + 思考，不含会话历史） hitlCognition *hitlCognitionState + // agentRuntimeCancel 当前 Eino ADK 原生 AgentCancelFunc 包装；取消任务时先触发它，再走 context 兜底。 + agentRuntimeCancel func(error) bool + agentRuntimeCancelVersion uint64 + + // agentTurnLoopInterrupt 当前 Eino TurnLoop 用户补充 push hook；中断并继续时优先将补充作为新 turn item 入队。 + agentTurnLoopInterrupt func(string) bool + agentTurnLoopInterruptVersion uint64 + cancel func(error) } @@ -220,6 +228,58 @@ func (m *AgentTaskManager) BindTaskCancel(conversationID string, cancel context. } } +// BindAgentRuntimeCancel 登记当前运行段的 Eino 原生 cancel hook。 +func (m *AgentTaskManager) BindAgentRuntimeCancel(conversationID string, cancel func(error) bool) func() { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || cancel == nil { + return func() {} + } + m.mu.Lock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + m.mu.Unlock() + return func() {} + } + t.agentRuntimeCancelVersion++ + version := t.agentRuntimeCancelVersion + t.agentRuntimeCancel = cancel + m.mu.Unlock() + + return func() { + m.mu.Lock() + defer m.mu.Unlock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil && cur.agentRuntimeCancelVersion == version { + cur.agentRuntimeCancel = nil + } + } +} + +// BindAgentTurnLoopInterrupt 登记当前运行任务的 Eino TurnLoop 用户补充入队 hook。 +func (m *AgentTaskManager) BindAgentTurnLoopInterrupt(conversationID string, push func(string) bool) func() { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || push == nil { + return func() {} + } + m.mu.Lock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + m.mu.Unlock() + return func() {} + } + t.agentTurnLoopInterruptVersion++ + version := t.agentTurnLoopInterruptVersion + t.agentTurnLoopInterrupt = push + m.mu.Unlock() + + return func() { + m.mu.Lock() + defer m.mu.Unlock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil && cur.agentTurnLoopInterruptVersion == version { + cur.agentTurnLoopInterrupt = nil + } + } +} + // ActiveMCPExecutionID 返回当前会话进行中的工具 executionId，无则空串。 func (m *AgentTaskManager) ActiveMCPExecutionID(conversationID string) string { conversationID = strings.TrimSpace(conversationID) @@ -402,13 +462,29 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, if cause == nil { cause = ErrTaskCancelled } + interruptPush := task.agentTurnLoopInterrupt + interruptNote := task.InterruptContinueNote + runtimeCancel := task.agentRuntimeCancel var toolCanceler func(string) if errors.Is(cause, ErrTaskCancelled) { toolCanceler = m.toolCanceler } m.mu.Unlock() - if cancel != nil { + if errors.Is(cause, multiagent.ErrInterruptContinue) && interruptPush != nil && interruptPush(interruptNote) { + m.mu.Lock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil { + cur.InterruptContinueNote = "" + } + m.mu.Unlock() + return true, nil + } + + runtimeHandled := false + if runtimeCancel != nil { + runtimeHandled = runtimeCancel(cause) + } + if cancel != nil && !runtimeHandled { cancel(cause) } if toolCanceler != nil { diff --git a/internal/handler/task_manager_tool_cancel_test.go b/internal/handler/task_manager_tool_cancel_test.go index 705aa068..51b411b0 100644 --- a/internal/handler/task_manager_tool_cancel_test.go +++ b/internal/handler/task_manager_tool_cancel_test.go @@ -32,6 +32,78 @@ func TestCancelTaskInvokesToolCancelerOnFullStop(t *testing.T) { } } +func TestCancelTaskUsesAgentRuntimeCancelAsPrimaryPath(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + tm.SetToolCanceler(func(conversationID string) { + if conversationID == "conv-native" { + order = append(order, "tool") + } + }) + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-native", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + unregister := tm.BindAgentRuntimeCancel("conv-native", func(err error) bool { + if !errors.Is(err, ErrTaskCancelled) { + t.Fatalf("runtime cancel got %v", err) + } + order = append(order, "runtime") + return true + }) + defer unregister() + + ok, err := tm.CancelTask("conv-native", ErrTaskCancelled) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"runtime", "tool"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } +} + +func TestCancelTaskFallsBackToContextWhenAgentRuntimeCancelMisses(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-fallback", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + unregister := tm.BindAgentRuntimeCancel("conv-fallback", func(err error) bool { + order = append(order, "runtime") + return false + }) + defer unregister() + + ok, err := tm.CancelTask("conv-fallback", ErrTaskCancelled) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"runtime", "context"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } +} + func TestCancelTaskSkipsToolCancelerOnInterruptContinue(t *testing.T) { tm := NewAgentTaskManager() called := false @@ -54,6 +126,80 @@ func TestCancelTaskSkipsToolCancelerOnInterruptContinue(t *testing.T) { } } +func TestCancelTaskPushesInterruptContinueToTurnLoopFirst(t *testing.T) { + tm := NewAgentTaskManager() + ctx, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-turn", "hello", cancel); err != nil { + t.Fatalf("StartTask: %v", err) + } + tm.SetInterruptContinueNote("conv-turn", "focus ssh") + + var gotNote string + unregister := tm.BindAgentTurnLoopInterrupt("conv-turn", func(note string) bool { + gotNote = note + return true + }) + defer unregister() + + ok, err := tm.CancelTask("conv-turn", multiagent.ErrInterruptContinue) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + if gotNote != "focus ssh" { + t.Fatalf("turn loop note = %q, want focus ssh", gotNote) + } + if cause := context.Cause(ctx); cause != nil { + t.Fatalf("context should not be cancelled when turn loop accepted interrupt, got %v", cause) + } + if note := tm.TakeInterruptContinueNote("conv-turn"); note != "" { + t.Fatalf("interrupt note should be consumed after turn loop push, got %q", note) + } +} + +func TestCancelTaskFallsBackWhenTurnLoopInterruptRejects(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-turn-fallback", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + tm.SetInterruptContinueNote("conv-turn-fallback", "fallback note") + unregisterTurn := tm.BindAgentTurnLoopInterrupt("conv-turn-fallback", func(note string) bool { + order = append(order, "turn") + if note != "fallback note" { + t.Fatalf("turn loop note = %q, want fallback note", note) + } + return false + }) + defer unregisterTurn() + unregisterRuntime := tm.BindAgentRuntimeCancel("conv-turn-fallback", func(err error) bool { + order = append(order, "runtime") + return false + }) + defer unregisterRuntime() + + ok, err := tm.CancelTask("conv-turn-fallback", multiagent.ErrInterruptContinue) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"turn", "runtime", "context"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } + if note := tm.TakeInterruptContinueNote("conv-turn-fallback"); note != "fallback note" { + t.Fatalf("interrupt note should remain for fallback rerun, got %q", note) + } +} + func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) { tm := NewAgentTaskManager() var gotCause error