diff --git a/CHANGELOG.md b/CHANGELOG.md index 28e92e04..50740449 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### Added +- 支持通过 `agent.task_timeout_minutes` 配置单代理、多代理、工作流及异步任务执行超时：未配置默认 600 分钟，`0` 表示不限时，负数拒绝加载；续跑保留重新计时行为，单次工具超时仍独立生效。 + - 新增单代理与多代理异步对话 API：接收后返回任务回执，客户端断线不停止执行；支持按用户隔离的幂等提交、持久状态与结果查询、显式取消以及服务重启后的中断状态。同步接口保持兼容。 - 启动时以统一最新 SQLite 定义自动检查全部项目管理表，缺表建表、缺字段补字段，无需维护逐字段迁移列表；保留已有数据和额外表/字段，结构补齐失败回滚并阻止启动。HITL 与通知表纳入统一初始化，历史漏洞外键迁移保留额外列、索引、触发器及依赖记录。 diff --git a/config.example.yaml b/config.example.yaml index 9e53e2e0..92db51b2 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -157,6 +157,8 @@ shodan: # Agent 配置 # 达到最大迭代次数时，AI 会自动总结测试结果 agent: + # 修改后重启服务，对新任务生效；例如 1440 = 24 小时，单次工具超时独立生效。 + task_timeout_minutes: 600 # 任务执行超时（分钟）；未配置=600，0=不限时，负数无效；续跑重新计时 max_iterations: 12000 # 全局最大迭代次数（单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用；agents/*.md 中 max_iterations>0 可单独覆盖） tool_timeout_minutes: 60 # 单次工具执行最大时长（分钟），超时自动终止；0 表示不限制（不推荐，易出现长时间挂起） tool_wait_timeout_seconds: 300 # 工具本轮最多等待（秒）；到时返回 execution_id，worker 继续后台执行，可用 wait_tool_execution 继续等待；0=等到完成 diff --git a/internal/config/agent_timeout.go b/internal/config/agent_timeout.go new file mode 100644 index 00000000..5e3efe85 --- /dev/null +++ b/internal/config/agent_timeout.go @@ -0,0 +1,27 @@ +package config + +import ( + "fmt" + "time" +) + +// EffectiveTaskTimeout returns the execution budget for each task segment. +// An omitted value retains the historical ten-hour limit; zero disables it. +func (a AgentConfig) EffectiveTaskTimeout() time.Duration { + if a.TaskTimeoutMinutes == nil { + return 600 * time.Minute + } + return time.Duration(*a.TaskTimeoutMinutes) * time.Minute +} + +func (a AgentConfig) validateTaskTimeout() error { + if a.TaskTimeoutMinutes == nil { + return nil + } + const maxMinutes = int64((1<<63 - 1) / time.Minute) + minutes := int64(*a.TaskTimeoutMinutes) + if minutes < 0 || minutes > maxMinutes { + return fmt.Errorf("agent.task_timeout_minutes must be between 0 and %d minutes", maxMinutes) + } + return nil +} diff --git a/internal/config/agent_timeout_test.go b/internal/config/agent_timeout_test.go new file mode 100644 index 00000000..4a1e7412 --- /dev/null +++ b/internal/config/agent_timeout_test.go @@ -0,0 +1,61 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" + + "gopkg.in/yaml.v3" +) + +func TestLoadTaskTimeout(t *testing.T) { + for _, tc := range []struct { + name, yaml string + want time.Duration + invalid bool + }{ + {"absent agent", "{}", 600 * time.Minute, false}, + {"absent field", "agent: {}", 600 * time.Minute, false}, + {"custom", "agent:\n task_timeout_minutes: 1440", 24 * time.Hour, false}, + {"unlimited", "agent:\n task_timeout_minutes: 0", 0, false}, + {"negative", "agent:\n task_timeout_minutes: -1", 0, true}, + {"overflow", "agent:\n task_timeout_minutes: 153722868", 0, true}, + } { + t.Run(tc.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.yaml") + if err := os.WriteFile(path, []byte(tc.yaml), 0600); err != nil { + t.Fatal(err) + } + cfg, err := Load(path) + if tc.invalid { + if err == nil || !strings.Contains(err.Error(), "agent.task_timeout_minutes") { + t.Fatalf("expected field validation error, got %v", err) + } + return + } + if err != nil { + t.Fatal(err) + } + if got := cfg.Agent.EffectiveTaskTimeout(); got != tc.want { + t.Fatalf("timeout = %v, want %v", got, tc.want) + } + // YAML saves must preserve explicit zero as distinct from omission. + data, err := yaml.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, data, 0600); err != nil { + t.Fatal(err) + } + again, err := Load(path) + if err != nil { + t.Fatal(err) + } + if again.Agent.EffectiveTaskTimeout() != tc.want { + t.Fatal("timeout changed after save/load") + } + }) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 776bd4e4..f4eaa2a9 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -1112,13 +1112,14 @@ type DatabaseConfig struct { } type AgentConfig struct { - MaxIterations int `yaml:"max_iterations" json:"max_iterations"` - ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长（分钟），超时自动终止，防止长时间挂起；0 表示不限制（不推荐） - ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数；到时返回 execution_id，worker 继续后台执行；0 表示等到完成 - ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数；0 表示默认 2 - ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发；0 表示默认 16 - ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断；0 表示默认 3；负数关闭 - ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数；0 表示默认 60 + TaskTimeoutMinutes *int `yaml:"task_timeout_minutes,omitempty" json:"task_timeout_minutes,omitempty"` // nil=600 minutes; 0=unlimited + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长（分钟），超时自动终止，防止长时间挂起；0 表示不限制（不推荐） + ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数；到时返回 execution_id，worker 继续后台执行；0 表示等到完成 + ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数；0 表示默认 2 + ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发；0 表示默认 16 + ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断；0 表示默认 3；负数关闭 + ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数；0 表示默认 60 // ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数（通用防挂死，不维护命令黑名单）；0=默认 300（5 分钟）；-1=关闭。 ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"` // WorkspaceRootDir 会话工作目录根路径（curl/wget 下载、read_file/glob/grep 本地分析）；空=tmp/workspace，其下按 projects/{id} 或 conversations/{id} 隔离。 @@ -1652,6 +1653,9 @@ func Load(path string) (*Config, error) { if err := yaml.Unmarshal(data, &cfg); err != nil { return nil, fmt.Errorf("解析配置文件失败: %w", err) } + if err := cfg.Agent.validateTaskTimeout(); err != nil { + return nil, err + } if cfg.ToolGuard != nil { if err := validateToolGuardYAML(data); err != nil { return nil, fmt.Errorf("调用拦截配置无效: %w", err) diff --git a/internal/handler/async_agent_runs.go b/internal/handler/async_agent_runs.go index e2a12ec2..ce25df11 100644 --- a/internal/handler/async_agent_runs.go +++ b/internal/handler/async_agent_runs.go @@ -10,7 +10,6 @@ import ( "net/http" "strings" "sync" - "time" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp" @@ -98,7 +97,7 @@ func (h *AgentHandler) submitAgentRunWithExecutor(c *gin.Context, multi bool, ex return } base, cancel := context.WithCancelCause(detachedAgentContext(c.Request.Context())) - ctx, timeoutCancel := context.WithTimeout(base, 600*time.Minute) + ctx, timeoutCancel := h.newTaskContext(base) task, err := h.tasks.StartTask(prep.ConversationID, req.Message, cancel) if err != nil { timeoutCancel() diff --git a/internal/handler/eino_empty_response_continue.go b/internal/handler/eino_empty_response_continue.go index 2c57df86..a946feaa 100644 --- a/internal/handler/eino_empty_response_continue.go +++ b/internal/handler/eino_empty_response_continue.go @@ -19,7 +19,7 @@ func (h *AgentHandler) rebindEinoRunningTask(parent context.Context, conversatio } baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(parent)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) - taskCtx, newTimeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, newTimeoutCancel := h.newTaskContext(baseCtx) h.tasks.UpdateTaskStatus(conversationID, "running") return baseCtx, cancelWithCause, taskCtx, newTimeoutCancel } diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index b3ed6207..3701190b 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -166,7 +166,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { var runErr error baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) if startedTask, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { var errorMsg string @@ -327,7 +327,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { timeoutCancel() baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) - taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel = h.newTaskContext(baseCtx) h.tasks.UpdateTaskStatus(conversationID, "running") continue } @@ -458,7 +458,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) defer cancelWithCause(nil) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) defer timeoutCancel() jsonTask, startErr := h.tasks.StartTask(prep.ConversationID, req.Message, cancelWithCause) if startErr != nil { diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index 6e9ecfa3..912de564 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -173,7 +173,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { var runErr error baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) if startedTask, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { var errorMsg string @@ -341,7 +341,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { timeoutCancel() baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) - taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel = h.newTaskContext(baseCtx) h.tasks.UpdateTaskStatus(conversationID, "running") continue } @@ -477,7 +477,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) defer cancelWithCause(nil) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) defer timeoutCancel() jsonTask, startErr := h.tasks.StartTask(prep.ConversationID, req.Message, cancelWithCause) if startErr != nil { diff --git a/internal/handler/task_deadline.go b/internal/handler/task_deadline.go index 37d7c8be..0448b261 100644 --- a/internal/handler/task_deadline.go +++ b/internal/handler/task_deadline.go @@ -3,6 +3,8 @@ package handler import ( "context" "errors" + + "cyberstrike-ai/internal/config" ) // taskDeadlineExceeded distinguishes the task's deadline from a child operation's @@ -11,3 +13,17 @@ import ( func taskDeadlineExceeded(taskCtx context.Context) bool { return taskCtx != nil && errors.Is(taskCtx.Err(), context.DeadlineExceeded) } + +// newTaskContext starts a fresh execution budget. Unlimited tasks remain cancellable. +// A parent's cancellation or earlier deadline is always respected. +func (h *AgentHandler) newTaskContext(parent context.Context) (context.Context, context.CancelFunc) { + var cfg config.AgentConfig + if h.config != nil { + cfg = h.config.Agent + } + timeout := cfg.EffectiveTaskTimeout() + if timeout == 0 { + return context.WithCancel(parent) + } + return context.WithTimeout(parent, timeout) +} diff --git a/internal/handler/task_timeout_test.go b/internal/handler/task_timeout_test.go new file mode 100644 index 00000000..3288500d --- /dev/null +++ b/internal/handler/task_timeout_test.go @@ -0,0 +1,122 @@ +package handler + +import ( + "context" + "errors" + "gopkg.in/yaml.v3" + "strconv" + "testing" + "time" + + "cyberstrike-ai/internal/config" +) + +func TestConfiguredTaskContext(t *testing.T) { + for _, minutes := range []int{-1, 0, 1440} { + t.Run(taskTimeoutCaseName(minutes), func(t *testing.T) { + h := &AgentHandler{} + want := 600 * time.Minute + if minutes >= 0 { + h.config = &config.Config{Agent: config.AgentConfig{TaskTimeoutMinutes: &minutes}} + want = time.Duration(minutes) * time.Minute + } + parent, cancelParent := context.WithCancelCause(context.Background()) + defer cancelParent(nil) + before := time.Now() + ctx, cancel := h.newTaskContext(parent) + defer cancel() + deadline, ok := ctx.Deadline() + if want == 0 { + if ok { + t.Fatal("unlimited task has a deadline") + } + } else if !ok || deadline.Before(before.Add(want)) || deadline.After(time.Now().Add(want)) { + t.Fatalf("unexpected deadline %v", deadline) + } + cancelParent(ErrTaskCancelled) + if !errors.Is(context.Cause(ctx), ErrTaskCancelled) { + t.Fatal("parent cancellation lost") + } + + ctx, cancel = h.newTaskContext(context.Background()) + cancel() + if !errors.Is(ctx.Err(), context.Canceled) { + t.Fatal("explicit cancellation lost") + } + + early := time.Now().Add(time.Minute) + parentDeadline, stop := context.WithDeadline(context.Background(), early) + defer stop() + ctx, cancel = h.newTaskContext(parentDeadline) + defer cancel() + if got, ok := ctx.Deadline(); !ok || !got.Equal(early) { + t.Fatal("earlier parent deadline lost") + } + }) + } +} + +func taskTimeoutCaseName(minutes int) string { + if minutes < 0 { + return "default" + } + if minutes == 0 { + return "unlimited" + } + return "custom" +} + +func TestRebindTaskRestartsConfiguredBudget(t *testing.T) { + for _, minutes := range []int{0, 1440} { + t.Run(taskTimeoutCaseName(minutes), func(t *testing.T) { + h := &AgentHandler{config: &config.Config{Agent: config.AgentConfig{TaskTimeoutMinutes: &minutes}}, tasks: NewAgentTaskManager()} + base, stop := context.WithCancelCause(context.Background()) + defer stop(nil) + if _, err := h.tasks.StartTask("resume", "test", stop); err != nil { + t.Fatal(err) + } + expired, cancelExpired := context.WithDeadline(base, time.Now().Add(-time.Second)) + before := time.Now() + _, stopNew, resumed, cancelNew := h.rebindEinoRunningTask(expired, "resume", cancelExpired) + defer stopNew(nil) + defer cancelNew() + if resumed.Err() != nil { + t.Fatalf("resumed task inherited expired context: %v", resumed.Err()) + } + deadline, ok := resumed.Deadline() + if minutes == 0 { + if ok { + t.Fatal("unlimited resume has deadline") + } + } else if !ok || deadline.Before(before.Add(24*time.Hour)) || deadline.After(time.Now().Add(24*time.Hour)) { + t.Fatalf("resume budget was not reset: %v", deadline) + } + stopNew(ErrTaskCancelled) + if !errors.Is(context.Cause(resumed), ErrTaskCancelled) { + t.Fatal("resumed task cannot be cancelled") + } + }) + } +} + +// Saving unrelated settings must not erase a manually configured task budget. +func TestSettingsSavePreservesTaskTimeout(t *testing.T) { + for _, value := range []string{"0", "1440"} { + var doc yaml.Node + if err := yaml.Unmarshal([]byte("agent:\n task_timeout_minutes: "+value+"\n"), &doc); err != nil { + t.Fatal(err) + } + updateAgentConfig(&doc, config.AgentConfig{MaxIterations: 100}) + data, err := yaml.Marshal(&doc) + if err != nil { + t.Fatal(err) + } + var cfg config.Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + t.Fatal(err) + } + if cfg.Agent.TaskTimeoutMinutes == nil || strconv.Itoa(*cfg.Agent.TaskTimeoutMinutes) != value { + t.Fatalf("settings save changed timeout: %s", data) + } + } +} diff --git a/internal/handler/workflow_integration.go b/internal/handler/workflow_integration.go index 3c40ea8c..042a3e8b 100644 --- a/internal/handler/workflow_integration.go +++ b/internal/handler/workflow_integration.go @@ -88,7 +88,7 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound( } baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(c.Request.Context())) defer cancelWithCause(nil) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) defer timeoutCancel() if startedTask, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { @@ -231,7 +231,7 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) defer cancelWithCause(nil) - taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + taskCtx, timeoutCancel := h.newTaskContext(baseCtx) defer timeoutCancel() if startedTask, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil {