pphua
9793b522d9
Add read-only task context inheritance, ordered task-level LLM failover, live Worker controls, cascading task cleanup, FIFO task concurrency, cancellation diagnostics, UI updates, tests, build script, and changelog.
610 lines
27 KiB
Go
610 lines
27 KiB
Go
package server
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Autumn-27/artex/agent"
|
||
"github.com/Autumn-27/artex/db"
|
||
"github.com/Autumn-27/norma/permission"
|
||
actool "github.com/Autumn-27/norma/tool"
|
||
)
|
||
|
||
// jsonResult marshals v to a JSON tool result.
|
||
func jsonResult(v any) (actool.Result, error) {
|
||
b, err := json.Marshal(v)
|
||
if err != nil {
|
||
return actool.Errorf(err.Error()), nil
|
||
}
|
||
return actool.Text(string(b)), nil
|
||
}
|
||
|
||
// 本文件实现 P2「跨任务编排工具集」(docs/跑分编排 §2 P2)。这些是 host 工具——需要
|
||
// 访问 Manager(任意任务的 Store)、Engine(暂停)、以及建任务流程,所以住在 server 层。
|
||
// 读类工具把「现有 per-task 工具」重定向到目标任务的 store 上跑(建一个临时 ToolSet
|
||
// 并 Call 其对应工具),从而复用完全相同的逻辑;控制类(spawn/pause)直接调 Manager/Engine。
|
||
// 它们像流量工具一样 seed 进 tools 表、按 agent 绑定(只绑给编排 agent 才可见)。
|
||
|
||
// hostTools is the runtime host-tool provider fed to ToolAugment: traffic tools
|
||
// (gated by capture) + cross-task orchestration tools + user-defined custom tools.
|
||
// The second return is the names of custom tools flagged `deferred` (schema
|
||
// withheld, routed via SearchExtraTools/ExecuteExtraTool). Per-agent binding still
|
||
// decides who actually sees any of them.
|
||
//
|
||
//nolint:unused // used as the hostTools provider in wireAgentAugment
|
||
func (s *Server) hostTools() ([]actool.CoreTool, map[string][]string) {
|
||
tools := append(s.m.HostTools(), s.orchestrationTools()...)
|
||
tools = append(tools, s.platformTools()...) // 平台操作工具(建改 skill/工具/MCP，给 Auto 用)
|
||
custom, err := s.customTools()
|
||
if err != nil {
|
||
log.Printf("[custom-tool] 加载失败: %v", err)
|
||
return tools, nil
|
||
}
|
||
tools = append(tools, custom...)
|
||
// deferred custom tools → name -> its bound agent keys. ToolAugment turns a
|
||
// name into a deferred entry only for agents it's actually bound to (so we don't
|
||
// advertise a tool the per-agent binding will drop from the callable set).
|
||
deferred := map[string][]string{}
|
||
rows, _ := s.m.pg.ListCustomTools()
|
||
for _, t := range rows {
|
||
if t.Deferred && t.Enabled {
|
||
deferred[t.Key] = t.Agents
|
||
}
|
||
}
|
||
return tools, deferred
|
||
}
|
||
|
||
// orchestrationTools returns the cross-task tool set. Bound per-agent via the
|
||
// tools table (default: no binding — opt-in for orchestration agents).
|
||
func (s *Server) orchestrationTools() []actool.CoreTool {
|
||
return []actool.CoreTool{
|
||
s.toolListTasks(),
|
||
s.toolListLLMProfiles(),
|
||
s.toolSpawnTask(),
|
||
s.toolPauseTask(),
|
||
s.toolGetTaskGraph(),
|
||
s.toolListTaskFindings(),
|
||
s.toolAddHint(),
|
||
s.toolGetWorkerTrace(),
|
||
s.toolListWorkerTraces(),
|
||
s.toolSearchWorkerTraces(),
|
||
s.toolGetTaskNodeDetail(),
|
||
s.toolUpdateFindingReport(),
|
||
}
|
||
}
|
||
|
||
// --- schema helpers ---
|
||
|
||
func strParam(desc string) map[string]any {
|
||
return map[string]any{"type": "string", "description": desc}
|
||
}
|
||
|
||
// parseProfileID reads an LLM profile id from a tool arg that may arrive as a JSON
|
||
// number (5) or a numeric string ("5"); returns 0 when absent/unparseable.
|
||
func parseProfileID(raw json.RawMessage) int64 {
|
||
if len(raw) == 0 {
|
||
return 0
|
||
}
|
||
var n int64
|
||
if json.Unmarshal(raw, &n) == nil {
|
||
return n
|
||
}
|
||
var str string
|
||
if json.Unmarshal(raw, &str) == nil {
|
||
v, _ := strconv.ParseInt(strings.TrimSpace(str), 10, 64)
|
||
return v
|
||
}
|
||
return 0
|
||
}
|
||
|
||
func objSchema(props map[string]any, required ...string) map[string]any {
|
||
m := map[string]any{"type": "object", "properties": props}
|
||
if len(required) > 0 {
|
||
req := make([]any, len(required))
|
||
for i, r := range required {
|
||
req[i] = r
|
||
}
|
||
m["required"] = req
|
||
}
|
||
return m
|
||
}
|
||
|
||
func roTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||
return actool.Build(actool.Spec{
|
||
Name: name, Description: desc, Schema: schema,
|
||
ReadOnly: func(json.RawMessage) bool { return true },
|
||
Concurrent: func(json.RawMessage) bool { return true },
|
||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||
return permission.Allowed()
|
||
},
|
||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||
return run(ctx, in)
|
||
},
|
||
})
|
||
}
|
||
|
||
func wrTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||
return actool.Build(actool.Spec{
|
||
Name: name, Description: desc, Schema: schema,
|
||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||
return permission.Allowed()
|
||
},
|
||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||
return run(ctx, in)
|
||
},
|
||
})
|
||
}
|
||
|
||
// delegateToTask resolves the `task_id` in the input, builds a ToolSet bound to
|
||
// that task's store, strips task_id, and calls the chosen per-task tool — so the
|
||
// cross-task read reuses the exact in-task logic against another task.
|
||
func (s *Server) delegateToTask(ctx context.Context, in json.RawMessage, pick func(*agent.ToolSet) actool.CoreTool) (actool.Result, error) {
|
||
var head struct {
|
||
TaskID string `json:"task_id"`
|
||
}
|
||
_ = json.Unmarshal(in, &head)
|
||
if strings.TrimSpace(head.TaskID) == "" {
|
||
return actool.Errorf("task_id 为必填"), nil
|
||
}
|
||
t, ok := s.m.Task(head.TaskID)
|
||
if !ok {
|
||
return actool.Errorf("task 不存在: " + head.TaskID), nil
|
||
}
|
||
var m map[string]json.RawMessage
|
||
_ = json.Unmarshal(in, &m)
|
||
delete(m, "task_id")
|
||
inner, _ := json.Marshal(m)
|
||
tsx := agent.NewToolSet(t.Store, "orchestrator")
|
||
if s.m.Assets() != nil {
|
||
tsx.SetAssetStore(s.m.Assets(), s.m.Assets().Companies())
|
||
}
|
||
tsx.SetNotify(t.Notify) // hint writes wake this task's planner (no-op for read tools)
|
||
return pick(tsx).Call(ctx, inner, nil)
|
||
}
|
||
|
||
// --- tools ---
|
||
|
||
func (s *Server) toolListTasks() actool.CoreTool {
|
||
return roTool("list_tasks",
|
||
"列出所有任务(id/描述/目标/状态/运行时长/父任务/LLM 配置)，编排 agent 用它掌握全局、看哪些任务卡太久、各自用哪个 LLM。运行时长：运行中=创建→现在，终态=创建→最后活动(秒)。llm_profile：任务 planner/worker 用的配置名，(激活配置)=跟随全局激活。",
|
||
objSchema(map[string]any{}),
|
||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||
lastAct, _ := s.m.PG().LastActivityAll()
|
||
// id -> name to resolve each task's pinned LLM profile.
|
||
profName := map[int64]string{}
|
||
if profs, err := s.m.pg.ListProfiles(); err == nil {
|
||
for _, p := range profs {
|
||
profName[p.ID] = p.Name
|
||
}
|
||
}
|
||
out := make([]map[string]any, 0)
|
||
for _, t := range s.m.List() {
|
||
status := s.deriveTaskStatus(t)
|
||
end := lastAct[t.ExpID]
|
||
if live := s.engine.LastActivity(t.ID); live > end {
|
||
end = live
|
||
}
|
||
dur := int64(0)
|
||
if status == "running" {
|
||
dur = time.Now().Unix() - t.CreatedAt
|
||
} else if end > t.CreatedAt {
|
||
dur = end - t.CreatedAt
|
||
}
|
||
row := map[string]any{"id": t.ID, "description": t.Description, "goal": t.Goal, "status": status, "run_seconds": dur}
|
||
if t.ParentRef != "" {
|
||
row["parent_ref"] = t.ParentRef
|
||
}
|
||
llmState := t.llmStateSnapshot()
|
||
if llmState.ProfileID == nil {
|
||
row["llm_profile"] = "(激活配置)"
|
||
} else if n, ok := profName[*llmState.ProfileID]; ok {
|
||
row["llm_profile"] = n
|
||
} else {
|
||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *llmState.ProfileID)
|
||
}
|
||
out = append(out, row)
|
||
}
|
||
return jsonResult(out)
|
||
})
|
||
}
|
||
|
||
// toolListLLMProfiles lists the available LLM profiles (name/model/active) so an
|
||
// orchestration agent can pick one for spawn_task's llm_profile. Never leaks keys.
|
||
func (s *Server) toolListLLMProfiles() actool.CoreTool {
|
||
return roTool("list_llm_profiles",
|
||
"列出可用的 LLM 配置(profile)：id、名称、模型、格式、是否为当前激活配置。用 id 给 spawn_task 的 llm_profile_id 参数指定子任务专属 LLM（如侦察用便宜模型、利用用强模型）。不含 API Key。",
|
||
objSchema(map[string]any{}),
|
||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||
profs, err := s.m.pg.ListProfiles()
|
||
if err != nil {
|
||
return actool.Errorf(err.Error()), nil
|
||
}
|
||
out := make([]map[string]any, 0, len(profs))
|
||
for _, p := range profs {
|
||
out = append(out, map[string]any{
|
||
"id": p.ID, "name": p.Name, "model": p.Model, "format": p.Format, "is_active": p.IsDefault,
|
||
})
|
||
}
|
||
return jsonResult(map[string]any{"profiles": out})
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolSpawnTask() actool.CoreTool {
|
||
return wrTool("spawn_task",
|
||
"新建一个子任务并启动探索引擎，返回 task_id。用于把一件事(如一道题/一个目标)派成独立任务。parent_ref 可选：填当前编排关联的父任务 id 做父子关联。",
|
||
objSchema(map[string]any{
|
||
"description": strParam("任务描述(简短标题)"),
|
||
"goal": strParam("任务目标(要达成什么)"),
|
||
"parent_ref": strParam("可选：父任务 id(做父子关联)"),
|
||
"llm_profile_id": map[string]any{"type": "integer", "description": "可选：指定本子任务 planner/worker 用的 LLM 配置 id(见 list_llm_profiles)；留空则继承父任务、再回退全局激活配置"},
|
||
"timeout_seconds": map[string]any{"type": "integer", "description": "可选：任务级超时(秒)。到点后触发优雅收尾并进入 timeout 终态；留空或 0 = 不限时"},
|
||
"plan_heartbeat_seconds": map[string]any{"type": "integer", "description": "可选：planner 心跳触发间隔(秒)。距上轮规划结束/任务开始满该值且期间无触发 → 触发一轮规划(兜底死锁 + 唤醒去监督飞行中的 worker)。留空或 0 = 默认 600(10min)；"},
|
||
"seed_first_intent": map[string]any{"type": "boolean", "description": "可选：对于简单任务可开启，创建时直接下发一条种子意图(内容=描述+目标)让 worker 免等首轮 planner 直接开跑测试；默认 false(走标准先规划再执行)。"},
|
||
}, "description", "goal"),
|
||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||
var a struct {
|
||
Description, Goal, ParentRef string
|
||
LLMProfileID json.RawMessage `json:"llm_profile_id"`
|
||
TimeoutSeconds int `json:"timeout_seconds"`
|
||
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"`
|
||
SeedFirstIntent bool `json:"seed_first_intent"`
|
||
}
|
||
_ = json.Unmarshal(in, &a)
|
||
if strings.TrimSpace(a.Description) == "" {
|
||
a.Description = "未命名任务"
|
||
}
|
||
if strings.TrimSpace(a.Goal) == "" {
|
||
return actool.Errorf("goal 为必填"), nil
|
||
}
|
||
if a.TimeoutSeconds < 0 {
|
||
a.TimeoutSeconds = 0
|
||
}
|
||
// LLM profile resolution: explicit id > inherit parent's pin > active(nil).
|
||
var pin *int64
|
||
if id := parseProfileID(a.LLMProfileID); id > 0 {
|
||
if _, ok := s.loadProfileConfig(id); !ok {
|
||
return actool.Errorf(fmt.Sprintf("LLM 配置 #%d 不存在或未设置 API Key", id)), nil
|
||
}
|
||
pin = &id
|
||
} else if a.ParentRef != "" {
|
||
if pt, ok := s.m.Task(a.ParentRef); ok {
|
||
pin = pt.LLMProfileID
|
||
}
|
||
}
|
||
t, err := s.m.CreateTask(a.Description, a.Goal, pin, a.TimeoutSeconds, a.PlanHeartbeatSeconds)
|
||
if err != nil {
|
||
return actool.Errorf(err.Error()), nil
|
||
}
|
||
if a.ParentRef != "" {
|
||
t.ParentRef = a.ParentRef
|
||
if id, e := strconv.ParseInt(t.ID, 10, 64); e == nil {
|
||
_ = s.m.PG().SetParentRef(id, a.ParentRef)
|
||
}
|
||
}
|
||
// 共享的建后流程,与 HTTP 建任务(server.go createTask)复用同一段 launchTask:
|
||
// seed + 后台可见地做目标分解(第0轮/LLM步骤/逐条goal) + engine.Run。
|
||
// seed_first_intent 默认 false(标准先规划再执行);简单任务可开启直接下发一 work 测试。
|
||
s.launchTask(t, a.Description+" "+a.Goal, a.SeedFirstIntent)
|
||
return actool.Text(fmt.Sprintf("task created: %s", t.ID)), nil
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolPauseTask() actool.CoreTool {
|
||
return wrTool("pause_task", "暂停指定任务(停止其 planner/worker 循环)。",
|
||
objSchema(map[string]any{"task_id": strParam("要暂停的任务 id")}, "task_id"),
|
||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||
var a struct {
|
||
TaskID string `json:"task_id"`
|
||
}
|
||
_ = json.Unmarshal(in, &a)
|
||
if _, ok := s.m.Task(a.TaskID); !ok {
|
||
return actool.Errorf("task 不存在: " + a.TaskID), nil
|
||
}
|
||
s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator)
|
||
s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask)
|
||
return actool.Text("task paused: " + a.TaskID), nil
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolGetTaskGraph() actool.CoreTool {
|
||
return roTool("get_task_graph", "读指定任务的探索图总览(同 graph_overview：资产计数/frontier/发现/覆盖等)，用 task_id 指定任务。",
|
||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GraphOverviewTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolListTaskFindings() actool.CoreTool {
|
||
return roTool("list_task_findings", "读指定任务的确认漏洞(含 flag/PoC；每条带 id/task_id/intent_id/vulnclass/severity/摘要/状态)，用 task_id 指定任务。",
|
||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListFindingsTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolAddHint() actool.CoreTool {
|
||
return wrTool("add_task_hint", "给指定任务注入战略提示(该任务的 planner 下轮生成意图时会读到)。\n"+
|
||
"★优先批量：多条提示放进 hints 数组一次提交（返回 ids 数组，与 hints 等长同序，失败项 id=0）；单条则省略 hints 直接给顶层 text。",
|
||
objSchema(map[string]any{
|
||
"task_id": strParam("任务 id"),
|
||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组，每个元素字段同顶层（text/asset_ids）。", "items": map[string]any{"type": "object"}},
|
||
"text": strParam("[单条] 提示内容"),
|
||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id（可选，0/1/多个；该任务内的资产 id）"},
|
||
}, "task_id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).AddHintTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolGetWorkerTrace() actool.CoreTool {
|
||
return roTool("get_task_worker_trace",
|
||
"看指定任务里某个 work(意图)的执行过程：get_task_worker_trace(task_id, intent_id) 看步骤摘要；再带 step_ids=[...] 取那几步完整内容(一次≤5)。",
|
||
objSchema(map[string]any{
|
||
"task_id": strParam("任务 id"),
|
||
"intent_id": map[string]any{"type": "integer", "description": "意图 id(该任务里的 work)"},
|
||
"step_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "可选：要取完整内容的步骤 id(≤5)"},
|
||
}, "task_id", "intent_id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GetWorkerTraceTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolListWorkerTraces() actool.CoreTool {
|
||
return roTool("list_task_worker_traces", "列出指定任务里跑过哪些 work(意图) + 各自步数，用于发现哪些 work 值得翻看(再用 get_task_worker_trace)。",
|
||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListWorkerTracesTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolSearchWorkerTraces() actool.CoreTool {
|
||
return roTool("search_task_worker_traces", "在指定任务里按关键字搜索所有 work 的执行过程(返回命中步骤摘要 + intent_id)。",
|
||
objSchema(map[string]any{"task_id": strParam("任务 id"), "q": strParam("搜索关键字")}, "task_id", "q"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).SearchWorkerTracesTool)
|
||
})
|
||
}
|
||
|
||
func (s *Server) toolGetTaskNodeDetail() actool.CoreTool {
|
||
return roTool("get_task_node_detail",
|
||
"读指定任务里某个探索图节点的完整内容(发现/事实/意图/目标：摘要 + 详情/证据/PoC)。id 为探索节点 id(如 report_finding 返回、或 list_task_findings 里的 id)。写漏洞报告前用它取该漏洞的完整证据。",
|
||
objSchema(map[string]any{
|
||
"task_id": strParam("任务 id"),
|
||
"id": map[string]any{"type": "integer", "description": "探索图节点 id(非资产 id)"},
|
||
}, "task_id", "id"),
|
||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||
return s.delegateToTask(ctx, in, (*agent.ToolSet).NodeDetailTool)
|
||
})
|
||
}
|
||
|
||
// toolUpdateFindingReport writes/overwrites a finding's detailed Markdown report.
|
||
// finding_id is the id report_finding returned ("finding recorded: <id>", the
|
||
// finding node id). The write (SetFindingReportByNodeID) is keyed by node_id and
|
||
// task-agnostic, so this host tool needs no task_id / exploration store.
|
||
func (s *Server) toolUpdateFindingReport() actool.CoreTool {
|
||
return wrTool("update_finding_report",
|
||
"为已登记的漏洞写入/更新【详细报告】(Markdown 全文,整段覆盖旧内容)。finding_id 传 report_finding 返回的那个 id(\"finding recorded: <id>\" 里的数字)。报告建议包含:漏洞概述、影响与危害、复现步骤、证据/PoC、修复建议。",
|
||
objSchema(map[string]any{
|
||
"finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||
"report": strParam("详细报告全文,Markdown 格式"),
|
||
}, "finding_id", "report"),
|
||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||
var a struct {
|
||
FindingID json.RawMessage `json:"finding_id"`
|
||
Report string `json:"report"`
|
||
}
|
||
_ = json.Unmarshal(in, &a)
|
||
nodeID := parseProfileID(a.FindingID) // 复用「数字或数字字符串」解析
|
||
if nodeID <= 0 {
|
||
return actool.Errorf("finding_id 无效"), nil
|
||
}
|
||
n, err := s.m.pg.SetFindingReportByNodeID(nodeID, a.Report)
|
||
if err != nil {
|
||
return actool.Errorf(err.Error()), nil
|
||
}
|
||
if n == 0 {
|
||
return actool.Errorf(fmt.Sprintf("未找到 finding_id=%d 对应的漏洞记录(先用 report_finding 登记)", nodeID)), nil
|
||
}
|
||
return actool.Text(fmt.Sprintf("finding %d report updated (%d chars)", nodeID, len(a.Report))), nil
|
||
})
|
||
}
|
||
|
||
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
|
||
func (s *Server) deriveTaskStatus(t *Task) string {
|
||
switch {
|
||
case isTerminalStatus(t.Status):
|
||
return t.Status
|
||
case t.Paused || s.engine.IsPaused(t.ID):
|
||
return "paused"
|
||
case s.engine.ReadyFor(t) && s.engine.Started(t.ID):
|
||
return "running"
|
||
}
|
||
return "created"
|
||
}
|
||
|
||
// orchestrationToolSeeds seeds the cross-task tools into the tools table so they
|
||
// are bindable per-agent (default: bound to nobody — opt-in for orchestration
|
||
// agents). First-insert only, like the traffic seeds.
|
||
func (s *Server) seedOrchestrationTools() {
|
||
// task-op + platform tools default-bind to the built-in Auto agent (它天生用来
|
||
// 操作平台)。SeedTool 首插入生效;老库已 seed 的行由 seedAutoDefaultBindings 补绑。
|
||
autoAgents, _ := json.Marshal([]string{"auto"})
|
||
for _, t := range s.orchestrationTools() {
|
||
schema, _ := json.Marshal(t.InputSchema())
|
||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||
}
|
||
for _, t := range s.platformTools() {
|
||
schema, _ := json.Marshal(t.InputSchema())
|
||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||
}
|
||
s.refreshBuiltinToolSchemas()
|
||
s.seedAutoDefaultBindings()
|
||
s.seedPlannerDefaultBindings()
|
||
s.seedAutoReportFindingBinding()
|
||
s.unbindGoalMetDefault()
|
||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||
// 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||
// pentest 一起 seed 好了（项目尚无旧库，不做迁移）。
|
||
}
|
||
|
||
// refreshBuiltinToolSchemas propagates code schema/description changes on the
|
||
// orchestration + platform tools into already-seeded rows ONCE per version flag —
|
||
// SeedTool is first-insert-only, so a new param (e.g. spawn_task 的 llm_profile) never
|
||
// reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag.
|
||
// Bump the flag whenever these tools' schemas/descriptions change in code.
|
||
func (s *Server) refreshBuiltinToolSchemas() {
|
||
const flag = "tool_schema_refresh_v5_spawn_seed_intent"
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
tools := append(s.orchestrationTools(), s.platformTools()...)
|
||
for _, t := range tools {
|
||
schema, _ := json.Marshal(t.InputSchema())
|
||
if err := s.m.pg.RefreshToolDefaults(t.Name(), t.Description(), schema); err != nil {
|
||
log.Printf("[tools] refresh %s schema failed: %v", t.Name(), err)
|
||
}
|
||
}
|
||
// 同时把 planner 的 goal_met 描述刷成代码默认：旧库 seed 的描述带“结束本轮规划”的
|
||
// 误导，会让 planner 把 goal_met 当成“结束空轮”的手段、刚开跑就误判整个任务完成。
|
||
for _, sd := range agent.BuiltinToolSeeds() {
|
||
if sd.Key != "goal_met" {
|
||
continue
|
||
}
|
||
schema, _ := json.Marshal(sd.Schema)
|
||
if err := s.m.pg.RefreshToolDefaults(sd.Key, sd.Desc, schema); err != nil {
|
||
log.Printf("[tools] refresh %s desc failed: %v", sd.Key, err)
|
||
}
|
||
}
|
||
_ = s.m.pg.SetSetting(flag, "true")
|
||
log.Printf("[tools] 已刷新 orchestration/platform 工具 schema 到代码默认(一次性)")
|
||
}
|
||
|
||
// unbindGoalMetDefault removes goal_met's default "planner" binding ONCE (guarded by
|
||
// a settings flag), so existing DBs match the new default of NO agent. goal_met bypasses
|
||
// per-goal prove_goal to declare the whole task done — powerful/risky and redundant with
|
||
// the prove_goal→auto-complete path — so it ships unbound; users can re-bind it per agent
|
||
// in the UI. A user's own binding to another agent is untouched (we only strip planner).
|
||
func (s *Server) unbindGoalMetDefault() {
|
||
const flag = "goal_met_unbind_default_v1"
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
if err := s.m.pg.RemoveAgentFromTool("planner", "goal_met"); err != nil {
|
||
log.Printf("[tools] goal_met 解绑 planner 失败: %v", err)
|
||
return
|
||
}
|
||
_ = s.m.pg.SetSetting(flag, "true")
|
||
}
|
||
|
||
// seedReporterAgent 预置一个「报告撰写」自定义 agent(builtin=false，可在 UI 编辑/删除)：
|
||
// 绑定 update_finding_report + 任务查询工具，并挂一个「report_finding 被调用即触发」的
|
||
// 触发器 —— 每登记一个漏洞就唤起它写详细报告。一次性(settings flag 守卫)：用户删掉后不再重建。
|
||
// 依赖：orchestration 工具已在本函数上方 SeedTool 入库，故绑定得上。
|
||
func (s *Server) seedReporterAgent() {
|
||
const flag = "reporter_agent_seed_v1"
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // 无论成功与否只尝试一次
|
||
|
||
if exist, _ := s.m.pg.GetAgentByKey("reporter"); exist != nil {
|
||
return // key 已被占用(用户手建过)——不覆盖
|
||
}
|
||
a, err := s.m.pg.CreateAgent("reporter", "报告撰写",
|
||
"漏洞详细报告撰写：发现漏洞时自动触发，查取证据与执行过程后写 Markdown 报告并回写。")
|
||
if err != nil {
|
||
log.Printf("[reporter] 创建 agent 失败: %v", err)
|
||
return
|
||
}
|
||
if err := s.m.pg.SeedPromptIfEmpty(a.ID, agent.ReporterDefaultPrompt); err != nil {
|
||
log.Printf("[reporter] seed prompt 失败: %v", err)
|
||
}
|
||
// 触发运行策略：parallel + none —— 一漏洞一报告、多个 finding 并发各写各的。
|
||
// merge 必须为 none：否则(默认 all)一波 finding 会被合并成一次运行，并行就没意义。
|
||
// maxParallel=5：同时最多 5 个报告会话，避免瞬时太多 LLM 调用。
|
||
if err := s.m.pg.SetAgentTriggerBehavior("reporter", "parallel", "none", 5); err != nil {
|
||
log.Printf("[reporter] 设置触发运行策略失败: %v", err)
|
||
}
|
||
// 绑定它需要的工具：写报告 + 读证据/执行过程/态势。
|
||
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{
|
||
"update_finding_report", "get_task_node_detail", "list_task_findings",
|
||
"get_task_worker_trace", "list_task_worker_traces", "search_task_worker_traces",
|
||
"get_task_graph",
|
||
}); err != nil {
|
||
log.Printf("[reporter] 绑定工具失败: %v", err)
|
||
}
|
||
// 触发器：report_finding 被调用即触发（工具返回 "finding recorded: <id>" 带上 finding_id，
|
||
// 任务 id 也在触发消息里）。
|
||
if _, err := s.m.pg.CreateTrigger(&db.AgentTrigger{
|
||
AgentKey: "reporter",
|
||
Enabled: true,
|
||
OnToolCall: true,
|
||
ToolNames: []string{"report_finding"},
|
||
ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请从触发上下文里取出 finding_id" +
|
||
"（工具返回 \"finding recorded: <id>\" 里的数字）与任务 id，按你的职责撰写该漏洞的详细报告，" +
|
||
"最后调用 update_finding_report(finding_id, report) 保存。",
|
||
}); err != nil {
|
||
log.Printf("[reporter] 创建触发器失败: %v", err)
|
||
}
|
||
log.Printf("[reporter] 已预置「报告撰写」agent + finding 触发器")
|
||
}
|
||
|
||
// seedAutoReportFindingBinding adds "auto" to report_finding's binding ONCE so
|
||
// conversation-context agents can call it without requiring an intent_id.
|
||
func (s *Server) seedAutoReportFindingBinding() {
|
||
const flag = "auto_report_finding_v1"
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
if err := s.m.pg.AddAgentToToolBinding("auto", []string{"report_finding"}); err != nil {
|
||
log.Printf("[auto] report_finding 默认绑定失败: %v", err)
|
||
return
|
||
}
|
||
_ = s.m.pg.SetSetting(flag, "true")
|
||
}
|
||
|
||
// seedPlannerDefaultBindings adds "planner" to report_finding's binding ONCE
|
||
// (guarded by a settings flag), so existing DBs — whose report_finding row was
|
||
// seeded as worker-only — also let the planner record findings. Fresh DBs already
|
||
// get it via PlannerTools(); this only backfills without overriding a user unbind.
|
||
func (s *Server) seedPlannerDefaultBindings() {
|
||
const flag = "planner_report_finding_v1"
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
if err := s.m.pg.AddAgentToToolBinding("planner", []string{"report_finding"}); err != nil {
|
||
log.Printf("[planner] report_finding 默认绑定失败: %v", err)
|
||
return
|
||
}
|
||
_ = s.m.pg.SetSetting(flag, "true")
|
||
}
|
||
|
||
// seedAutoDefaultBindings adds "auto" to the task-op + platform tools' bindings
|
||
// ONCE (guarded by a settings flag), so existing DBs whose tool rows were seeded
|
||
// before Auto existed still give Auto its default toolset — without re-adding it
|
||
// after a user deliberately unbinds.
|
||
func (s *Server) seedAutoDefaultBindings() {
|
||
const flag = "auto_default_bindings_v3" // v3: 替换旧资产工具名，加入 insert_assets/add_company_scope
|
||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||
return
|
||
}
|
||
keys := make([]string, 0, len(platformToolKeys)+12)
|
||
for _, t := range s.orchestrationTools() {
|
||
keys = append(keys, t.Name())
|
||
}
|
||
keys = append(keys, platformToolKeys...)
|
||
// 资产工具：Auto 操作平台常要看/登记资产、管理公司范围。
|
||
keys = append(keys, "insert_assets", "add_company_scope", "list_assets")
|
||
if err := s.m.pg.AddAgentToToolBinding("auto", keys); err != nil {
|
||
log.Printf("[auto] 默认绑定失败: %v", err)
|
||
return
|
||
}
|
||
_ = s.m.pg.SetSetting(flag, "true")
|
||
}
|