mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-13 11:40:53 +08:00
A worker hitting a wall only a human can clear (captcha, email/SMS OTP, 2FA,
a decision) calls ask_operator with a precise question and an optional
screenshot. The question lands in the intent's transcript as a user-role turn
with an attachment card (the same {text, attachments} shape chat uploads use,
so the existing transcript UI renders it natively), and the intent pauses. The
operator replies through the existing message-to-worker flow; the reply resumes
the run with the answer as the next input. Images are whitelisted (png/jpg/
jpeg/webp/gif), capped at 10MB, and copied into the task uploads dir.
Closes #41
949 lines
48 KiB
Go
949 lines
48 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
|
|
}
|
|
|
|
// This file implements the P2 "cross-task orchestration tool set" (docs/benchmark-orchestration §2 P2). These are host tools — they need
|
|
// access to the Manager (any task's Store), the Engine (pause), and the task-creation flow, so they live in the server layer.
|
|
// Read tools redirect "the existing per-task tools" to run against the target task's store (build a temporary ToolSet
|
|
// and Call its corresponding tool), reusing exactly the same logic; control tools (spawn/pause) call Manager/Engine directly.
|
|
// Like the traffic tools, they are seeded into the tools table and bound per agent (visible only to agents they are bound to).
|
|
|
|
// 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.findingRetestTools()...)
|
|
tools = append(tools, s.platformTools()...) // platform operation tools (create/edit skills/tools/MCP, for Auto)
|
|
custom, err := s.customTools()
|
|
if err != nil {
|
|
log.Printf("[custom-tool] load failed: %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(),
|
|
s.toolGetFindingTraffic(),
|
|
s.toolBindFindingTraffic(),
|
|
}
|
|
}
|
|
|
|
// --- 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 is required"), nil
|
|
}
|
|
t, ok := s.m.Task(head.TaskID)
|
|
if !ok {
|
|
return actool.Errorf("task not found: " + 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) // generic wake-up (write ops without a dedicated callback go through it; no-op for read tools)
|
|
tsx.SetNotifyHint(t.NotifyHint) // add_hint → records a "human added N strategic hints: ..." entry and wakes the planner
|
|
return pick(tsx).Call(ctx, inner, nil)
|
|
}
|
|
|
|
// --- tools ---
|
|
|
|
func (s *Server) toolListTasks() actool.CoreTool {
|
|
return roTool("list_tasks",
|
|
"List all tasks (id/description/goal/status/run duration/parent task/LLM profile). The orchestration agent uses it to keep the big picture, see which tasks have been stuck too long, and which LLM each task uses. Run duration: running = created→now, terminal = created→last activity (seconds). llm_profile: the profile name used by the task's planner/worker; (active profile) = follows the global active one.",
|
|
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"] = "(active profile)"
|
|
} else if n, ok := profName[*llmState.ProfileID]; ok {
|
|
row["llm_profile"] = n
|
|
} else {
|
|
row["llm_profile"] = fmt.Sprintf("#%d(deleted)", *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",
|
|
"List the available LLM profiles: id, name, model, format, and whether one is the currently active profile. Use the id for spawn_task's llm_profile_id parameter to give a subtask its own LLM (e.g. a cheap model for recon, a strong model for exploitation). Never includes API keys.",
|
|
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",
|
|
"Create a subtask and start its exploration engine; returns task_id. Use it to spin one piece of work (e.g. a challenge/a goal) off into an independent task. parent_ref is optional: set the parent task id of the current orchestration to link parent and child.",
|
|
objSchema(map[string]any{
|
|
"description": strParam("task description (short title)"),
|
|
"goal": strParam("task goal (what to achieve)"),
|
|
"parent_ref": strParam("optional: parent task id (for parent-child linking)"),
|
|
"source_task_ids": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": fmt.Sprintf("optional: ids of source tasks to inherit read-only (at most %d). The subtask may read the assets/conclusions those tasks already explored as its starting point; unlike parent_ref's plain parent-child pointer, this is content inheritance.", db.MaxTaskSourceCount)},
|
|
"llm_profile_id": map[string]any{"type": "integer", "description": "optional: the LLM profile id this subtask's planner/worker uses (see list_llm_profiles); empty inherits the parent task, then falls back to the global active profile"},
|
|
"timeout_seconds": map[string]any{"type": "integer", "description": "optional: task-level timeout (seconds). When it fires, a graceful wrap-up runs and the task enters the timeout terminal state; empty or 0 = no time limit"},
|
|
"plan_heartbeat_seconds": map[string]any{"type": "integer", "description": "optional: planner heartbeat interval (seconds). When this much time has passed since the last planning round ended/the task started with no trigger in between → run one planning round (deadlock backstop + wakes the planner to supervise in-flight workers). Empty or 0 = default 600 (10min);"},
|
|
"seed_first_intent": map[string]any{"type": "boolean", "description": "optional: enable for simple tasks to issue one seed intent at creation (content = description + goal) so a worker can start testing right away without waiting for the first planner round; default false (the standard plan-then-execute flow)."},
|
|
}, "description", "goal"),
|
|
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
|
var a struct {
|
|
Description string `json:"description"`
|
|
Goal string `json:"goal"`
|
|
ParentRef string `json:"parent_ref"`
|
|
SourceTaskIDs []string `json:"source_task_ids"`
|
|
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 = "Unnamed task"
|
|
}
|
|
if strings.TrimSpace(a.Goal) == "" {
|
|
return actool.Errorf("goal is required"), nil
|
|
}
|
|
if a.TimeoutSeconds < 0 {
|
|
a.TimeoutSeconds = 0
|
|
}
|
|
// Read-only source-task inheritance: count cap + each id valid/deduped/existing, the same validation rules as HTTP task creation.
|
|
if len(a.SourceTaskIDs) > db.MaxTaskSourceCount {
|
|
return actool.Errorf(fmt.Sprintf("at most %d source tasks may be selected", db.MaxTaskSourceCount)), nil
|
|
}
|
|
sourceIDs := make([]int64, 0, len(a.SourceTaskIDs))
|
|
seenSources := map[int64]bool{}
|
|
for _, raw := range a.SourceTaskIDs {
|
|
id, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64)
|
|
if err != nil || id <= 0 || seenSources[id] {
|
|
return actool.Errorf("invalid or duplicate source task id"), nil
|
|
}
|
|
if _, ok := s.m.Task(strconv.FormatInt(id, 10)); !ok {
|
|
return actool.Errorf(fmt.Sprintf("source task #%d not found", id)), nil
|
|
}
|
|
seenSources[id] = true
|
|
sourceIDs = append(sourceIDs, id)
|
|
}
|
|
// 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 profile #%d not found or has no API key set", id)), nil
|
|
}
|
|
pin = &id
|
|
} else if a.ParentRef != "" {
|
|
if pt, ok := s.m.Task(a.ParentRef); ok {
|
|
pin = pt.LLMProfileID
|
|
}
|
|
}
|
|
var llmIDs []int64
|
|
if pin != nil {
|
|
llmIDs = []int64{*pin}
|
|
}
|
|
t, err := s.m.CreateTaskWithOptions(a.Description, a.Goal, db.TaskCreateOptions{
|
|
SourceTaskIDs: sourceIDs,
|
|
LLMProfileIDs: llmIDs,
|
|
TimeoutSeconds: a.TimeoutSeconds,
|
|
PlanHeartbeatSeconds: 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)
|
|
}
|
|
}
|
|
// Shared post-create flow: HTTP task creation (server.go createTask) reuses this same launchTask:
|
|
// seed + a visibly-in-background goal decomposition (round 0 / LLM steps / per-goal) + engine.Run.
|
|
// seed_first_intent defaults to false (the standard plan-then-execute flow); simple tasks may enable it to issue one work test directly.
|
|
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", "Pause the given task (stops its planner/worker loops).",
|
|
objSchema(map[string]any{"task_id": strParam("the id of the task to pause")}, "task_id"),
|
|
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
|
var a struct {
|
|
TaskID string `json:"task_id"`
|
|
}
|
|
_ = json.Unmarshal(in, &a)
|
|
t, ok := s.m.Task(a.TaskID)
|
|
if !ok {
|
|
return actool.Errorf("task not found: " + a.TaskID), nil
|
|
}
|
|
if _, err := s.applyTaskControlWithCause(t, "pause", agent.AbortPausedByOrchestrator); err != nil {
|
|
return actool.Errorf(err.Error()), nil
|
|
}
|
|
return actool.Text("task paused: " + a.TaskID), nil
|
|
})
|
|
}
|
|
|
|
func (s *Server) toolGetTaskGraph() actool.CoreTool {
|
|
return roTool("get_task_graph", "Read the exploration-graph overview of the given task (same as graph_overview: asset counts / frontier / findings / coverage etc.), selecting the task via task_id.",
|
|
objSchema(map[string]any{"task_id": strParam("task 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", "Read the confirmed vulnerabilities of the given task (with flags/PoCs; each carries id/task_id/intent_id/vulnclass/severity/summary/status), selecting the task via task_id.",
|
|
objSchema(map[string]any{"task_id": strParam("task 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", "Inject a strategic hint into the given task (that task's planner will read it when generating intents next round).\n"+
|
|
"★Batch first: submit several hints at once in the hints array (returns an ids array, same length and order as hints, failed entries id=0); for a single hint, omit hints and pass the top-level text directly.",
|
|
objSchema(map[string]any{
|
|
"task_id": strParam("task id"),
|
|
"hints": map[string]any{"type": "array", "description": "[Prefer this] Array of hints; each element has the same fields as the top level (text/asset_ids/traffic_refs).", "items": objSchema(map[string]any{"text": strParam("hint text"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()})},
|
|
"text": strParam("[single] hint text"),
|
|
"traffic_refs": agent.HintTrafficSchema(),
|
|
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "anchored asset ids (optional, 0/1/many; asset ids within that task)"},
|
|
}, "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",
|
|
"View how one work (intent) in the given task ran: get_task_worker_trace(task_id, intent_id) shows step summaries; then pass step_ids=[...] to fetch the full content of those steps (at most 5 per call; passing more returns only the first 5).",
|
|
objSchema(map[string]any{
|
|
"task_id": strParam("task id"),
|
|
"intent_id": map[string]any{"type": "integer", "description": "intent id (the work within that task)"},
|
|
"step_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "optional: step ids to fetch full content for (at most 5 per call; passing more returns only the first 5, the rest are listed in omitted_step_ids)"},
|
|
}, "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", "List which works (intents) have run in the given task + each one's step count, to find which works are worth reviewing (then use get_task_worker_trace).",
|
|
objSchema(map[string]any{"task_id": strParam("task 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", "Search the execution traces of all works in the given task by keyword (returns matching step summaries + intent_id).",
|
|
objSchema(map[string]any{"task_id": strParam("task id"), "q": strParam("search keyword")}, "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",
|
|
"Read the full content of one exploration-graph node in the given task (finding/fact/intent/goal: summary + details/evidence/PoC). id is the exploration node id (e.g. as returned by report_finding, or the id in list_task_findings). Use it to fetch a vulnerability's full evidence before writing its report.",
|
|
objSchema(map[string]any{
|
|
"task_id": strParam("task id"),
|
|
"id": map[string]any{"type": "integer", "description": "exploration-graph node id (not an asset 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",
|
|
"Write/update the [detailed report] of a registered vulnerability (full Markdown text, replaces the old content wholesale). Pass as finding_id the id report_finding returned (the number in \"finding recorded: <id>\"). A report should cover: vulnerability overview, impact and harm, reproduction steps, evidence/PoC, and fix recommendations.",
|
|
objSchema(map[string]any{
|
|
"finding_id": map[string]any{"type": "integer", "description": "target vulnerability id (the id returned by report_finding)"},
|
|
"report": strParam("full report text, Markdown format"),
|
|
"evidence_version": map[string]any{"type": "integer", "description": "evidence version returned by get_finding_traffic; prevents the report from overwriting newer evidence changes"},
|
|
}, "finding_id", "report"),
|
|
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
|
var a struct {
|
|
EvidenceVersion *int64 `json:"evidence_version"`
|
|
FindingID json.RawMessage `json:"finding_id"`
|
|
Report string `json:"report"`
|
|
}
|
|
_ = json.Unmarshal(in, &a)
|
|
nodeID := parseProfileID(a.FindingID) // reuses the "number or numeric string" parser
|
|
if nodeID <= 0 {
|
|
return actool.Errorf("invalid finding_id"), nil
|
|
}
|
|
n, err := s.m.pg.SetFindingReportVersionByNodeID(ctx, nodeID, a.Report, a.EvidenceVersion)
|
|
if err != nil {
|
|
return actool.Errorf(err.Error()), nil
|
|
}
|
|
if n == 0 {
|
|
return actool.Errorf(fmt.Sprintf("no vulnerability record found for finding_id=%d (register it with report_finding first)", 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 {
|
|
lifecycle := t.lifecycleSnapshot()
|
|
switch {
|
|
case isTerminalStatus(lifecycle.Status):
|
|
return lifecycle.Status
|
|
case lifecycle.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 (it exists to
|
|
// operate the platform). SeedTool takes effect on first insert; rows already seeded in old DBs get the binding backfilled by seedAutoDefaultBindings.
|
|
autoAgents, _ := json.Marshal([]string{"auto"})
|
|
for _, t := range s.orchestrationTools() {
|
|
schema, _ := json.Marshal(t.InputSchema())
|
|
bindings := autoAgents
|
|
if t.Name() == "bind_finding_traffic" {
|
|
bindings = json.RawMessage(`["reporter"]`)
|
|
}
|
|
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, bindings)
|
|
}
|
|
for _, t := range s.platformTools() {
|
|
schema, _ := json.Marshal(t.InputSchema())
|
|
bindings := autoAgents
|
|
if t.Name() == "mcp_status" {
|
|
// health visibility is for the executors too, not just the platform agent
|
|
bindings = json.RawMessage(`["auto","worker","pentest"]`)
|
|
}
|
|
if t.Name() == "ask_operator" {
|
|
// operator input requests are an executor concern (captcha/OTP walls)
|
|
bindings = json.RawMessage(`["worker","pentest"]`)
|
|
}
|
|
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, bindings)
|
|
}
|
|
s.refreshBuiltinToolSchemas()
|
|
s.seedAutoDefaultBindings()
|
|
s.seedPlannerDefaultBindings()
|
|
s.seedPlannerListAssetsBinding()
|
|
s.seedCompanyScopeRebind()
|
|
s.seedWorkerReadToolsUnbind() // list_facts/list_companies/list_worker_traces unbound from worker's defaults (one-time)
|
|
s.seedWorkerReadbackRebind() // repairs an old migration's wrong removal: re-binds search_all_worker_traces/get_worker_trace/node_detail back onto worker (one-time)
|
|
s.seedTaskCredentialBindings() // save/list_task_credential onto worker + pentest (one-time, #3)
|
|
s.seedAutoReportFindingBinding()
|
|
s.unbindGoalMetDefault()
|
|
s.reseedGoalsPrompt() // goals prompt gains the "extract operation constraints" step → old DBs append a new default version (one-time)
|
|
s.reseedMainAgentPrompt() // mainagent prompt gains "after goals are met, ask via add_intent whether to register a goal" (one-time)
|
|
s.reseedPlannerPrompt() // planner prompt: rewrites the "zero intents" justification rules + adds quantitative acceptance checks (one-time)
|
|
s.reseedWorkerPrompt() // worker prompt: adds the evidence bar for negative conclusions (one-time)
|
|
s.seedReporterAgent() // seeds the "report writer" agent + tool bindings + finding trigger (one-time)
|
|
s.upgradeReporterTriggerMessage() // backfill migration for old DBs: makes the reporter send evidence_version back (one-time)
|
|
s.seedFindingTrafficTools() // adds optional evidence params and read-only evidence tools, preserving user config
|
|
s.seedFindingWorkflowTools()
|
|
// Note: pentest's default tool bindings need no migration — BuiltinToolSeeds already seeds
|
|
// list_assets/insert_assets/report_finding/list_findings/list_companies together with
|
|
// pentest at fresh initialization (no old DBs exist in this project, so no migration).
|
|
}
|
|
|
|
// 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's 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_v7_list_facts_paging"
|
|
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)
|
|
}
|
|
}
|
|
// Also refresh some built-in agent tools to the code defaults:
|
|
// - goal_met: the description seeded in old DBs carries the misleading "end this planning round", making the planner treat it as
|
|
// a way to "end an empty round" and wrongly declare the whole task done right after start.
|
|
// - insert_assets: new related param (marks whether the asset relates to the current task and decides whether it counts toward coverage),
|
|
// SeedTool is first-insert-only, so schemas already seeded in old DBs would otherwise miss this new param.
|
|
// - list_facts: now paginated with new limit/before/q params; the empty schema already seeded in old DBs would otherwise
|
|
// show "no parameters" on the tools page, and the model would never see those param descriptions.
|
|
refreshBuiltin := map[string]bool{"goal_met": true, "insert_assets": true, "list_facts": true}
|
|
for _, sd := range agent.BuiltinToolSeeds() {
|
|
if !refreshBuiltin[sd.Key] {
|
|
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] refreshed orchestration/platform tool schemas to code defaults (one-time)")
|
|
}
|
|
|
|
// 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] failed to unbind goal_met from planner: %v", err)
|
|
return
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
// reseedGoalsPrompt refreshes the goals decomposer's prompt to [the current code default] — the default body gained the step
|
|
// "extract operation constraints (set_constraints) before decomposing goals", and SeedPromptIfEmpty is first-insert-only, so old DBs
|
|
// with an existing version 1 never got it. Here we use version management to [append a new version] and switch to it (ResetPromptToDefault),
|
|
// old versions stay in history, and users who customized can recover from the version record. settings flag guard → done only once;
|
|
// bump this flag whenever the default changes again. Fresh DBs need nothing (SeedPromptIfEmpty already seeds the latest default).
|
|
func (s *Server) reseedGoalsPrompt() {
|
|
const flag = "goals_prompt_constraint_step_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // attempt only once regardless of outcome
|
|
a, err := s.m.pg.GetAgentByKey("goals")
|
|
if err != nil || a == nil {
|
|
return // fresh DB without agent rows yet: seedPrompts seeds the latest default directly, this migration is unnecessary
|
|
}
|
|
tmpl := agent.BuiltinPromptSeeds()["goals"]
|
|
if tmpl == "" {
|
|
return
|
|
}
|
|
// Fresh DB: seedPrompts already seeded the latest default → the current version equals the code default, no duplicate version needed.
|
|
if cur, err := s.m.pg.CurrentPrompt(a.ID); err == nil && cur == tmpl {
|
|
return
|
|
}
|
|
if _, err := s.m.pg.ResetPromptToDefault(a.ID, tmpl); err != nil {
|
|
log.Printf("[prompts] failed to refresh goals prompt to the new default: %v", err)
|
|
return
|
|
}
|
|
log.Printf("[prompts] goals prompt: appended a new default version (adds the extract-constraints step, one-time)")
|
|
}
|
|
|
|
// reseedMainAgentPrompt refreshes the mainagent prompt to [the current code default] — the default body gained the guidance "after all goals are
|
|
// met, when injecting an intent directly via add_intent, ask the human whether to register it as a formal goal", and SeedPromptIfEmpty is first-insert
|
|
// only, so versions already in old DBs never got it. Version management [appends a new version] and switches to it (ResetPromptToDefault); old versions stay
|
|
// in history, and users who customized can recover from the version record. settings flag guard → done only once. Fresh DBs need nothing
|
|
// (SeedPromptIfEmpty already seeds the latest default). Structurally identical to reseedGoalsPrompt.
|
|
func (s *Server) reseedMainAgentPrompt() {
|
|
const flag = "mainagent_prompt_goalless_intent_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // attempt only once regardless of outcome
|
|
a, err := s.m.pg.GetAgentByKey("mainagent")
|
|
if err != nil || a == nil {
|
|
return // fresh DB without agent rows yet: seedPrompts seeds the latest default directly, this migration is unnecessary
|
|
}
|
|
tmpl := agent.BuiltinPromptSeeds()["mainagent"]
|
|
if tmpl == "" {
|
|
return
|
|
}
|
|
// Fresh DB: seedPrompts already seeded the latest default → the current version equals the code default, no duplicate version needed.
|
|
if cur, err := s.m.pg.CurrentPrompt(a.ID); err == nil && cur == tmpl {
|
|
return
|
|
}
|
|
if _, err := s.m.pg.ResetPromptToDefault(a.ID, tmpl); err != nil {
|
|
log.Printf("[prompts] failed to refresh mainagent prompt to the new default: %v", err)
|
|
return
|
|
}
|
|
log.Printf("[prompts] mainagent prompt: appended a new default version (adds the ask-about-registering-a-goal-after-met step, one-time)")
|
|
}
|
|
|
|
// reseedPlannerPrompt refreshes the planner prompt to [the current code default] — the default body was slimmed down and "restraint" demoted to
|
|
// dedup-only, adding "depth over coverage" and the "hard floor: with unmet goals and no running intents, output is mandatory", and bounding the re-check of negative conclusions.
|
|
// Bump the flag below (currently v2) whenever the default changes materially so existing old DBs refresh once more. SeedPromptIfEmpty is first-insert-only; versions already in old DBs never got it, so version management
|
|
// [appends a new version] and switches to it (ResetPromptToDefault); old versions stay in history, and users who customized can recover from the version record.
|
|
// settings flag guard → done only once. Fresh DBs need nothing (SeedPromptIfEmpty already seeds the latest default). Structurally identical to
|
|
// reseedGoalsPrompt.
|
|
func (s *Server) reseedPlannerPrompt() {
|
|
const flag = "planner_prompt_compact_realistic_v2"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // attempt only once regardless of outcome
|
|
a, err := s.m.pg.GetAgentByKey("planner")
|
|
if err != nil || a == nil {
|
|
return // fresh DB without agent rows yet: seedPrompts seeds the latest default directly, this migration is unnecessary
|
|
}
|
|
tmpl := agent.BuiltinPromptSeeds()["planner"]
|
|
if tmpl == "" {
|
|
return
|
|
}
|
|
// Fresh DB: seedPrompts already seeded the latest default → the current version equals the code default, no duplicate version needed.
|
|
if cur, err := s.m.pg.CurrentPrompt(a.ID); err == nil && cur == tmpl {
|
|
return
|
|
}
|
|
if _, err := s.m.pg.ResetPromptToDefault(a.ID, tmpl); err != nil {
|
|
log.Printf("[prompts] failed to refresh planner prompt to the new default: %v", err)
|
|
return
|
|
}
|
|
log.Printf("[prompts] planner prompt: appended a new default version (slimmed rework + restraint demoted to dedup + depth-first + bounded negative re-checks, one-time)")
|
|
}
|
|
|
|
// reseedWorkerPrompt refreshes the worker prompt to [the current code default] — the record_fact section of the default body dropped the whole "for negative conclusions,
|
|
// write observations + tentative readings" sentence and decoupled confidence (observed/inferred) from "whether this intent's options are exhausted" (these easily mislead the planner),
|
|
// and tightened per-item facts arrays to the rare exception of "fully independent of each other, impossible to merge". Bumped the flag to v3 so existing old DBs refresh once more.
|
|
// SeedPromptIfEmpty is first-insert-only; versions already in old DBs never got it, so version management [appends a new version] and switches to it; old versions stay in history and can be recovered.
|
|
// settings flag guard → done only once. Fresh DBs need nothing. Structurally identical to reseedGoalsPrompt.
|
|
func (s *Server) reseedWorkerPrompt() {
|
|
const flag = "worker_prompt_compact_v4"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // attempt only once regardless of outcome
|
|
a, err := s.m.pg.GetAgentByKey("worker")
|
|
if err != nil || a == nil {
|
|
return // fresh DB without agent rows yet: seedPrompts seeds the latest default directly, this migration is unnecessary
|
|
}
|
|
tmpl := agent.BuiltinPromptSeeds()["worker"]
|
|
if tmpl == "" {
|
|
return
|
|
}
|
|
// Fresh DB: seedPrompts already seeded the latest default → the current version equals the code default, no duplicate version needed.
|
|
if cur, err := s.m.pg.CurrentPrompt(a.ID); err == nil && cur == tmpl {
|
|
return
|
|
}
|
|
if _, err := s.m.pg.ResetPromptToDefault(a.ID, tmpl); err != nil {
|
|
log.Printf("[prompts] failed to refresh worker prompt to the new default: %v", err)
|
|
return
|
|
}
|
|
log.Printf("[prompts] worker prompt: appended a new default version (context-lookup section narrowed to list_assets/list_findings, dropping list_facts/node_detail/asset_neighbors, one-time)")
|
|
}
|
|
|
|
// reporterToolCallMessage must unconditionally require reading get_finding_traffic once before writing the report.
|
|
// That tool is read-only and "does not depend on the capture switch": manually bound evidence is readable whether auto-binding is on or off. If this said
|
|
// "read only when auto-binding is enabled", with the default-off config the reporter would not pass evidence_version,
|
|
// SetFindingReportVersionByNodeID would then write -1 with legacy semantics, and the vulnerability details and Markdown export
|
|
// would be stuck on "evidence changed, report needs updating" forever, with no UI entry able to clear it.
|
|
const reporterToolCallMessage = "A vulnerability was just registered via report_finding. Read finding_id (the standalone vulnerability record ID) and finding_node_id (the exploration node ID) from the returned JSON, " +
|
|
"then read the current evidence list and its version with get_finding_traffic(finding_id) first (an empty list is normal — write the report as usual); " +
|
|
"if your run instructions enable auto-binding, verify and associate this vulnerability's traffic before reading. Use finding_node_id for node details." +
|
|
"Finally call update_finding_report(finding_id=finding_node_id, report, evidence_version=the version actually read) to save; " +
|
|
"evidence_version must be passed, or the report will be permanently marked as needing an update. Do not mix up the two kinds of id."
|
|
|
|
// Legacy trigger message (0.3.8 and earlier). Only records still matching it verbatim are migrated/overwritten; user-modified ones stay as-is.
|
|
const reporterToolCallMessageV1 = "A vulnerability was just registered via report_finding. Take finding_id out of the trigger context" +
|
|
" (the number in the tool's returned \"finding recorded: <id>\") and the task id, write that vulnerability's detailed report per your duties," +
|
|
" and finally call update_finding_report(finding_id, report) to save it."
|
|
|
|
// upgradeReporterTriggerMessage refreshes reporter trigger messages that still carry the default copy in old DBs to the new version.
|
|
// seedReporterAgent is guarded by reporter_agent_seed_v1 and only writes the trigger when creating the agent, so
|
|
// upgraded DBs never got the new copy — seedFindingTrafficTools backfilled the tool schema with
|
|
// evidence_version, but nothing told the reporter to use it. One-time, and it only overwrites unmodified copy.
|
|
func (s *Server) upgradeReporterTriggerMessage() {
|
|
const flag = "reporter_trigger_evidence_version_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // attempt once
|
|
triggers, err := s.m.pg.ListTriggersFor("reporter")
|
|
if err != nil {
|
|
log.Printf("[reporter] failed to read triggers: %v", err)
|
|
return
|
|
}
|
|
for _, t := range triggers {
|
|
if !t.OnToolCall || t.ToolCallMessage != reporterToolCallMessageV1 {
|
|
continue // user-modified or not a finding trigger — leave it alone.
|
|
}
|
|
t.ToolCallMessage = reporterToolCallMessage
|
|
if err := s.m.pg.UpdateTrigger(t); err != nil {
|
|
log.Printf("[reporter] failed to upgrade trigger message: %v", err)
|
|
return
|
|
}
|
|
log.Printf("[reporter] trigger message upgraded to read and return evidence_version")
|
|
}
|
|
}
|
|
|
|
// seedReporterAgent presets a "report writer" custom agent (builtin=false, editable/deletable in the UI):
|
|
// it binds update_finding_report + the task-query tools and carries a trigger that fires "whenever report_finding is called"
|
|
// — every registered vulnerability wakes it to write the detailed report. One-time (settings flag guard): once the user deletes it, it is never recreated.
|
|
// Dependency: the orchestration tools were seeded via SeedTool just above this function, so the binding can attach.
|
|
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") }() // attempt only once regardless of outcome
|
|
|
|
if exist, _ := s.m.pg.GetAgentByKey("reporter"); exist != nil {
|
|
return // key already taken (user created it) — do not overwrite
|
|
}
|
|
a, err := s.m.pg.CreateAgent("reporter", "Report writer",
|
|
"Detailed vulnerability report writing: triggers automatically when a vulnerability is found, fetches evidence and execution traces, then writes a Markdown report and writes it back.")
|
|
if err != nil {
|
|
log.Printf("[reporter] failed to create agent: %v", err)
|
|
return
|
|
}
|
|
if err := s.m.pg.SeedPromptIfEmpty(a.ID, agent.ReporterDefaultPrompt); err != nil {
|
|
log.Printf("[reporter] seed prompt failed: %v", err)
|
|
}
|
|
// Trigger run policy: parallel + none — one report per vulnerability, several findings concurrently each writing its own.
|
|
// merge must be none: otherwise (default all) a wave of findings merges into one run and parallelism is pointless.
|
|
// maxParallel=5: at most 5 report sessions at once, avoiding a burst of too many LLM calls.
|
|
if err := s.m.pg.SetAgentTriggerBehavior("reporter", "parallel", "none", 5); err != nil {
|
|
log.Printf("[reporter] failed to set trigger run policy: %v", err)
|
|
}
|
|
// Bind the tools it needs: report writing + reading evidence/execution traces/situation.
|
|
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] failed to bind tools: %v", err)
|
|
}
|
|
// Trigger: fires whenever report_finding is called (the tool returns "finding recorded: <id>" carrying the finding_id,
|
|
// and the task id is also in the trigger message).
|
|
if _, err := s.m.pg.CreateTrigger(&db.AgentTrigger{
|
|
AgentKey: "reporter",
|
|
Enabled: true,
|
|
OnToolCall: true,
|
|
ToolNames: []string{"report_finding"},
|
|
ToolCallMessage: reporterToolCallMessage,
|
|
}); err != nil {
|
|
log.Printf("[reporter] failed to create trigger: %v", err)
|
|
}
|
|
log.Printf("[reporter] seeded the 'report writer' agent + finding trigger")
|
|
}
|
|
|
|
// 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 default binding failed: %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 default binding failed: %v", err)
|
|
return
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
// seedPlannerListAssetsBinding adds "planner" to list_assets's binding ONCE
|
|
// (guarded by a settings flag), so existing DBs — whose list_assets row was seeded
|
|
// as auto/pentest-only — also let the planner query the asset store by DSL. Fresh
|
|
// DBs already get it via PlannerTools(); this only backfills without overriding a
|
|
// user unbind.
|
|
func (s *Server) seedPlannerListAssetsBinding() {
|
|
const flag = "planner_list_assets_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
if err := s.m.pg.AddAgentToToolBinding("planner", []string{"list_assets"}); err != nil {
|
|
log.Printf("[planner] list_assets default binding failed: %v", err)
|
|
return
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
// seedCompanyScopeRebind changes add_company_scope's default binding ONCE on
|
|
// existing DBs (guarded by a settings flag): the tool moves off worker and onto
|
|
// planner — defining a company's asset scope is a planning/main/auto concern, not
|
|
// something a worker does mid-exploration. Fresh DBs already get planner via
|
|
// PlannerTools() and lack worker via WorkerTools(); this only backfills old rows.
|
|
// One-shot + flag-guarded so a user who later re-binds worker isn't overridden.
|
|
func (s *Server) seedCompanyScopeRebind() {
|
|
const flag = "company_scope_rebind_v1" // worker→planner default-binding switch
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
if err := s.m.pg.AddAgentToToolBinding("planner", []string{"add_company_scope"}); err != nil {
|
|
log.Printf("[planner] add_company_scope default binding failed: %v", err)
|
|
return
|
|
}
|
|
if err := s.m.pg.RemoveAgentFromTool("worker", "add_company_scope"); err != nil {
|
|
log.Printf("[worker] failed to unbind add_company_scope: %v", err)
|
|
return
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
// seedWorkerReadToolsUnbind strips the read-context tools off worker's default
|
|
// binding ONCE on existing DBs (guarded by a settings flag): a worker executes one
|
|
// intent and writes back — reading facts/companies and listing all workers' traces is
|
|
// a planning/main concern, not the executor's. Fresh DBs already lack these via
|
|
// WorkerTools(); this only backfills old rows without overriding a user who
|
|
// deliberately re-binds worker. Each RemoveAgentFromTool is per-tool +
|
|
// membership-guarded, so planner/mainagent bindings of the same tool are untouched.
|
|
//
|
|
// NOTE: search_all_worker_traces / get_worker_trace / node_detail are intentionally NOT
|
|
// unbound — worker owns them for cross-work look-back + node drill-down (see WorkerTools).
|
|
// They used to be in this list back when worker lacked them; seedWorkerReadbackRebind
|
|
// repairs DBs whose old run stripped them.
|
|
func (s *Server) seedWorkerReadToolsUnbind() {
|
|
const flag = "worker_readtools_unbind_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
for _, k := range []string{
|
|
"list_facts", "list_companies", "list_worker_traces",
|
|
} {
|
|
if err := s.m.pg.RemoveAgentFromTool("worker", k); err != nil {
|
|
log.Printf("[worker] failed to unbind %s from worker: %v", k, err)
|
|
return // on error the flag is not set; retried on next startup
|
|
}
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
// seedWorkerReadbackRebind re-binds the cross-work look-back / drill-down tools onto
|
|
// worker ONCE (guarded by a settings flag): an earlier seedWorkerReadToolsUnbind wrongly
|
|
// stripped search_all_worker_traces / get_worker_trace / node_detail from worker after
|
|
// they had been added to WorkerTools(), so any DB that ran that migration lost them.
|
|
// Fresh DBs already have them via WorkerTools() and this is a harmless no-op there.
|
|
// One-shot + flag-guarded so a user who later deliberately unbinds them isn't overridden.
|
|
// seedTaskCredentialBindings binds the task-credential tools onto worker and
|
|
// pentest ONCE (guarded by a settings flag), so existing DBs seeded before #3
|
|
// still get them — fresh DBs already have them via WorkerTools()/the pentest
|
|
// catalog seed. A user who deliberately unbinds them is not overridden.
|
|
func (s *Server) seedTaskCredentialBindings() {
|
|
const flag = "task_credential_bindings_v1"
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
for _, agentKey := range []string{"worker", "pentest"} {
|
|
if err := s.m.pg.AddAgentToToolBinding(agentKey, []string{"save_task_credential", "list_task_credentials"}); err != nil {
|
|
log.Printf("[%s] task-credential default binding failed: %v", agentKey, err)
|
|
return // flag not set; retried on next startup
|
|
}
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|
|
|
|
func (s *Server) seedWorkerReadbackRebind() {
|
|
const flag = "worker_readback_rebind_v2" // v2: adds node_detail
|
|
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
|
return
|
|
}
|
|
if err := s.m.pg.AddAgentToToolBinding("worker", []string{
|
|
"search_all_worker_traces", "get_worker_trace", "node_detail",
|
|
}); err != nil {
|
|
log.Printf("[worker] failed to re-bind look-back/detail tools: %v", err)
|
|
return // on error the flag is not set; retried on next startup
|
|
}
|
|
_ = 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: replaces old asset tool names, adds 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...)
|
|
// Asset tools: Auto operating the platform often views/registers assets and manages company scopes.
|
|
keys = append(keys, "insert_assets", "add_company_scope", "list_assets")
|
|
if err := s.m.pg.AddAgentToToolBinding("auto", keys); err != nil {
|
|
log.Printf("[auto] default binding failed: %v", err)
|
|
return
|
|
}
|
|
_ = s.m.pg.SetSetting(flag, "true")
|
|
}
|