Autumn-27
52aba07846
主 agent 之前只能 add_hint/add_intent,无法对已在运行的 work 实时纠偏。 现将 planner 的 steer_work(不打断、不丢进展地注入纠偏指令)默认下放给主 agent。 - MainAgentTools 加入 steerWorkTool - MainAgent 加 steerWork 字段与 SetSteerWork,Chat 里接进 ToolSet - 全局与 per-task router 两处主 agent 构建均 SetSteerWork(engine.SteerWork) - 主 agent 提示词「操舵」段补充 steer_work 用法引导(仅影响新 seed) 作用域安全:steer_work 仍校验 intent_id 属于本任务(关联任务只读)。 Claude-Session: https://claude.ai/code/session_019QXSCDmzimwnjntBNCgUD7
154 lines
8.4 KiB
Go
154 lines
8.4 KiB
Go
package agent
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
|
||
"github.com/Autumn-27/artex/db"
|
||
"github.com/Autumn-27/norma/agentcore"
|
||
"github.com/Autumn-27/norma/llm"
|
||
"github.com/Autumn-27/norma/permission"
|
||
actool "github.com/Autumn-27/norma/tool"
|
||
"github.com/Autumn-27/norma/transcript"
|
||
)
|
||
|
||
// MainAgent is the thin human-interface orchestrator (docs §4.2 / §7). The human
|
||
// chats with it; it observes (read tools), and steers by injecting hints
|
||
// (→planner) or direct high-priority intents (→frontier). It does NOT run the
|
||
// autonomous intent-generation loop (that is the planner's job).
|
||
type MainAgent struct {
|
||
prov llm.Provider
|
||
model string
|
||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||
window int // context window in tokens (for compaction)
|
||
windowFn func() int // optional dynamic task-chain minimum
|
||
maxTurns int // max agent turns per run (0 = unlimited)
|
||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||
steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off)
|
||
}
|
||
|
||
func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store, window, maxTurns int) *MainAgent {
|
||
return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns}
|
||
}
|
||
|
||
func (m *MainAgent) SetCompactionWindowResolver(fn func() int) { m.windowFn = fn }
|
||
|
||
func (m *MainAgent) compactionWindow() int {
|
||
if m.windowFn != nil {
|
||
return m.windowFn()
|
||
}
|
||
return m.window
|
||
}
|
||
|
||
// SetProxy points the main agent's WebFetch at the recording proxy plus the CA
|
||
// cert it trusts to verify HTTPS through it (empty addr = direct).
|
||
func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert }
|
||
|
||
// SetWebSearch selects the web_search backend for the main agent (off by default).
|
||
func (m *MainAgent) SetWebSearch(o WebSearchOpts) { m.webSearch = o }
|
||
|
||
// SetSteerWork wires the engine callback that lets the main agent's steer_work
|
||
// tool inject a mid-run course-correction into a running work (nil = tool off).
|
||
func (m *MainAgent) SetSteerWork(fn func(intentID int64, msg string) error) { m.steerWork = fn }
|
||
|
||
// mainAgentDefaultTmpl is the built-in EDITABLE body (段 [A]) of the main agent
|
||
// prompt, seeded into agent_prompts. Goal is a {{.Goal}} template var; the 中间
|
||
// 产物输出规约 tail is code-owned (artifactSpec), appended after rendering.
|
||
const mainAgentDefaultTmpl = `你是一个授权渗透测试系统的"主 agent"，是人类操作员的接口。你不亲自探索、也不自主连续生成意图（那是规划者的工作）。你的职责：
|
||
|
||
1. 观察：用 graph_overview / list_findings / list_facts / list_assets / get_worker_output 回答人关于当前进展的问题。
|
||
2. 操舵（把人的意图落到系统）：
|
||
- 人想"改方向/强调某类漏洞/重点某区域" → 用 add_hint 写提示（规划者下次会读到）。
|
||
- 人想"立刻测某个具体目标" → 用 add_intent 直接注入一条高优先级意图（priority 8-10）。
|
||
- 人想"对某条正在运行的意图(work)实时纠偏（别再走 X、聚焦 Y）" → 用 steer_work（不打断、不丢已有进展，worker 下一步动作前生效）；先用 get_worker_output 看它在干嘛。方向整个错了则改用 add_intent 另下新意图。
|
||
- 人想"新增一个要达成的最终目标" → 用 set_goals 增补目标。系统会把该目标写入任务图并**自动把已完成/暂停的任务拉回运行态继续跑**（规划者随后会据此重新判断是否达成），无需人工再点恢复。
|
||
3. 用人话简洁回复，说明你做了什么。
|
||
|
||
当前任务目标：{{.Goal}}
|
||
|
||
不要编造发现；只根据工具返回的真实数据回答。`
|
||
|
||
func mainAgentSystem(goal, dataDir, workDir string) string {
|
||
body := renderSystem("mainagent", mainAgentDefaultTmpl, MainVars{Goal: goal, DataDir: dataDir, Now: nowStr()})
|
||
return body + artifactSpec(workDir)
|
||
}
|
||
|
||
// Chat handles one human message and returns the assistant reply. emit, if
|
||
// 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, resume func(), notifyGoal func([]string)) (string, error) {
|
||
tsx := NewToolSet(ts, "human")
|
||
if as != nil {
|
||
tsx.SetAssetStore(as, as.Companies())
|
||
}
|
||
tsx.SetTaskID(taskID)
|
||
tsx.SetCoverageEnabled(as == nil || as.CoverageEnabled(taskID))
|
||
tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced)
|
||
tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running
|
||
tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标：…」触发
|
||
tsx.steerWork = m.steerWork // enable steer_work tool (nil = unavailable)
|
||
// 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash）
|
||
// 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets（不入 prompt）。
|
||
base := append(tsx.DropCoverageTools(tsx.MainAgentTools()), actool.DefaultTools()...)
|
||
ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
|
||
tools, def, cleanup := AugmentTools(ctx, "mainagent", base)
|
||
tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参
|
||
defer cleanup()
|
||
// 本任务的工作目录 <workDir>/tasks/<taskID>，先建好。
|
||
mainDir := ensureRunDir(m.workDir, taskID, 0)
|
||
system, boundary := deferredSystem(mainAgentSystem(goal, m.workDir, mainDir), def)
|
||
opts := agentcore.Options{
|
||
Provider: m.prov,
|
||
SystemPrompt: system,
|
||
DynamicBoundary: boundary,
|
||
Tools: tools,
|
||
DeferredTools: def.Deferred,
|
||
UnlockSet: def.Unlock,
|
||
PermissionMode: permission.ModeBypass,
|
||
EnableWebFetch: true, // 走记录代理留痕；载入代理 CA 验证 MITM 重签的 HTTPS 证书
|
||
WebFetchProxy: m.proxyAddr,
|
||
WebFetchCACert: m.proxyCACert,
|
||
// 联网搜索(可选)。ddgs 无需 key；brave-free 需 BraveKey；tavily 需 TavilyKey。
|
||
// WebSearchProxy 是独立出口代理(http/https/socks5)，与记录流量的 MITM 代理无关；空则直连。
|
||
EnableWebSearch: m.webSearch.Enabled,
|
||
WebSearchBackend: m.webSearch.Backend,
|
||
BraveSearchAPIKey: m.webSearch.BraveKey,
|
||
TavilySearchAPIKey: m.webSearch.TavilyKey,
|
||
WebSearchProxy: m.webSearch.Proxy,
|
||
BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA
|
||
WorkingDir: mainDir, // 本任务工作目录 <workDir>/tasks/<taskID>
|
||
ToolOutputDir: cmdOutDir(mainDir),
|
||
MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management)
|
||
Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window
|
||
Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢
|
||
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。
|
||
Settlement: wrapupSettlement("mainagent", nil),
|
||
}
|
||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task
|
||
opts.Transcript = m.tx
|
||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||
}
|
||
s := agentcore.NewSession(opts)
|
||
defer s.Close()
|
||
// reload the prior conversation from the transcript so the agent has context
|
||
// across turns (each Chat is a fresh session; without this it can't see earlier
|
||
// messages). First turn: no file yet → Resume loads nothing and proceeds.
|
||
if m.tx != nil {
|
||
_ = s.Resume(opts.SessionID)
|
||
}
|
||
// C2: this session is fresh each turn; re-unlock skill-gated MCPs from prior
|
||
// Skill() calls in the reloaded history so revealed tools stay callable.
|
||
seedUnlockFromHistory(s.Messages(), def.UnlockSkill)
|
||
text, _, err := captureRunSession(ctx, s, message, func(r db.Activity) {
|
||
if emit != nil {
|
||
r.Worker = "mainagent"
|
||
emit(r)
|
||
}
|
||
})
|
||
return text, err
|
||
}
|