From 1b0d5ea49eb72c7362a7d6d14a125d7332a975d5 Mon Sep 17 00:00:00 2001 From: Ed1s0nZ Date: Tue, 6 Oct 2026 03:24:14 +0800 Subject: [PATCH] fix: enforce configured workflow tool execution timeouts --- docs/bug-sheet-review.md | 4 ++ internal/mcp/execution_deadline.go | 14 ++++ internal/mcp/execution_deadline_test.go | 51 ++++++++++++++ internal/mcp/execution_service.go | 18 +++++ internal/workflow/agent_progress_test.go | 25 +++++++ internal/workflow/nodes.go | 28 +++++++- internal/workflow/tool_timeout_test.go | 90 ++++++++++++++++++++++++ 7 files changed, 229 insertions(+), 1 deletion(-) create mode 100644 internal/mcp/execution_deadline.go create mode 100644 internal/mcp/execution_deadline_test.go create mode 100644 internal/workflow/agent_progress_test.go create mode 100644 internal/workflow/tool_timeout_test.go diff --git a/docs/bug-sheet-review.md b/docs/bug-sheet-review.md index 969c2c90..e1e8e58c 100644 --- a/docs/bug-sheet-review.md +++ b/docs/bug-sheet-review.md @@ -136,3 +136,7 @@ ### 侧栏标题及仪表盘项目标签（表格第 2、18 行） 最近对话标题的 flex 子元素填满剩余空间，但原生按钮默认文本居中，导致标题偏向中间。增加明确左对齐。第 2 行截图未提供其他可识别重叠证据，未重新设计侧栏。仪表盘项目标签使用居中 inline-flex 时长文字会首尾裁切且不产生省略号，改为普通块级文本截断，完整名称保留既有 title。实际样式表与代表性 DOM 在 1280px 和 375px 视口下验证：标题左对齐，长项目名称尾部省略，相邻文本可见。此验证不覆盖完整应用交互。 + +### 工作流进度及已有工具超时配置核对（表格第 29 行，继续核对） + +当前 Agent 回调会转发 iteration、tool_call、tool_result 并附带工作流及节点标识，新增回归测试确认此路径；原截图没有运行日志，尚不能确认历史零进度原因。发现已有工具节点 timeout_seconds 虽校验却未执行，现将配置期限传给 MCP 后台执行服务的硬超时，兼容后台等待分离机制，不影响未设置期限的任务。实际慢 MCP 工具测试确认 1 秒期限触发失败；执行服务测试确认等待方取消不抹掉明确期限、不延长已有更短硬超时、期限后迟到成功保留 hard_timeout 及结束时间。工作流和 MCP 包测试通过。此修复不宣称解决 Agent 节点历史全程零进度，也未新增 Agent 超时配置。 diff --git a/internal/mcp/execution_deadline.go b/internal/mcp/execution_deadline.go new file mode 100644 index 00000000..f59d0172 --- /dev/null +++ b/internal/mcp/execution_deadline.go @@ -0,0 +1,14 @@ +package mcp + +import ( + "context" + "time" +) + +type toolExecutionDeadlineKey struct{} + +// WithToolExecutionDeadline sets an execution limit independently of a caller's +// bounded wait. Detached workers retain this explicit limit. +func WithToolExecutionDeadline(ctx context.Context, deadline time.Time) context.Context { + return context.WithValue(ctx, toolExecutionDeadlineKey{}, deadline) +} diff --git a/internal/mcp/execution_deadline_test.go b/internal/mcp/execution_deadline_test.go new file mode 100644 index 00000000..66ae10de --- /dev/null +++ b/internal/mcp/execution_deadline_test.go @@ -0,0 +1,51 @@ +package mcp + +import ( + "context" + "go.uber.org/zap" + "testing" + "time" +) + +func TestExplicitExecutionDeadlineSurvivesDetachedWait(t *testing.T) { + service := NewExecutionService(nil, zap.NewNop()) + waitCtx, cancelWait := context.WithCancel(context.Background()) + ctx := WithToolExecutionDeadline(waitCtx, time.Now().Add(50*time.Millisecond)) + started := make(chan struct{}) + handle, err := service.Submit(ctx, ExecutionRequest{ToolName: "test-only-deadline", Run: func(runCtx context.Context) (*ToolResult, error) { + if _, ok := runCtx.Deadline(); !ok { + t.Error("execution deadline lost") + } + close(started) + <-runCtx.Done() + return &ToolResult{}, nil + }}) + if err != nil { + t.Fatal(err) + } + <-started + cancelWait() + snapshot, err := service.Wait(context.Background(), handle.ID, time.Second) + if err != nil { + t.Fatal(err) + } + if snapshot.Execution.Status != ToolExecutionStatusHardTimeout || snapshot.Execution.EndTime == nil { + t.Fatalf("deadline did not close worker: %#v", snapshot.Execution) + } +} + +func TestExplicitExecutionDeadlineDoesNotExtendExistingHardTimeout(t *testing.T) { + service := NewExecutionService(nil, zap.NewNop()) + ctx := WithToolExecutionDeadline(context.Background(), time.Now().Add(time.Hour)) + handle, err := service.Submit(ctx, ExecutionRequest{ToolName: "test-only-short-limit", HardTimeout: 20 * time.Millisecond, Run: func(ctx context.Context) (*ToolResult, error) { <-ctx.Done(); return nil, ctx.Err() }}) + if err != nil { + t.Fatal(err) + } + snapshot, err := service.Wait(context.Background(), handle.ID, time.Second) + if err != nil { + t.Fatal(err) + } + if snapshot.Execution.Status != ToolExecutionStatusHardTimeout { + t.Fatalf("hard timeout extended: %#v", snapshot.Execution) + } +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go index 5e858603..36072dbf 100644 --- a/internal/mcp/execution_service.go +++ b/internal/mcp/execution_service.go @@ -153,6 +153,19 @@ func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*E } runCtx := detachedExecutionContext(ctx) + var executionDeadline time.Time + if ctx != nil { + executionDeadline, _ = ctx.Value(toolExecutionDeadlineKey{}).(time.Time) + } + if !executionDeadline.IsZero() { + remaining := time.Until(executionDeadline) + if remaining <= 0 { + remaining = time.Nanosecond + } + if req.HardTimeout <= 0 || remaining < req.HardTimeout { + req.HardTimeout = remaining + } + } var cancel context.CancelFunc if req.HardTimeout > 0 { runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout) @@ -252,6 +265,11 @@ func (s *ExecutionService) markEntryRunning(entry *executionEntry) { func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) { id := entry.exec.ID + if err == nil && ctx != nil && errors.Is(ctx.Err(), context.DeadlineExceeded) { + if deadline, ok := ctx.Value(toolExecutionDeadlineKey{}).(time.Time); ok && !deadline.IsZero() { + err = context.DeadlineExceeded + } + } var blockedErr *toolGuardBlockError if errors.As(err, &blockedErr) { result, err = blockedErr.result, nil diff --git a/internal/workflow/agent_progress_test.go b/internal/workflow/agent_progress_test.go new file mode 100644 index 00000000..24d61c51 --- /dev/null +++ b/internal/workflow/agent_progress_test.go @@ -0,0 +1,25 @@ +package workflow + +import "testing" + +func TestWorkflowAgentProgressForwardsToolActivityWithNodeIdentity(t *testing.T) { + state := newWorkflowLocalState(nil, "test-only-run") + node := graphNode{ID: "agent-1", Type: "agent", Label: "Agent"} + events := []string{} + callback := workflowAgentProgress(func(kind, message string, data interface{}) { + events = append(events, kind) + m, ok := data.(map[string]interface{}) + if !ok || m["workflowNodeId"] != node.ID || m["workflowRunId"] != state.WorkflowRunID { + t.Fatalf("missing workflow identity: %#v", data) + } + }, state, node) + for _, kind := range []string{"iteration", "tool_call", "tool_result"} { + callback(kind, "test-only activity", map[string]interface{}{"iteration": 1, "einoScope": "main"}) + } + for _, kind := range []string{"response_start", "response_delta", "response", "done"} { + callback(kind, "", nil) + } + if len(events) != 3 || events[0] != "iteration" || events[1] != "tool_call" || events[2] != "tool_result" { + t.Fatalf("activity not forwarded or nested terminal leaked: %v", events) + } +} diff --git a/internal/workflow/nodes.go b/internal/workflow/nodes.go index eb0342e2..8351f9cd 100644 --- a/internal/workflow/nodes.go +++ b/internal/workflow/nodes.go @@ -4,10 +4,12 @@ import ( "context" "fmt" "strings" + "time" "unicode/utf8" "cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/multiagent" ) @@ -74,7 +76,16 @@ func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *Workf "args": toolArgs, }) } - result, err := args.Agent.ExecuteMCPToolForConversation(ctx, args.ConversationID, toolName, toolArgs) + toolCtx, cancel, err := workflowToolContext(ctx, cfgString(node.Config, "timeout_seconds")) + if err != nil { + errText := err.Error() + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "error": errText}), false, "failed", errText + } + defer cancel() + result, err := args.Agent.ExecuteMCPToolForConversation(toolCtx, args.ConversationID, toolName, toolArgs) + if err == nil && toolCtx.Err() != nil { + err = toolCtx.Err() + } if err != nil { errText := err.Error() return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "arguments": toolArgs, "error": errText}), false, "failed", errText @@ -130,6 +141,21 @@ func truncateWorkflowToolOutput(output string, maxBytes int, executionID string) return output[:head] + marker + output[tailStart:] } +func workflowToolContext(parent context.Context, raw string) (context.Context, context.CancelFunc, error) { + if strings.TrimSpace(raw) == "" { + ctx, cancel := context.WithCancel(parent) + return ctx, cancel, nil + } + seconds, err := parsePositiveInt(raw) + if err != nil || int64(seconds) > int64((1<<63-1)/int64(time.Second)) { + return nil, nil, fmt.Errorf("工具节点超时时间必须是有效正整数秒") + } + ctx, cancel := context.WithTimeout(parent, time.Duration(seconds)*time.Second) + deadline, _ := ctx.Deadline() + ctx = mcp.WithToolExecutionDeadline(ctx, deadline) + return ctx, cancel, nil +} + func runAgentNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { if args.AppCfg == nil || args.Agent == nil { errText := "Agent 节点执行失败：应用配置或 Agent 为空" diff --git a/internal/workflow/tool_timeout_test.go b/internal/workflow/tool_timeout_test.go new file mode 100644 index 00000000..c98c3828 --- /dev/null +++ b/internal/workflow/tool_timeout_test.go @@ -0,0 +1,90 @@ +package workflow + +import ( + "context" + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "go.uber.org/zap" + "strings" + "testing" + "time" +) + +func TestWorkflowToolContextHonorsConfiguredAndParentDeadlines(t *testing.T) { + before := time.Now() + ctx, cancel, err := workflowToolContext(context.Background(), "2") + if err != nil { + t.Fatal(err) + } + defer cancel() + deadline, ok := ctx.Deadline() + if !ok || deadline.Before(before.Add(2*time.Second)) || deadline.After(time.Now().Add(2*time.Second)) { + t.Fatalf("unexpected deadline: %v", deadline) + } + parent, stop := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer stop() + child, cleanup, err := workflowToolContext(parent, "2") + if err != nil { + t.Fatal(err) + } + defer cleanup() + <-child.Done() + if child.Err() != context.DeadlineExceeded { + t.Fatalf("parent deadline lost: %v", child.Err()) + } +} + +func TestWorkflowToolContextRejectsInvalidAndPreservesUnset(t *testing.T) { + for _, raw := range []string{"0", "-1", "1.5", "9223372036854775807"} { + if _, _, err := workflowToolContext(context.Background(), raw); err == nil { + t.Fatalf("accepted %q", raw) + } + } + ctx, cancel, err := workflowToolContext(context.Background(), "") + if err != nil { + t.Fatal(err) + } + if _, ok := ctx.Deadline(); ok { + t.Fatal("invented deadline for unset configuration") + } + cancel() + if ctx.Err() != context.Canceled { + t.Fatal("context cancellation lost") + } +} + +func TestRunToolNodeAppliesTimeoutToActualMCPExecution(t *testing.T) { + logger := zap.NewNop() + server := mcp.NewServer(logger) + observedDeadline := make(chan time.Time, 1) + server.RegisterTool(mcp.Tool{Name: "test-only-slow-tool", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + deadline, ok := ctx.Deadline() + if !ok { + t.Error("tool did not receive configured deadline") + } + observedDeadline <- deadline + <-ctx.Done() + return nil, ctx.Err() + }) + ag := agent.NewAgent(&config.OpenAIConfig{APIKey: "test-only-key", Model: "test-only-model"}, &config.AgentConfig{}, server, nil, logger, 1) + node := graphNode{ID: "tool-1", Type: "tool", Config: map[string]interface{}{"tool_name": "test-only-slow-tool", "timeout_seconds": "1"}} + parent, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + start := time.Now() + output, proceed, status, errText := runToolNode(parent, RunArgs{Agent: ag, ConversationID: "test-only-conversation"}, node, newWorkflowLocalState(nil, "test-only-run")) + if proceed || status != "failed" || !strings.Contains(errText, "deadline exceeded") { + t.Fatalf("timeout became success: %v %s %s %#v", proceed, status, errText, output) + } + if time.Since(start) >= 3*time.Second { + t.Fatal("configured timeout not applied") + } + select { + case deadline := <-observedDeadline: + if deadline.After(start.Add(2 * time.Second)) { + t.Fatal("tool inherited only parent timeout") + } + default: + t.Fatal("actual MCP handler was not executed") + } +}