From 2ac176b96a3aba6ccc7e2dba991070b393e6bcf0 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Mon, 17 Aug 2026 09:01:19 -0400 Subject: [PATCH] =?UTF-8?q?feat:=20set=5Fgoals=20=E5=8D=87=E7=BA=A7?= =?UTF-8?q?=E4=B8=BA=E5=8F=97=E7=AE=A1=E9=A2=86=E5=9F=9F=E5=B7=A5=E5=85=B7?= =?UTF-8?q?(=E5=8F=AF=20web=20=E7=AE=A1=E7=90=86,=E9=BB=98=E8=AE=A4?= =?UTF-8?q?=E7=BB=91=20goals+mainagent)+=20=E4=B8=BB=20agent=20=E5=8A=A0?= =?UTF-8?q?=E7=9B=AE=E6=A0=87=E5=A4=8D=E6=B4=BB=E4=BB=BB=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - set_goals 从目标拆解器内联提交工具,改为入 tools 目录的领域工具:直接写 goal 节点、 连任务根(RelSpawns)、notify 唤醒 planner;支持单条/批量;origin 按来源区分 (拆解器=goals、主 agent 运行时=human)。新 key 走 SeedTool 首插入,旧库无需迁移。 - 默认绑定 goals + mainagent:目录 seed 新增 goals 项并入 order,MainAgentTools 纳入; DecomposeGoals 统一改用同一受管工具直接落库,createGoals 只保留无产出兜底。 - 主 agent 通过 set_goals 新增目标时,把终态(done/failed/timeout)/暂停的任务拉回 running 继续跑:新增 ToolSet.resumeTask 回调(仅 mainagent 接线),经 resumeTaskForGoal 置 running + engine.Run(幂等/重启后重起循环)+ 解除暂停。拆解器/worker 不触发。 - mainagent 提示词补充 set_goals 用法。 --- agent/goals.go | 55 +++++++++++----------- agent/mainagent.go | 6 ++- agent/toolcatalog.go | 5 +- agent/tools.go | 103 ++++++++++++++++++++++++++++++++++++++++++ agent/tools_insert.go | 2 + server/goals.go | 55 ++++++++++++++++------ server/server.go | 3 +- 7 files changed, 184 insertions(+), 45 deletions(-) diff --git a/agent/goals.go b/agent/goals.go index 3674215..1ce68fa 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -72,43 +72,30 @@ type GoalSpec struct { // // as + taskID, when non-nil/positive, wire the add_task_scope tool so the // decomposer can register the explicit asset scope it extracts from the goal. -func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc string, as *db.AssetStore, taskID int64, emit func(db.Activity)) []GoalSpec { +// +// ts is the task's exploration store: set_goals writes the decomposed goal nodes +// straight into it (the same managed tool the main agent uses to add goals at +// runtime). The returned specs are read back from the store so callers can emit +// per-goal activity and detect the "LLM produced nothing" case for their fallback. +func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { prov, err := c.NewProvider() if err != nil { return nil } - var out struct { - Goals []GoalSpec `json:"goals"` - } - captured := false - submitTool := actool.Build(actool.Spec{ - Name: "set_goals", - Description: "提交拆分后的目标列表。", - Schema: map[string]any{"type": "object", "properties": map[string]any{ - "goals": map[string]any{"type": "array", "items": map[string]any{ - "type": "object", "properties": map[string]any{ - "text": map[string]any{"type": "string", "description": "一个独立可验证的目标"}, - "vulnclass": map[string]any{"type": "string", "description": "对应漏洞类(若明确)，如 SQLi/IDOR；业务逻辑目标可留空"}, - }, "required": []string{"text"}, - }}, - }, "required": []string{"goals"}}, - Permissions: func(context.Context, json.RawMessage, acperm.Context) acperm.Decision { return acperm.Allowed() }, - Run: func(_ context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) { - _ = json.Unmarshal(in, &out) - captured = true - return actool.Text("ok"), nil - }, - }) + // worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link + // each goal under the task root. This is the catalog's real set_goals tool, so a + // web-edited description/schema on it applies here too. + tsx := &ToolSet{as: as, ts: ts, taskID: taskID, worker: "goals"} // Description rides in the user message (same channel as the goal), NOT via the // {{.EngagementDescription}} template var — else a prompt that references the var // would inject the description twice. System prompt stays pure static instructions. sys := renderSystem("goals", goalsDefaultTmpl, GoalsVars{DataDir: dataDir, Now: nowStr()}) - tools := []actool.CoreTool{submitTool} + tools := []actool.CoreTool{tsx.setGoals()} // Wire add_task_scope only when we have a real asset store + task to write to. // The scope-extraction tail is appended in lockstep so the prompt never asks for // a tool that isn't present. if as != nil && taskID > 0 { - tools = append(tools, (&ToolSet{as: as, taskID: taskID}).addTaskScope()) + tools = append(tools, tsx.addTaskScope()) sys += goalsScopeTail } userMsg := "任务目标：\n" + goalText @@ -131,8 +118,22 @@ func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc strin DisableBackgroundTasks: true, MaxTurns: 6, }, userMsg, captureEmit) - if !captured { + // set_goals persisted the goals directly; read them back so the caller sees what + // was written (empty slice ⇒ the LLM produced nothing ⇒ caller falls back). + if ts == nil { return nil } - return out.Goals + nodes, _ := ts.ListByKind(db.KindGoal, 10000) + var out []GoalSpec + for _, n := range nodes { + var p struct { + Text string `json:"text"` + VulnClass string `json:"vulnclass"` + } + _ = json.Unmarshal(n.Payload, &p) + if strings.TrimSpace(p.Text) != "" { + out = append(out, GoalSpec{Text: p.Text, VulnClass: p.VulnClass}) + } + } + return out } diff --git a/agent/mainagent.go b/agent/mainagent.go index 9269a25..0d440b7 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -48,6 +48,7 @@ const mainAgentDefaultTmpl = `你是一个授权渗透测试系统的"主 agent" 2. 操舵（把人的意图落到系统）： - 人想"改方向/强调某类漏洞/重点某区域" → 用 add_hint 写提示（规划者下次会读到）。 - 人想"立刻测某个具体目标" → 用 add_intent 直接注入一条高优先级意图（priority 8-10）。 + - 人想"新增一个要达成的最终目标" → 用 set_goals 增补目标。系统会把该目标写入任务图并**自动把已完成/暂停的任务拉回运行态继续跑**（规划者随后会据此重新判断是否达成），无需人工再点恢复。 3. 用人话简洁回复，说明你做了什么。 当前任务目标：{{.Goal}} @@ -63,13 +64,14 @@ func mainAgentSystem(goal, dataDir, workDir string) string { // non-nil, receives each execution step (thinking / tool_use / tool_result / // text / result) so the main-agent session shows its work — exactly like the // worker/planner sessions — not just the final answer. -func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify func()) (string, error) { +func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify, resume func()) (string, error) { tsx := NewToolSet(ts, "human") if as != nil { tsx.SetAssetStore(as, as.Companies()) } tsx.SetTaskID(taskID) - tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) + tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) + tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running // 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash） base := append(tsx.MainAgentTools(), actool.DefaultTools()...) tools, def, cleanup := AugmentTools(ctx, "mainagent", base) diff --git a/agent/toolcatalog.go b/agent/toolcatalog.go index c539b21..92e5c46 100644 --- a/agent/toolcatalog.go +++ b/agent/toolcatalog.go @@ -37,6 +37,9 @@ func builtinToolsByAgent() map[string][]actool.CoreTool { "mainagent": ts.MainAgentTools(), "planner": ts.PlannerTools(), "worker": ts.WorkerTools(), + // goals（目标拆解器）默认绑 set_goals：它靠这个工具把拆出的目标写进图。 + // 与 mainagent 共用同一受管工具，web 端可改描述/schema、按 agent 勾选。 + "goals": {ts.setGoals()}, // auto 默认绑漏洞上报 + 资产管理工具，其他域工具可在 UI 按需勾选。 // 新库由此 seed 写入；老库由 seedAutoDefaultBindings 迁移。 "auto": {ts.addFinding(), ts.insertAssets(), ts.addCompanyScope(), ts.listAssets(), ts.listCompanies()}, @@ -50,7 +53,7 @@ func builtinToolsByAgent() map[string][]actool.CoreTool { // list_assets 三者都有）合成一条，Agents 取并集。顺序稳定（mainagent→planner→worker）。 func BuiltinToolSeeds() []ToolSeed { byAgent := builtinToolsByAgent() - order := []string{"mainagent", "planner", "worker", "auto", "pentest"} + order := []string{"mainagent", "goals", "planner", "worker", "auto", "pentest"} type acc struct { tool actool.CoreTool diff --git a/agent/tools.go b/agent/tools.go index 6a9d10c..3430ef4 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -77,8 +77,18 @@ type ToolSet struct { // carrying (intentID, summary) so the round can spell out which intent found what. // Wired for workers; nil elsewhere → falls back to notify (bare wake). notifyFinding func(intentID int64, summary string) + // resumeTask, if set, revives the task after a graph change that should make a + // stopped task run again (currently: set_goals adds a goal). It flips a terminal/ + // paused task back to running and (re)starts the engine loops — a plain notify() + // can't, because the planner's terminal gate swallows wakes. Wired ONLY for the + // main agent (human steering); nil for the goals decomposer and workers. + resumeTask func() } +// SetResumeTask wires the task-revive callback (see ToolSet.resumeTask). Set only by +// the main-agent chat, so runtime-added goals can pull a finished task back to running. +func (t *ToolSet) SetResumeTask(fn func()) { t.resumeTask = fn } + // SetNotify wires the planner-wake callback (see ToolSet.notify). Set by callers // that hold the task handle (main-agent chat, cross-task orchestration). func (t *ToolSet) SetNotify(fn func()) { t.notify = fn } @@ -799,6 +809,99 @@ func (t *ToolSet) addOneHint(it hintItem) (int64, error) { return id, err } +type goalItem struct { + Text string `json:"text"` + VulnClass string `json:"vulnclass"` +} + +// addOneGoal 挂一条 goal 节点(open)到探索图:连到任务根(origin fact,rel spawns), +// 并唤醒 planner(运行时新增的目标要尽快被判定)。origin 取 t.worker(缺省 system): +// goals 拆解器写入的记 "goals"、主 agent 运行时新增的记 "human"。 +func (t *ToolSet) addOneGoal(it goalItem) (int64, error) { + text := strings.TrimSpace(it.Text) + if text == "" { + return 0, fmt.Errorf("text 不能为空") + } + payload := map[string]any{"text": text} + if vc := strings.TrimSpace(it.VulnClass); vc != "" { + payload["vulnclass"] = vc + } + origin := t.worker + if origin == "" { + origin = "system" + } + id, err := t.ts.AddNode(db.KindGoal, payload, 0, "open", origin, nil) + if err != nil { + return 0, err + } + if of, _ := t.ts.OriginFactID(); of > 0 && id > 0 { + _ = t.ts.Link(of, db.RelSpawns, id) // goals descend from the task root (origin fact) + } + if t.notify != nil { + t.notify() // wake the planner so a runtime-added goal is judged promptly (debounced) + } + return id, nil +} + +// setGoals 给【本任务】新增探索目标(goal 节点)。既是目标拆解器的提交工具,也是主 +// agent 运行时补目标的工具——同一个受管工具,可在 web 端改描述/schema、按 agent 绑定。 +func (t *ToolSet) setGoals() actool.CoreTool { + return writeTool("set_goals", + "给【本任务】新增探索目标(goal)。目标=最终可交付/可核验的结果,不是攻击步骤或侦察动作。\n"+ + "★优先批量:多个目标放进 goals 数组一次提交,返回 ids 与之等长同序(失败项 id=0,详情见 errors)。单条则省略 goals 直接给顶层 text。\n"+ + "vulnclass 可选:对应漏洞类(如 SQLi/IDOR),业务逻辑类目标留空。目标是否达成由系统判定标记 met,本工具只负责新增。", + obj(map[string]any{ + "goals": map[string]any{"type": "array", "description": "【优先用这个】要新增的目标数组,按顺序处理。每个元素:text(必填,一个独立可验证的最终目标)+ vulnclass(可选)。返回 ids 与本数组等长、同序。", "items": map[string]any{"type": "object"}}, + "text": str("[单条] 一个独立可验证的最终目标"), + "vulnclass": str("[单条] 对应漏洞类(若明确),如 SQLi/IDOR;业务逻辑目标可留空"), + }), + func(_ context.Context, in json.RawMessage) (actool.Result, error) { + if t.ts == nil { + return actool.Errorf("set_goals 未启用: ExplorationStore 未初始化"), nil + } + var a struct { + Goals []goalItem `json:"goals"` + goalItem // 单条模式:顶层 text/vulnclass + } + _ = json.Unmarshal(in, &a) + batch := len(a.Goals) > 0 + items := a.Goals + if !batch { + items = []goalItem{a.goalItem} + } + + ids := make([]int64, len(items)) + errs := map[string]string{} + added := 0 + for i, it := range items { + id, err := t.addOneGoal(it) + if err != nil { + errs[strconv.Itoa(i)] = err.Error() + continue + } + ids[i] = id + added++ + } + // 主 agent 运行时新增目标 → 把已完成/暂停的任务拉回 running 继续跑(终态门会吞 + // 掉普通 notify,必须显式复活)。仅 mainagent 接了此回调;拆解器/worker 为 nil。 + if added > 0 && t.resumeTask != nil { + t.resumeTask() + } + + if !batch { // 单条:保持原返回 + if e, bad := errs["0"]; bad { + return actool.Errorf(e), nil + } + return actool.Text(fmt.Sprintf("goal added: %d", ids[0])), nil + } + out := map[string]any{"ids": ids} + if len(errs) > 0 { + out["errors"] = errs + } + return jsonResult(out) + }) +} + func (t *ToolSet) addHint() actool.CoreTool { return writeTool("add_hint", "把人类/主 agent 的战略提示挂到探索图，规划者下次生成意图时会读到它。\n"+ "★优先批量：多条提示放进 hints 数组一次提交（比逐条调用省往返）。返回 ids 数组，与 hints 等长同序（失败项 id=0，详情见 errors）。单条则省略 hints 直接给顶层 text。", diff --git a/agent/tools_insert.go b/agent/tools_insert.go index d442bf2..c808988 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -595,6 +595,8 @@ func (t *ToolSet) MainAgentTools() []actool.CoreTool { return []actool.CoreTool{ t.graphOverview(), t.listFindings(), t.listFacts(), t.nodeDetail(), t.getWorkerOutput(), t.getWorkerTrace(), t.searchAllWorkerTraces(), t.addHint(), t.addIntent(), + // set_goals：人可在运行时给本任务补一个新的最终目标（规划者据此重判是否达成）。 + t.setGoals(), // asset management (handlers guard nil store internally) t.insertAssets(), t.addCompanyScope(), t.listAssets(), t.addFinding(), t.recordFact(), diff --git a/server/goals.go b/server/goals.go index aefb30e..279a2e9 100644 --- a/server/goals.go +++ b/server/goals.go @@ -46,6 +46,35 @@ func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) { }() } +// resumeTaskForGoal 让「主 agent 通过 set_goals 新增目标」的任务能继续跑:把终态 +// (done/failed/timeout)拉回 running、解除暂停,并(重)启动引擎循环 + 唤醒 planner。 +// 已在 running 且未暂停的任务:只剩 Run 里的一次 Notify,近乎无副作用。 +// +// 为什么必须显式复活:planner 循环的终态门(engine.go)会吞掉普通 notify——光新增 goal +// 节点 + Notify 唤不醒已判完成的任务;重启后终态任务的 goroutine 也可能已不在,故还要 Run。 +func (s *Server) resumeTaskForGoal(t *Task) { + if t == nil { + return + } + // 终态 → 拉回 running(SetTaskStatus 会清 completed_at 并同步内存 handle 的 Status, + // 使 planner 循环的终态门放行)。 + if isTerminalStatus(t.Status) { + if err := s.m.SetTaskStatus(t.ID, "running"); err != nil { + log.Printf("[chat] task %s set_goals 复活置 running 失败: %v", t.ID, err) + } + } + // 确保引擎循环存活:已 started 只会 Notify;未 started(重启后未恢复的终态任务)则重起循环。 + s.engine.Run(s.ctx, t) + // 暂停中 → 解除暂停(清引擎内存态 + Notify),并持久化 paused=false(存活过重启)。 + if t.Paused || s.engine.IsPaused(t.ID) { + s.engine.Resume(t) + t.Paused = false + if err := s.m.SetTaskPaused(t.ID, false); err != nil { + log.Printf("[chat] task %s set_goals 复活解除暂停失败: %v", t.ID, err) + } + } +} + // createGoals materializes the goal node(s) under the task root (rel objective). // Decomposition is done ENTIRELY by the LLM (the project requires an LLM). There is // no rule-based fallback splitter — it only ever produced garbage (shredded URLs, @@ -65,31 +94,29 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity as = s.m.Assets() } taskID, _ := strconv.ParseInt(t.ID, 10, 64) - for _, g := range agent.DecomposeGoals(ctx, cfg, s.m.dir, t.Goal, t.Description, as, taskID, emit) { + // DecomposeGoals persists each decomposed goal via the set_goals tool straight + // into t.Store; it returns the specs it wrote so we can emit activity + spot the + // empty (LLM error / no provider) case below. No write-back needed here anymore. + for _, g := range agent.DecomposeGoals(ctx, cfg, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) { if strings.TrimSpace(g.Text) != "" { specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) } } } if len(specs) == 0 { - // No rule-based splitting: use the raw task goal verbatim as a single goal - // (only reached on an LLM error — normal runs get real decomposed goals). + // No decomposed goals (LLM error / no provider): use the raw task goal verbatim + // as a single goal so the task still has something to judge against. This is the + // only path that writes here — decomposed goals are already persisted by the tool. if g := strings.TrimSpace(t.Goal); g != "" { log.Printf("[goals] task %s: LLM 目标拆解无产出，回退为「原始目标作为单目标」", t.ID) + origin, _ := t.Store.OriginFactID() + id, _ := t.Store.AddNode(db.KindGoal, map[string]any{"text": g}, 0, "open", "system", nil) + if origin > 0 && id > 0 { + _ = t.Store.Link(origin, db.RelSpawns, id) // goal descends from the task root (origin fact) + } specs = []goalSpec{{Text: g}} } } - origin, _ := t.Store.OriginFactID() - for _, g := range specs { - payload := map[string]any{"text": g.Text} - if g.VulnClass != "" { - payload["vulnclass"] = g.VulnClass - } - id, _ := t.Store.AddNode(db.KindGoal, payload, 0, "open", "system", nil) - if origin > 0 && id > 0 { - _ = t.Store.Link(origin, db.RelSpawns, id) // goals descend from the task root (origin fact) - } - } return specs } diff --git a/server/server.go b/server/server.go index 439d249..0a13c12 100644 --- a/server/server.go +++ b/server/server.go @@ -2052,7 +2052,8 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // is the captured "result" step — no separate reply emit (would duplicate). emit := func(rec db.Activity) { s.engine.emitActivity(t, rec) } maTaskID, _ := strconv.ParseInt(t.ID, 10, 64) - if _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, req.Message, emit, t.Notify); err != nil && ctx.Err() == nil { + resume := func() { s.resumeTaskForGoal(t) } // set_goals 新增目标 → 把任务拉回 running + if _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, req.Message, emit, t.Notify, resume); err != nil && ctx.Err() == nil { s.engine.emitActivity(t, db.Activity{Worker: "mainagent", Kind: "text", IsError: true, Summary: "（主 Agent 出错：" + err.Error() + "）"}) } }()