Merge origin/main into codex/btw-side-question
合入 main 的 norma v0.3.7、会话分组与流量证据更新，保留 /btw 快照和输入布局。同步自动放行、关闭流量绑定的测试断言，并等待归档测试夹具的旁路后台退出。
This commit is contained in:
ji ruo committed
commit
2930c4b915
94 files changed
+9737
-1241
No files matched your search
@@ -210,6 +210,7 @@ func seedPrompts(pg *db.DB) {
|
||||
// flag and, if kept, wrapped so the model sees the DB-overridden description/schema
|
||||
// and缺省入参 get injected. MCP/skill/host tools have no row and pass through.
|
||||
func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) {
|
||||
agent.FindingTrafficBindingEnabled = func() bool { return pg.GetBool(settingAgentTrafficBinding, false) }
|
||||
// Seed the built-in domain tools (first-insert only; DO NOTHING preserves edits).
|
||||
// No startup prune: rows we didn't seed are left alone so future user-defined
|
||||
// custom tools (system=false, added via the UI) survive restarts.
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestConversationListRunningState(t *testing.T) {
|
||||
s, _ := newRetestServer(t)
|
||||
var ids []int64
|
||||
for _, key := range []string{"auto", "reporter", "retester"} {
|
||||
c, err := s.m.pg.CreateConversation(key, "runtime status test", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ids = append(ids, c.ID)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
s.chatMu.Lock()
|
||||
clear(s.chatBusy)
|
||||
s.chatMu.Unlock()
|
||||
for _, id := range ids {
|
||||
_, _ = s.m.pg.Exec(`DELETE FROM conversations WHERE id=$1`, id)
|
||||
}
|
||||
})
|
||||
check := func(want map[int64]bool) {
|
||||
t.Helper()
|
||||
w := retestRequest(s.pgListConversations, http.MethodGet, 0, "")
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", w.Code, w.Body)
|
||||
}
|
||||
var body struct {
|
||||
Conversations []conversationListItem `json:"conversations"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := 0
|
||||
for _, item := range body.Conversations {
|
||||
if expected, ok := want[item.ID]; ok {
|
||||
found++
|
||||
if item.Running != expected || item.Title != "runtime status test" {
|
||||
t.Fatalf("conversation %d: running=%v want=%v title=%q", item.ID, item.Running, expected, item.Title)
|
||||
}
|
||||
}
|
||||
}
|
||||
if found != len(want) {
|
||||
t.Fatalf("found %d of %d conversations", found, len(want))
|
||||
}
|
||||
}
|
||||
check(map[int64]bool{ids[0]: false, ids[1]: false, ids[2]: false})
|
||||
s.chatMu.Lock()
|
||||
s.chatBusy[s.convBusyKey(ids[1])] = true
|
||||
s.chatBusy[s.convBusyKey(ids[2])] = true
|
||||
s.chatBusy["unrelated-task"] = true
|
||||
s.chatMu.Unlock()
|
||||
check(map[int64]bool{ids[0]: false, ids[1]: true, ids[2]: true})
|
||||
s.chatMu.Lock()
|
||||
clear(s.chatBusy)
|
||||
s.chatMu.Unlock()
|
||||
check(map[int64]bool{ids[0]: false, ids[1]: false, ids[2]: false})
|
||||
}
|
||||
+73
-6
@@ -46,6 +46,11 @@ func decodeConversationRequest(w http.ResponseWriter, r *http.Request, value any
|
||||
// conversation_activities and the browser POLLS ?since=cursor for live updates
|
||||
// (no per-conversation SSE broadcaster needed).
|
||||
|
||||
type conversationListItem struct {
|
||||
*db.Conversation
|
||||
Running bool `json:"running"`
|
||||
}
|
||||
|
||||
func (s *Server) pgListConversations(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
@@ -56,7 +61,13 @@ func (s *Server) pgListConversations(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"conversations": cs})
|
||||
items := make([]conversationListItem, 0, len(cs))
|
||||
s.chatMu.Lock()
|
||||
for _, c := range cs {
|
||||
items = append(items, conversationListItem{Conversation: c, Running: s.chatBusy[s.convBusyKey(c.ID)]})
|
||||
}
|
||||
s.chatMu.Unlock()
|
||||
writeJSON(w, 200, map[string]any{"conversations": items})
|
||||
}
|
||||
|
||||
func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -448,7 +459,8 @@ func (s *Server) pgSendConversationMessage(w http.ResponseWriter, r *http.Reques
|
||||
// conversation_activities. Shared by the chat HTTP handler and the P3 scheduler.
|
||||
// busyKey clears when the run ends (best-effort in-flight marker).
|
||||
func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) {
|
||||
go s.runConversationSync(c, msg, busyKey)
|
||||
ctx, cancel := s.conversationRunContext(c.ID, busyKey)
|
||||
go s.runConversationTurn(ctx, cancel, c, msg, busyKey)
|
||||
}
|
||||
|
||||
// runConversationSync runs ONE agent turn on a conversation and BLOCKS until the
|
||||
@@ -456,12 +468,23 @@ func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) {
|
||||
// fire-and-forget chat path; the P3 trigger queue calls it directly so it can wait
|
||||
// for completion before starting the next queued fire for the same agent.
|
||||
func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) {
|
||||
ctx, cancel := s.conversationRunContext(c.ID, busyKey)
|
||||
s.runConversationTurn(ctx, cancel, c, msg, busyKey)
|
||||
}
|
||||
|
||||
// Register cancellation before returning 202, so an immediate stop/delete cannot
|
||||
// miss a background goroutine which has not started yet.
|
||||
func (s *Server) conversationRunContext(id int64, busyKey string) (context.Context, context.CancelCauseFunc) {
|
||||
// Per-run cancellable context so a manual stop (pgStopConversation) can abort
|
||||
// just this session. Registered under chatMu so the stop handler can find it.
|
||||
ctx, cancel := context.WithCancelCause(intercept.WithConvID(s.ctx, c.ID))
|
||||
ctx, cancel := context.WithCancelCause(intercept.WithConvID(s.ctx, id))
|
||||
s.chatMu.Lock()
|
||||
s.chatCancel[busyKey] = cancel
|
||||
s.chatMu.Unlock()
|
||||
return ctx, cancel
|
||||
}
|
||||
|
||||
func (s *Server) runConversationTurn(ctx context.Context, cancel context.CancelCauseFunc, c *db.Conversation, msg, busyKey string) {
|
||||
defer func() {
|
||||
cancel(agent.AbortChatTurnFinished)
|
||||
s.chatMu.Lock()
|
||||
@@ -469,9 +492,48 @@ func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) {
|
||||
delete(s.chatCancel, busyKey)
|
||||
s.chatMu.Unlock()
|
||||
}()
|
||||
// Only the first turn executes a historical retest. Follow-up conversation
|
||||
// turns may explain the sealed result; the result tool refuses to overwrite it.
|
||||
finishStatus, finishReason := "failed", "复测未能启动"
|
||||
if c.AgentKey == db.FindingRetestAgentKey {
|
||||
// Read without the run cancellation so an immediate stop still seals pending.
|
||||
r, err := s.m.pg.FindingRetestForConversation(context.Background(), c.ID)
|
||||
if err != nil {
|
||||
// The sealing defer below needs r.ID, which we do not have here. Seal by
|
||||
// conversation instead, otherwise the row stays 'pending' forever.
|
||||
log.Printf("[conv %d] load retest: %v", c.ID, err)
|
||||
if err := s.m.pg.FailPendingRetestForConversation(c.ID, "复测状态读取失败，请重新发起"); err != nil {
|
||||
log.Printf("[conv %d] seal retest: %v", c.ID, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if r != nil && r.Status == "pending" {
|
||||
defer func() {
|
||||
if ctx.Err() != nil {
|
||||
finishStatus, finishReason = "stopped", "复测已停止或服务已关闭"
|
||||
}
|
||||
s.finishRetest(r.ID, finishStatus, finishReason)
|
||||
}()
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
started, err := s.m.pg.StartFindingRetest(ctx, r.ID)
|
||||
if err != nil {
|
||||
finishReason = err.Error()
|
||||
return
|
||||
}
|
||||
if !started {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
// precedence: the conversation agent's own binding → this conversation's pin → global.
|
||||
ca := s.resolveChatAgent(c)
|
||||
if ca == nil {
|
||||
finishReason = s.chatUnavailableReason()
|
||||
return
|
||||
}
|
||||
pg := s.m.pg
|
||||
@@ -489,9 +551,14 @@ func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) {
|
||||
}
|
||||
// On a manual stop ctx is cancelled; Chat already emits a clean "已手动停止"
|
||||
// step, so skip the raw-error entry — only surface genuine failures.
|
||||
if _, err := ca.Chat(ctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil && ctx.Err() == nil {
|
||||
_, _ = pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "text", IsError: true,
|
||||
Summary: "（出错：" + err.Error() + "）", Detail: err.Error()})
|
||||
if _, err := ca.Chat(ctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil {
|
||||
finishReason = err.Error()
|
||||
if ctx.Err() == nil {
|
||||
_, _ = pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "text", IsError: true,
|
||||
Summary: "（出错：" + err.Error() + "）", Detail: err.Error()})
|
||||
}
|
||||
} else {
|
||||
finishStatus, finishReason = "completed", ""
|
||||
}
|
||||
_ = pg.TouchConversation(c.ID)
|
||||
}
|
||||
|
||||
+12
-3
@@ -292,12 +292,18 @@ func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO {
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
TrafficCount int `json:"traffic_count"`
|
||||
EvidenceVersion int64 `json:"evidence_version"`
|
||||
ReportEvidenceVersion int64 `json:"report_evidence_version"`
|
||||
ReportStale bool `json:"report_stale"`
|
||||
TrafficBindings []db.FindingTrafficBinding `json:"traffic_bindings,omitempty"`
|
||||
|
||||
ID string `json:"id"`
|
||||
FindingID string `json:"finding_id,omitempty"` // standalone findings-table id — the handle for status updates
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Name string `json:"name,omitempty"` // 漏洞名称;为空时前端回退展示 vulnclass
|
||||
Severity string `json:"severity"` // critical | high | medium | low
|
||||
Status string `json:"status"` // pending | in_progress | confirmed | resolved | false_positive | ignored | duplicate | risk_accepted
|
||||
Status string `json:"status"` // pending | in_progress | confirmed | resolved | fixed | false_positive | ignored | duplicate | risk_accepted
|
||||
Summary string `json:"summary"`
|
||||
Evidence string `json:"evidence"`
|
||||
Report string `json:"report,omitempty"` // 详细报告(Markdown);仅详情接口返回,列表为空
|
||||
@@ -399,6 +405,7 @@ func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int
|
||||
if m, ok := meta[n.ID]; ok {
|
||||
d.FindingID = i64s(m.ID)
|
||||
d.Status = m.Status
|
||||
d.TrafficCount = m.TrafficCount
|
||||
d.Assets = findingAssetDTOs(m.AssetIDs, assets)
|
||||
}
|
||||
out = append(out, d)
|
||||
@@ -426,8 +433,10 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO {
|
||||
status = db.FindingPending
|
||||
}
|
||||
d := FindingDTO{
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
TrafficCount: f.TrafficCount, EvidenceVersion: f.EvidenceVersion, ReportEvidenceVersion: f.ReportEvidenceVersion,
|
||||
ReportStale: f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion, TrafficBindings: f.TrafficBindings,
|
||||
VulnClass: f.VulnClass,
|
||||
Name: f.Name,
|
||||
Severity: f.Severity,
|
||||
|
||||
+76
-4
@@ -527,9 +527,52 @@ func (e *Engine) drainSteer(intentID int64) (string, bool) {
|
||||
// before each tool call it drains a queued course-correction (if any) and blocks the
|
||||
// call, handing the message back to the model — which re-plans its next step instead
|
||||
// of running the tool. No queued message → the guard behaves exactly as before.
|
||||
// 它同时负责「空转回合」的续跑，见 Stop。
|
||||
type steerHooks struct {
|
||||
inner harness.HookRunner
|
||||
drain func() (string, bool)
|
||||
// nudges 是本条意图已注入的空转续跑次数，上限 limit。指针:harness 持有的是
|
||||
// steerHooks 的值拷贝，计数必须共享同一份。
|
||||
nudges *atomic.Int64
|
||||
// limit 是空转续跑的次数上限，由 Engine.emptyTurnNudgeLimit() 从「空响应重试
|
||||
// 次数」解析而来。<=0 = 不介入(用户显式关掉了这层)。
|
||||
limit int
|
||||
// label 形如 "worker-1 · #42"，只用于日志。
|
||||
label string
|
||||
}
|
||||
|
||||
// 空转回合(只有思考、既无正文也无工具调用)续跑次数的默认值，与 SDK 空响应重试的
|
||||
// 内置默认(norma/llm/openai.go 的 emptyResponseRetries)保持一致——两层共用同一个
|
||||
// 旋钮，不配置时的行为也该对齐。解析见 Engine.emptyTurnNudgeLimit。
|
||||
//
|
||||
// 注意这个数是「一条意图的总量」，不是「连续几次」:harness 自身的 stopHookActive
|
||||
// 已经限死了连续空转只推一次——推完那一轮若还是空转，Stop 钩子不会再被调到，run 直接
|
||||
// 收场;只有真正发生过一次工具回合，配额才刷新(norma/harness/query.go:534)。所以这个
|
||||
// 闸挡的是「工具 → 空转 → 推 → 工具 → 空转」这种病态循环，别让它把意图预算耗光。
|
||||
const defaultEmptyTurnNudges = 2
|
||||
|
||||
// emptyTurnNudge 是空转回合注入的续跑指令。
|
||||
//
|
||||
// 这种回合在 harness 眼里是一次自然结束(stop_reason=end_turn 且无 tool_use)，五层
|
||||
// LLM 重试一层都不适用——它不是错误，是模型「想完了但没动手」。SDK 的空响应重试也
|
||||
// 够不着:它以「有没有 yield 过事件」判空，而思考增量本身就是事件(norma/llm/openai.go
|
||||
// 的 SEThinkingDelta)，所以 thinking-only 不算空。何况那层是原样重发整个 prompt，
|
||||
// 对这种由上下文形状决定的空转，重发只会让模型再想一遍。这里换成追加一条指令，让它
|
||||
// 带着已经产出的思考继续，输入变了才有理由给出不同的行为。
|
||||
const emptyTurnNudge = "【空转提醒】你上一轮只输出了思考过程，既没有给出正文回复，也没有调用任何工具，" +
|
||||
"这一轮等于没有产出。请直接执行你刚才想好的下一步：要么调用工具，要么给出结论文字。不要重复思考。"
|
||||
|
||||
// isThinkingOnlyTurn reports whether the latest assistant turn produced neither
|
||||
// text nor a tool call — i.e. the model spent the whole round thinking.
|
||||
func isThinkingOnlyTurn(messages []llm.Message) bool {
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
m := messages[i]
|
||||
if m.Role != llm.RoleAssistant {
|
||||
continue
|
||||
}
|
||||
return strings.TrimSpace(m.Text()) == "" && len(m.ToolUses()) == 0
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (h steerHooks) PreToolUse(ctx context.Context, name string, input []byte) (bool, string, []byte) {
|
||||
@@ -549,11 +592,31 @@ func (h steerHooks) PostToolUse(ctx context.Context, name string, input, result
|
||||
}
|
||||
}
|
||||
|
||||
// Stop 在 guard 原有语义之上补一层「空转回合」续跑:模型只输出了思考、既没给正文
|
||||
// 也没调工具时，harness 会把它当成自然结束并以空 summary 收场(query.go 的
|
||||
// ReasonCompleted + asst.Text())，一条本来还没做完的意图就这样断在半路。此时注入
|
||||
// 一条续跑指令，让模型带着已有思考接着走。
|
||||
func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []string, string) {
|
||||
var (
|
||||
prevent bool
|
||||
blocking []string
|
||||
msg string
|
||||
)
|
||||
if h.inner != nil {
|
||||
return h.inner.Stop(ctx, messages)
|
||||
prevent, blocking, msg = h.inner.Stop(ctx, messages)
|
||||
}
|
||||
return false, nil, ""
|
||||
// inner 已经决定硬停、或已经要注入自己的续跑消息 → 尊重它，不再叠加。
|
||||
// limit<=0 = 用户把「空响应重试次数」配成了 -1，即显式关掉这层。
|
||||
if prevent || len(blocking) > 0 || h.nudges == nil || h.limit <= 0 || !isThinkingOnlyTurn(messages) {
|
||||
return prevent, blocking, msg
|
||||
}
|
||||
n := h.nudges.Add(1)
|
||||
if n > int64(h.limit) {
|
||||
log.Printf("[work %s] 空转回合(仅思考、无正文无工具)已达续跑上限 %d，放行收场", h.label, h.limit)
|
||||
return prevent, blocking, msg
|
||||
}
|
||||
log.Printf("[work %s] 空转回合(仅思考、无正文无工具)，注入续跑指令 (%d/%d)", h.label, n, h.limit)
|
||||
return false, []string{emptyTurnNudge}, ""
|
||||
}
|
||||
|
||||
// KillWork cancels the in-flight work running intentID (planner's kill_work tool).
|
||||
@@ -956,8 +1019,17 @@ func (e *Engine) runIntent(ctx context.Context, t *Task, name string, worker *ag
|
||||
a.NodeID, a.Worker = &nid, name
|
||||
emit(a)
|
||||
}
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit)
|
||||
hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }}
|
||||
label := fmt.Sprintf("%s · #%d", name, iid)
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, label, taskEmit)
|
||||
// nudges 有意建在 model_error 重跑循环之外:空转续跑的上限是「这条意图」的总量，
|
||||
// 重跑一轮不该把额度清零重来。
|
||||
hooks := steerHooks{
|
||||
inner: t.Guard.Hooks(),
|
||||
drain: func() (string, bool) { return e.drainSteer(iid) },
|
||||
nudges: &atomic.Int64{},
|
||||
limit: e.emptyTurnNudgeLimit(),
|
||||
label: label,
|
||||
}
|
||||
wTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.BeginLLMCall(t.ID)
|
||||
var reason harness.TerminalReason
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// 空转回合(仅思考、无正文无工具)的识别与续跑，见 steerHooks.Stop。
|
||||
|
||||
func assistantThinking(text string) llm.Message {
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
|
||||
{Type: llm.BlockThinking, Thinking: text, Signature: "sig"},
|
||||
}}
|
||||
}
|
||||
|
||||
func TestIsThinkingOnlyTurn(t *testing.T) {
|
||||
toolUse := llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
|
||||
{Type: llm.BlockThinking, Thinking: "先扫端口"},
|
||||
{Type: llm.BlockToolUse, ID: "t1", Name: "run_nuclei"},
|
||||
}}
|
||||
cases := []struct {
|
||||
name string
|
||||
msgs []llm.Message
|
||||
want bool
|
||||
}{
|
||||
{"仅思考", []llm.Message{llm.UserText("开始"), assistantThinking("想想")}, true},
|
||||
{"思考+工具", []llm.Message{llm.UserText("开始"), toolUse}, false},
|
||||
{"思考+正文", []llm.Message{assistantThinking("想想"), {
|
||||
Role: llm.RoleAssistant,
|
||||
Content: []llm.ContentBlock{{Type: llm.BlockThinking, Thinking: "x"}, llm.TextBlock("结论")},
|
||||
}}, false},
|
||||
{"正文只有空白字符", []llm.Message{{
|
||||
Role: llm.RoleAssistant,
|
||||
Content: []llm.ContentBlock{{Type: llm.BlockThinking, Thinking: "x"}, llm.TextBlock(" \n ")},
|
||||
}}, true},
|
||||
{"完全空的 assistant 回合", []llm.Message{{Role: llm.RoleAssistant}}, true},
|
||||
// 工具结果是 user 角色，判定必须回溯到它前面那条 assistant，而不是就近误判。
|
||||
{"最后一条是工具结果", []llm.Message{toolUse, {
|
||||
Role: llm.RoleUser,
|
||||
Content: []llm.ContentBlock{{Type: llm.BlockToolResult, ToolUseID: "t1"}},
|
||||
}}, false},
|
||||
{"没有 assistant 消息", []llm.Message{llm.UserText("开始")}, false},
|
||||
{"空历史", nil, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
if got := isThinkingOnlyTurn(c.msgs); got != c.want {
|
||||
t.Fatalf("isThinkingOnlyTurn = %v, want %v", got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// fakeHooks 是一个可编程的 inner HookRunner，用来验证 steerHooks 对 inner 决定的尊重。
|
||||
type fakeHooks struct {
|
||||
prevent bool
|
||||
blocking []string
|
||||
msg string
|
||||
}
|
||||
|
||||
func (f fakeHooks) PreToolUse(context.Context, string, []byte) (bool, string, []byte) {
|
||||
return false, "", nil
|
||||
}
|
||||
func (f fakeHooks) PostToolUse(context.Context, string, []byte, []byte, bool) {}
|
||||
func (f fakeHooks) Stop(context.Context, []llm.Message) (bool, []string, string) {
|
||||
return f.prevent, f.blocking, f.msg
|
||||
}
|
||||
|
||||
func TestSteerHooksStopNudgesEmptyTurn(t *testing.T) {
|
||||
empty := []llm.Message{assistantThinking("我应该先枚举子域名")}
|
||||
|
||||
t.Run("空转回合注入续跑指令", func(t *testing.T) {
|
||||
h := steerHooks{nudges: &atomic.Int64{}, limit: defaultEmptyTurnNudges, label: "worker-1 · #1"}
|
||||
prevent, blocking, _ := h.Stop(context.Background(), empty)
|
||||
if prevent {
|
||||
t.Fatal("空转回合不应硬停")
|
||||
}
|
||||
if len(blocking) != 1 || blocking[0] != emptyTurnNudge {
|
||||
t.Fatalf("blocking = %v, want [emptyTurnNudge]", blocking)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("有正文或工具时不介入", func(t *testing.T) {
|
||||
h := steerHooks{nudges: &atomic.Int64{}, limit: defaultEmptyTurnNudges}
|
||||
normal := []llm.Message{{
|
||||
Role: llm.RoleAssistant,
|
||||
Content: []llm.ContentBlock{llm.TextBlock("已完成扫描，未发现开放端口")},
|
||||
}}
|
||||
if _, blocking, _ := h.Stop(context.Background(), normal); blocking != nil {
|
||||
t.Fatalf("正常收场被误判为空转: %v", blocking)
|
||||
}
|
||||
if n := h.nudges.Load(); n != 0 {
|
||||
t.Fatalf("未介入时不应计数, got %d", n)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("达到上限后放行收场", func(t *testing.T) {
|
||||
const limit = 5 // 用户把「空响应重试次数」配成 5
|
||||
h := steerHooks{nudges: &atomic.Int64{}, limit: limit}
|
||||
for i := 1; i <= limit; i++ {
|
||||
if _, blocking, _ := h.Stop(context.Background(), empty); len(blocking) != 1 {
|
||||
t.Fatalf("第 %d 次应仍在配额内, blocking = %v", i, blocking)
|
||||
}
|
||||
}
|
||||
if _, blocking, _ := h.Stop(context.Background(), empty); blocking != nil {
|
||||
t.Fatalf("超出上限仍在注入: %v", blocking)
|
||||
}
|
||||
})
|
||||
|
||||
// 「空响应重试次数」配 -1 = 关掉这层，emptyTurnNudgeLimit 解析成 0。
|
||||
t.Run("配置关闭时不介入", func(t *testing.T) {
|
||||
h := steerHooks{nudges: &atomic.Int64{}, limit: 0}
|
||||
if _, blocking, _ := h.Stop(context.Background(), empty); blocking != nil {
|
||||
t.Fatalf("已关闭仍在注入: %v", blocking)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("inner 决定硬停时不叠加", func(t *testing.T) {
|
||||
h := steerHooks{inner: fakeHooks{prevent: true, msg: "guard 拒绝收场"}, nudges: &atomic.Int64{}, limit: defaultEmptyTurnNudges}
|
||||
prevent, blocking, msg := h.Stop(context.Background(), empty)
|
||||
if !prevent || msg != "guard 拒绝收场" || blocking != nil {
|
||||
t.Fatalf("inner 的硬停被改写: prevent=%v blocking=%v msg=%q", prevent, blocking, msg)
|
||||
}
|
||||
if n := h.nudges.Load(); n != 0 {
|
||||
t.Fatalf("让位给 inner 时不应消耗配额, got %d", n)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("inner 已要续跑时不叠加", func(t *testing.T) {
|
||||
h := steerHooks{inner: fakeHooks{blocking: []string{"guard 的续跑理由"}}, nudges: &atomic.Int64{}, limit: defaultEmptyTurnNudges}
|
||||
_, blocking, _ := h.Stop(context.Background(), empty)
|
||||
if len(blocking) != 1 || blocking[0] != "guard 的续跑理由" {
|
||||
t.Fatalf("inner 的续跑消息被改写: %v", blocking)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("未装计数器时行为不变", func(t *testing.T) {
|
||||
h := steerHooks{limit: defaultEmptyTurnNudges} // 例如未来其他调用点忘了传 nudges
|
||||
if _, blocking, _ := h.Stop(context.Background(), empty); blocking != nil {
|
||||
t.Fatalf("无计数器时不应注入: %v", blocking)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func (s *Server) listActiveFindingRetests(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
items, err := pg.ListActiveFindingRetests(r.Context())
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"retests": items})
|
||||
}
|
||||
|
||||
func (s *Server) listFindingRetests(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok || id <= 0 {
|
||||
writeErr(w, 400, "bad finding id")
|
||||
return
|
||||
}
|
||||
f, err := pg.GetFinding(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if f == nil {
|
||||
writeErr(w, 404, "finding not found")
|
||||
return
|
||||
}
|
||||
items, err := pg.ListFindingRetests(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"retests": items})
|
||||
}
|
||||
|
||||
func (s *Server) startFindingRetest(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok || id <= 0 {
|
||||
writeErr(w, 400, "bad finding id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Notes string `json:"notes"`
|
||||
}
|
||||
if !decodeConversationRequest(w, r, &req) {
|
||||
return
|
||||
}
|
||||
req.Notes = strings.TrimSpace(req.Notes)
|
||||
if utf8.RuneCountInString(req.Notes) > 4000 {
|
||||
writeErr(w, 400, "复测补充说明最多 4000 个字符")
|
||||
return
|
||||
}
|
||||
f, err := pg.GetFinding(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if f == nil {
|
||||
writeErr(w, 404, "finding not found")
|
||||
return
|
||||
}
|
||||
a, err := pg.GetAgentByKey(db.FindingRetestAgentKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if a == nil || !a.Enabled {
|
||||
writeErr(w, 409, "漏洞复测 Agent 不存在或未启用，请在 Agent 管理中配置 retester")
|
||||
return
|
||||
}
|
||||
for _, key := range []string{"get_finding_retest_context", "record_finding_retest_result"} {
|
||||
t, err := pg.GetTool(key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if t == nil || !t.Enabled || !slices.Contains(t.Agents, a.Key) {
|
||||
writeErr(w, 409, "请为复测 Agent 启用并绑定工具："+key)
|
||||
return
|
||||
}
|
||||
}
|
||||
if s.resolveChatAgent(&db.Conversation{AgentKey: a.Key}) == nil {
|
||||
writeErr(w, 503, s.chatUnavailableReason())
|
||||
return
|
||||
}
|
||||
if s.ctx.Err() != nil {
|
||||
writeErr(w, 503, "服务正在停止")
|
||||
return
|
||||
}
|
||||
retest, conv, created, err := pg.CreateFindingRetest(r.Context(), id, req.Notes)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
writeErr(w, 404, "finding not found")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if created {
|
||||
busyKey := s.convBusyKey(conv.ID)
|
||||
s.chatMu.Lock()
|
||||
s.chatBusy[busyKey] = true
|
||||
s.chatMu.Unlock()
|
||||
s.runConversation(conv, retest.InitialMessage(), busyKey)
|
||||
}
|
||||
code := http.StatusOK
|
||||
if created {
|
||||
code = http.StatusAccepted
|
||||
}
|
||||
writeJSON(w, code, map[string]any{"retest": retest, "created": created})
|
||||
}
|
||||
|
||||
// Both tools derive the finding from server-owned conversation context. Tool
|
||||
// arguments cannot redirect a result into a different finding or conversation.
|
||||
func (s *Server) findingRetestTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
roTool("get_finding_retest_context", "读取当前复测会话关联的漏洞证据快照、复测状态、补充说明与当前任务约束。无参数，只能读取本会话。",
|
||||
objSchema(map[string]any{}), func(ctx context.Context, _ json.RawMessage) (actool.Result, error) {
|
||||
r, err := s.m.pg.FindingRetestForConversation(ctx, intercept.ConvIDFromContext(ctx))
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if r == nil {
|
||||
return actool.Errorf("当前会话未关联复测记录，请从漏洞详情发起复测"), nil
|
||||
}
|
||||
var constraints []db.Constraint
|
||||
f, err := s.m.pg.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if f != nil && f.TaskID != nil {
|
||||
if task, ok := s.m.Task(strconv.FormatInt(*f.TaskID, 10)); ok {
|
||||
constraints, err = task.Store.ListConstraints()
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return jsonResult(map[string]any{"retest": r, "current_constraints": constraints})
|
||||
}),
|
||||
wrTool("record_finding_retest_result", "为当前复测会话保存唯一结论；原漏洞证据与报告保持不变。会话成功结束且结论为 fixed 时，系统自动将漏洞状态改为已修复；其他结论保留原状态。必须提供本次实际检查的证据，无法确认时写明阻塞原因。",
|
||||
objSchema(map[string]any{
|
||||
"verdict": map[string]any{"type": "string", "enum": []string{"reproduced", "fixed", "inconclusive"}},
|
||||
"summary": strParam("本次复测结论摘要"),
|
||||
"evidence": strParam("Markdown：本次实际步骤、观察、对照、结论依据；无法确认则列出已检查内容和阻塞原因"),
|
||||
}, "verdict", "summary", "evidence"), func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
Verdict string `json:"verdict"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence string `json:"evidence"`
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if err := s.m.pg.RecordFindingRetestResult(ctx, intercept.ConvIDFromContext(ctx), a.Verdict, a.Summary, a.Evidence); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(map[string]any{"saved": true, "verdict": a.Verdict})
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Seed the editable agent atomically, without an automatic discovery trigger.
|
||||
// Once seeded, user edits/deletion survive restarts; a pre-existing key is kept.
|
||||
func (s *Server) seedFindingRetester() error {
|
||||
for _, t := range s.findingRetestTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
bindings, _ := json.Marshal([]string{db.FindingRetestAgentKey})
|
||||
if err := s.m.pg.SeedTool(t.Name(), t.Description(), schema, bindings); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
const flag = "finding_retester_seed_v1"
|
||||
tx, err := s.m.pg.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
// Lock this migration, including concurrent server initialization.
|
||||
if _, err = tx.Exec(`SELECT pg_advisory_xact_lock(7337741010)`); err != nil {
|
||||
return err
|
||||
}
|
||||
var done string
|
||||
err = tx.QueryRow(`SELECT value FROM settings WHERE key=$1`, flag).Scan(&done)
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
if done == "true" {
|
||||
return nil
|
||||
}
|
||||
var id int64
|
||||
err = tx.QueryRow(`INSERT INTO agents(key,name,description,role,builtin,enabled)
|
||||
VALUES ($1,'漏洞复测','从漏洞详情手动启动，读取原证据并保存独立复测结论。','assistant',false,true)
|
||||
ON CONFLICT (key) DO NOTHING RETURNING id`, db.FindingRetestAgentKey).Scan(&id)
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
if id > 0 {
|
||||
var pid int64
|
||||
if err = tx.QueryRow(`INSERT INTO agent_prompts(agent_id,version,template_text,note,updated_by)
|
||||
VALUES ($1,1,$2,'内置默认','system') RETURNING id`, id, agent.RetesterDefaultPrompt).Scan(&pid); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = tx.Exec(`UPDATE agents SET current_prompt_id=$1 WHERE id=$2`, pid, id); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if _, err = tx.Exec(`INSERT INTO settings(key,value) VALUES ($1,'true') ON CONFLICT(key) DO UPDATE SET value='true'`, flag); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (s *Server) finishRetest(id int64, status, reason string) {
|
||||
if err := s.m.pg.FinishFindingRetest(id, status, reason); err != nil {
|
||||
// Surface failure to the conversation runner rather than inventing a result.
|
||||
log.Printf("[retest %d] finish: %v", id, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"iter"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
type retestProvider struct {
|
||||
complete func(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error)
|
||||
}
|
||||
|
||||
func (p retestProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
return p.complete(ctx, req)
|
||||
}
|
||||
func (p retestProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
yield(llm.StreamEvent{}, errors.New("test expects non-streaming"))
|
||||
}
|
||||
}
|
||||
|
||||
func newRetestServer(t *testing.T) (*Server, int64) {
|
||||
t.Helper()
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skip("test postgres not configured")
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
td := t.TempDir()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := &Server{ctx: ctx, m: &Manager{pg: pg, dir: td, tasks: map[string]*Task{}},
|
||||
chatBusy: map[string]bool{}, chatCancel: map[string]context.CancelCauseFunc{}}
|
||||
if err := s.seedFindingRetester(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fid, err := pg.AddFinding(0, 0, "retest-server", "复测测试", "high", "summary", "original proof", "worker", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
oldAugment, oldResolve, oldPrompt := agent.ToolAugment, agent.ToolResolve, agent.PromptOverride
|
||||
agent.ToolAugment = func(context.Context, string) ([]actool.CoreTool, agent.DeferredInfo, func()) {
|
||||
return s.findingRetestTools(), agent.DeferredInfo{}, func() {}
|
||||
}
|
||||
agent.ToolResolve = nil
|
||||
agent.PromptOverride = func(string) (string, bool) { return agent.RetesterDefaultPrompt, true }
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
waitRetestIdle(t, s)
|
||||
agent.ToolAugment, agent.ToolResolve, agent.PromptOverride = oldAugment, oldResolve, oldPrompt
|
||||
_, _ = pg.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid)
|
||||
_, _ = pg.DeleteFinding(fid)
|
||||
pg.Close()
|
||||
})
|
||||
return s, fid
|
||||
}
|
||||
|
||||
func setRetestProvider(s *Server, p retestProvider) {
|
||||
s.chatAgent = agent.NewChatAgent(p, "test", s.m.dir, nil, 10000)
|
||||
s.chatAgent.SetNonStreaming(func() bool { return true })
|
||||
}
|
||||
|
||||
func retestRequest(handler http.HandlerFunc, method string, id int64, body string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest(method, "/", strings.NewReader(body))
|
||||
r.SetPathValue("id", strconv.FormatInt(id, 10))
|
||||
w := httptest.NewRecorder()
|
||||
handler(w, r)
|
||||
return w
|
||||
}
|
||||
|
||||
func waitRetestIdle(t *testing.T, s *Server) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
s.chatMu.Lock()
|
||||
n := len(s.chatBusy)
|
||||
s.chatMu.Unlock()
|
||||
if n == 0 {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Error("conversation did not become idle")
|
||||
}
|
||||
|
||||
func TestRetestActiveStatusLifecycle(t *testing.T) {
|
||||
s, fid := newRetestServer(t)
|
||||
check := func(want string, conversationID int64) {
|
||||
t.Helper()
|
||||
w := retestRequest(s.listActiveFindingRetests, http.MethodGet, 0, "")
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Retests []db.ActiveFindingRetest `json:"retests"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var matches []db.ActiveFindingRetest
|
||||
for _, item := range body.Retests {
|
||||
if item.FindingID == fid {
|
||||
matches = append(matches, item)
|
||||
}
|
||||
}
|
||||
if want == "" {
|
||||
if len(matches) != 0 {
|
||||
t.Fatalf("terminal retest still active: %+v", matches)
|
||||
}
|
||||
} else if len(matches) != 1 || matches[0].Status != want || matches[0].ConversationID != conversationID {
|
||||
t.Fatalf("active=%+v want status=%s conversation=%d", matches, want, conversationID)
|
||||
}
|
||||
for _, field := range []string{`"snapshot"`, `"evidence"`, `"notes"`, `"summary"`} {
|
||||
if strings.Contains(w.Body.String(), field) {
|
||||
t.Fatalf("status response leaks %s", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
check("", 0)
|
||||
for _, terminal := range []string{"completed", "failed", "stopped"} {
|
||||
r, c, _, err := s.m.pg.CreateFindingRetest(t.Context(), fid, "private notes")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
check("pending", c.ID)
|
||||
if _, err := s.m.pg.StartFindingRetest(t.Context(), r.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
check("running", c.ID)
|
||||
if terminal == "completed" {
|
||||
if err := s.m.pg.RecordFindingRetestResult(t.Context(), c.ID, "fixed", "summary", "proof"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
check("running", c.ID) // A staged verdict is still running until the turn ends.
|
||||
}
|
||||
if err := s.m.pg.FinishFindingRetest(r.ID, terminal, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
check("", 0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetestHTTPThroughConversationAndTools(t *testing.T) {
|
||||
s, fid := newRetestServer(t)
|
||||
entered, release := make(chan struct{}), make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
setRetestProvider(s, retestProvider{complete: func(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
n := calls.Add(1)
|
||||
if n == 1 {
|
||||
close(entered)
|
||||
select {
|
||||
case <-release:
|
||||
case <-ctx.Done():
|
||||
return llm.Message{}, "", llm.Usage{}, ctx.Err()
|
||||
}
|
||||
}
|
||||
name, input := "get_finding_retest_context", `{}`
|
||||
if n == 2 {
|
||||
serialized, _ := json.Marshal(req.Messages)
|
||||
if !bytes.Contains(serialized, []byte("original proof")) {
|
||||
t.Error("agent did not receive source evidence")
|
||||
}
|
||||
name, input = "record_finding_retest_result", `{"verdict":"inconclusive","summary":"缺少测试登录态","evidence":"已检查原证据；当前缺少有效登录态，无法确认修复状态。"}`
|
||||
}
|
||||
if n > 2 {
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("复测结论已保存")}}, "end_turn", llm.Usage{}, nil
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{{Type: llm.BlockToolUse, ID: strconv.Itoa(int(n)), Name: name, Input: json.RawMessage(input)}}}, "tool_use", llm.Usage{}, nil
|
||||
}})
|
||||
w := retestRequest(s.startFindingRetest, "POST", fid, `{"notes":"仅验证原接口"}`)
|
||||
if w.Code != 202 {
|
||||
t.Fatalf("create: %d %s", w.Code, w.Body)
|
||||
}
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("provider not called")
|
||||
}
|
||||
dup := retestRequest(s.startFindingRetest, "POST", fid, `{}`)
|
||||
if dup.Code != 200 || !strings.Contains(dup.Body.String(), `"created":false`) {
|
||||
t.Fatalf("duplicate %d %s", dup.Code, dup.Body)
|
||||
}
|
||||
close(release)
|
||||
waitRetestIdle(t, s)
|
||||
rows, err := s.m.pg.ListFindingRetests(fid)
|
||||
if err != nil || len(rows) != 1 || rows[0].Status != "completed" || rows[0].Verdict != "inconclusive" {
|
||||
t.Fatalf("history=%+v err=%v", rows, err)
|
||||
}
|
||||
if calls.Load() != 3 {
|
||||
t.Fatal("unexpected LLM calls", calls.Load())
|
||||
}
|
||||
f, _ := s.m.pg.GetFinding(fid)
|
||||
if f.Status != db.FindingPending || f.Evidence != "original proof" {
|
||||
t.Fatal("source modified")
|
||||
}
|
||||
listing := retestRequest(s.listFindingRetests, "GET", fid, "")
|
||||
if listing.Code != 200 || strings.Contains(listing.Body.String(), `"snapshot"`) {
|
||||
t.Fatalf("history %d %s", listing.Code, listing.Body)
|
||||
}
|
||||
res, err := s.findingRetestTools()[1].Call(intercept.WithConvID(t.Context(), *rows[0].ConversationID), json.RawMessage(`{"verdict":"fixed","summary":"bad overwrite","evidence":"bad"}`), nil)
|
||||
if err != nil || !res.IsError {
|
||||
t.Fatal("sealed result accepted", res, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetestStopsFailuresAndNoVerdict(t *testing.T) {
|
||||
for _, mode := range []string{"stop", "failure", "no_verdict"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
s, fid := newRetestServer(t)
|
||||
setRetestProvider(s, retestProvider{complete: func(ctx context.Context, _ llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
if mode == "stop" {
|
||||
<-ctx.Done()
|
||||
return llm.Message{}, "", llm.Usage{}, ctx.Err()
|
||||
}
|
||||
if mode == "failure" {
|
||||
return llm.Message{}, "", llm.Usage{}, errors.New("test model failed")
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("未保存结论")}}, "end_turn", llm.Usage{}, nil
|
||||
}})
|
||||
w := retestRequest(s.startFindingRetest, "POST", fid, `{}`)
|
||||
if w.Code != 202 {
|
||||
t.Fatalf("create %d %s", w.Code, w.Body)
|
||||
}
|
||||
if mode == "stop" {
|
||||
var body struct {
|
||||
Retest db.FindingRetest `json:"retest"`
|
||||
}
|
||||
_ = json.Unmarshal(w.Body.Bytes(), &body)
|
||||
stop := retestRequest(s.pgStopConversation, "POST", *body.Retest.ConversationID, `{}`)
|
||||
if stop.Code != 200 || !strings.Contains(stop.Body.String(), "stopping") {
|
||||
t.Fatalf("immediate stop: %d %s", stop.Code, stop.Body)
|
||||
}
|
||||
}
|
||||
waitRetestIdle(t, s)
|
||||
rows, _ := s.m.pg.ListFindingRetests(fid)
|
||||
want := "failed"
|
||||
if mode == "stop" {
|
||||
want = "stopped"
|
||||
}
|
||||
if len(rows) != 1 || rows[0].Status != want || rows[0].Error == "" {
|
||||
t.Fatalf("expected %s: %+v", want, rows)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetestFixedUpdatesFindingThroughConversation(t *testing.T) {
|
||||
s, fid := newRetestServer(t)
|
||||
var calls atomic.Int32
|
||||
setRetestProvider(s, retestProvider{complete: func(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
if calls.Add(1) == 1 {
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{{Type: llm.BlockToolUse, ID: "fixed-result", Name: "record_finding_retest_result", Input: json.RawMessage(`{"verdict":"fixed","summary":"修复生效","evidence":"原触发条件失效，正常对照仍可用。"}`)}}}, "tool_use", llm.Usage{}, nil
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("结论已保存")}}, "end_turn", llm.Usage{}, nil
|
||||
}})
|
||||
w := retestRequest(s.startFindingRetest, "POST", fid, `{}`)
|
||||
if w.Code != http.StatusAccepted {
|
||||
t.Fatalf("create: %d %s", w.Code, w.Body)
|
||||
}
|
||||
waitRetestIdle(t, s)
|
||||
f, err := s.m.pg.GetFinding(fid)
|
||||
if err != nil || f.Status != db.FindingFixed || f.Evidence != "original proof" {
|
||||
t.Fatalf("fixed finding=%+v err=%v", f, err)
|
||||
}
|
||||
// The new manual status is also accepted by the existing update API.
|
||||
w = retestRequest(s.patchFinding, "PATCH", fid, `{"status":"fixed"}`)
|
||||
if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"status":"fixed"`) {
|
||||
t.Fatalf("manual fixed: %d %s", w.Code, w.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetestValidationSeedingAndScope(t *testing.T) {
|
||||
s, fid := newRetestServer(t)
|
||||
for _, tc := range []struct {
|
||||
id int64
|
||||
body string
|
||||
code int
|
||||
}{{0, `{}`, 400}, {fid, `{`, 400}, {fid, `{"notes":"` + strings.Repeat("文", 4001) + `"}`, 400}, {999999999, `{}`, 404}, {fid, `{}`, 503}} {
|
||||
w := retestRequest(s.startFindingRetest, "POST", tc.id, tc.body)
|
||||
if w.Code != tc.code {
|
||||
t.Fatalf("code=%d want=%d %s", w.Code, tc.code, w.Body)
|
||||
}
|
||||
}
|
||||
rows, _ := s.m.pg.ListFindingRetests(fid)
|
||||
if len(rows) != 0 {
|
||||
t.Fatal("invalid request created state")
|
||||
}
|
||||
for _, tool := range s.findingRetestTools() {
|
||||
res, err := tool.Call(t.Context(), json.RawMessage(`{"verdict":"fixed","summary":"x","evidence":"y"}`), nil)
|
||||
if err != nil || !res.IsError {
|
||||
t.Fatal("unassociated conversation accepted", tool.Name(), res, err)
|
||||
}
|
||||
}
|
||||
a, _ := s.m.pg.GetAgentByKey(db.FindingRetestAgentKey)
|
||||
if a == nil || a.Builtin {
|
||||
t.Fatal("missing editable retester")
|
||||
}
|
||||
var triggers int
|
||||
if err := s.m.pg.QueryRow(`SELECT count(*) FROM agent_triggers WHERE agent_key=$1`, db.FindingRetestAgentKey).Scan(&triggers); err != nil || triggers != 0 {
|
||||
t.Fatal("unexpected automatic trigger", triggers, err)
|
||||
}
|
||||
if _, err := s.m.pg.SavePrompt(a.ID, "customized retest prompt", "test", "test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.seedFindingRetester(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var prompt string
|
||||
if err := s.m.pg.QueryRow(`SELECT p.template_text FROM agent_prompts p JOIN agents a ON a.current_prompt_id=p.id WHERE a.key=$1`, a.Key).Scan(&prompt); err != nil || prompt != "customized retest prompt" {
|
||||
t.Fatal("seed overwrote prompt", prompt, err)
|
||||
}
|
||||
_, _ = s.m.pg.SavePrompt(a.ID, agent.RetesterDefaultPrompt, "restore", "test")
|
||||
}
|
||||
@@ -0,0 +1,414 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/evidence"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func (s *Server) evidenceStore() *evidence.Store {
|
||||
return evidence.New(s.m.pg, s.m.traffic, filepath.Join(s.m.dir, "evidence"))
|
||||
}
|
||||
|
||||
// Add only new optional properties; preserve edited descriptions, existing
|
||||
// properties, agent bindings and disabled flags. The one-time flag also keeps
|
||||
// subsequent user unbinding of the evidence reader intact.
|
||||
func (s *Server) seedFindingTrafficTools() {
|
||||
const flag = "finding_traffic_tools_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{"report_finding", "update_finding_report"} {
|
||||
var schema any
|
||||
if key == "update_finding_report" {
|
||||
schema = s.toolUpdateFindingReport().InputSchema()
|
||||
} else {
|
||||
for _, seed := range agent.BuiltinToolSeeds() {
|
||||
if seed.Key == key {
|
||||
schema = seed.Schema
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
raw, err := json.Marshal(schema)
|
||||
if err != nil {
|
||||
log.Printf("[evidence] tool schema: %v", err)
|
||||
return
|
||||
}
|
||||
var obj map[string]json.RawMessage
|
||||
if err = json.Unmarshal(raw, &obj); err != nil {
|
||||
return
|
||||
}
|
||||
// BuiltinToolSeeds stores Schema as RawMessage; both forms marshal as JSON.
|
||||
var properties map[string]json.RawMessage
|
||||
if err = json.Unmarshal(obj["properties"], &properties); err != nil {
|
||||
return
|
||||
}
|
||||
name := "traffic_refs"
|
||||
if key == "update_finding_report" {
|
||||
name = "evidence_version"
|
||||
}
|
||||
if len(properties[name]) == 0 {
|
||||
return
|
||||
}
|
||||
_, err = s.m.pg.Exec(`UPDATE tools SET schema=jsonb_set(schema,ARRAY['properties',$2::text],$3::jsonb,true),updated_at=now()
|
||||
WHERE key=$1 AND system AND NOT(COALESCE(schema->'properties','{}'::jsonb) ? $2)`, key, name, string(properties[name]))
|
||||
if err != nil {
|
||||
log.Printf("[evidence] upgrade tool %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if reporter, _ := s.m.pg.GetAgentByKey("reporter"); reporter != nil {
|
||||
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"get_finding_traffic"}); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
func (s *Server) registerFindingTraffic(mux *http.ServeMux) {
|
||||
base := "/api/exploration/findings/{id}/traffic"
|
||||
mux.HandleFunc("GET "+base, s.getFindingTraffic)
|
||||
mux.HandleFunc("POST "+base, s.bindFindingTraffic)
|
||||
mux.HandleFunc("PATCH "+base+"/{binding_id}", s.editFindingTraffic)
|
||||
mux.HandleFunc("DELETE "+base+"/{binding_id}", s.editFindingTraffic)
|
||||
mux.HandleFunc("PUT "+base+"/order", s.editFindingTraffic)
|
||||
mux.HandleFunc("GET "+base+"/{binding_id}", s.getFindingTrafficDetail)
|
||||
mux.HandleFunc("GET "+base+"/{binding_id}/body", s.getFindingTrafficBody)
|
||||
}
|
||||
|
||||
func evidenceError(w http.ResponseWriter, err error) {
|
||||
status := http.StatusUnprocessableEntity
|
||||
if errors.Is(err, db.ErrEvidenceConflict) || errors.Is(err, db.ErrTaskArchiveState) {
|
||||
status = http.StatusConflict
|
||||
}
|
||||
if errors.Is(err, db.ErrFindingNotFound) || errors.Is(err, db.ErrEvidenceNotFound) {
|
||||
status = http.StatusNotFound
|
||||
}
|
||||
writeErr(w, status, err.Error())
|
||||
}
|
||||
|
||||
func (s *Server) findingTrafficAccess(w http.ResponseWriter, r *http.Request, write bool) (int64, bool) {
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
writeErr(w, 400, "invalid finding id")
|
||||
return 0, false
|
||||
}
|
||||
f, err := s.m.pg.GetFinding(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return 0, false
|
||||
}
|
||||
if f == nil {
|
||||
writeErr(w, 404, "finding not found")
|
||||
return 0, false
|
||||
}
|
||||
if taskID := r.URL.Query().Get("context_task"); taskID != "" {
|
||||
task := s.m.ResolveTask(taskID)
|
||||
if task == nil {
|
||||
writeErr(w, 404, "context task not found")
|
||||
return 0, false
|
||||
}
|
||||
_, inherited, allowed := findingProvenanceInTask(task, f.TaskID)
|
||||
if !allowed {
|
||||
writeErr(w, 404, "finding not available in task context")
|
||||
return 0, false
|
||||
}
|
||||
if write && inherited {
|
||||
writeErr(w, 403, "继承漏洞只读，请在来源任务中修改")
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
func trafficSummary(in *db.FindingTraffic) *db.FindingTraffic {
|
||||
out := *in
|
||||
out.Bindings = append([]db.FindingTrafficBinding{}, in.Bindings...)
|
||||
for i := range out.Bindings {
|
||||
out.Bindings[i].Snapshot.ReqHead = ""
|
||||
out.Bindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
out, err := s.m.pg.GetFindingTraffic(r.Context(), id)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, trafficSummary(out))
|
||||
}
|
||||
|
||||
func (s *Server) bindFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, true)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Refs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil {
|
||||
writeErr(w, 400, "invalid body")
|
||||
return
|
||||
}
|
||||
if len(body.Refs) == 0 {
|
||||
writeErr(w, 400, "请选择流量")
|
||||
return
|
||||
}
|
||||
out, err := s.evidenceStore().Bind(r.Context(), id, body.Refs)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, trafficSummary(out))
|
||||
}
|
||||
|
||||
func (s *Server) editFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, true)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Version *int64 `json:"version"`
|
||||
Role *string `json:"role"`
|
||||
Note *string `json:"note"`
|
||||
Order []string `json:"binding_ids"`
|
||||
}
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil || body.Version == nil {
|
||||
writeErr(w, 400, "version 和有效请求体必填")
|
||||
return
|
||||
}
|
||||
var order []int64
|
||||
bindingID := int64(0)
|
||||
if r.Method == http.MethodPut {
|
||||
if body.Order == nil {
|
||||
writeErr(w, 400, "binding_ids 必填")
|
||||
return
|
||||
}
|
||||
order = []int64{}
|
||||
for _, raw := range body.Order {
|
||||
v, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || v <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
order = append(order, v)
|
||||
}
|
||||
} else {
|
||||
var err error
|
||||
bindingID, err = strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bindingID <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
}
|
||||
err := s.m.pg.EditFindingTraffic(r.Context(), id, bindingID, *body.Version, body.Role, body.Note, r.Method == http.MethodDelete, order)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
s.getFindingTraffic(w, r)
|
||||
}
|
||||
|
||||
type evidencePreview struct {
|
||||
Content string `json:"content"`
|
||||
Offset int64 `json:"offset"`
|
||||
Total int64 `json:"total"`
|
||||
NextOffset int64 `json:"next_offset"`
|
||||
Truncated bool `json:"truncated"`
|
||||
Binary bool `json:"binary"`
|
||||
}
|
||||
|
||||
func readEvidencePreview(store *evidence.Store, snapshot db.TrafficEvidenceSnapshot, side string, offset, length int64) (out evidencePreview, err error) {
|
||||
if offset < 0 || length < 0 {
|
||||
return out, errors.New("offset / length 不能为负数")
|
||||
}
|
||||
if length == 0 || length > 8192 {
|
||||
length = 8192
|
||||
}
|
||||
f, total, err := store.OpenBody(snapshot, side)
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
defer f.Close()
|
||||
if offset > total {
|
||||
return out, errors.New("offset 超出正文长度")
|
||||
}
|
||||
if _, err = f.Seek(offset, io.SeekStart); err != nil {
|
||||
return out, err
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(f, length))
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
// Leave an incomplete trailing UTF-8 rune for the next page. Explicit byte
|
||||
// offsets are still accepted; complete downloads always retain original bytes.
|
||||
if offset+int64(len(raw)) < total && len(raw) >= utf8.UTFMax {
|
||||
start := len(raw) - 1
|
||||
for start > 0 && raw[start]&0xc0 == 0x80 {
|
||||
start--
|
||||
}
|
||||
if !utf8.FullRune(raw[start:]) {
|
||||
raw = raw[:start]
|
||||
}
|
||||
}
|
||||
out = evidencePreview{Offset: offset, Total: total, NextOffset: offset + int64(len(raw)), Truncated: offset+int64(len(raw)) < total,
|
||||
Binary: bytes.IndexByte(raw, 0) >= 0 || (offset == 0 && !utf8.Valid(raw))}
|
||||
if out.Binary {
|
||||
out.Content = fmt.Sprintf("[二进制正文，%d 字节；请下载查看]", total)
|
||||
} else {
|
||||
out.Content = string(bytes.ToValidUTF8(raw, []byte("�")))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTrafficDetail(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bid <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
store := s.evidenceStore()
|
||||
var result any
|
||||
err = store.WithBinding(r.Context(), id, bid, func(b db.FindingTrafficBinding) error {
|
||||
req, err := readEvidencePreview(store, b.Snapshot, "request", 0, 8192)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := readEvidencePreview(store, b.Snapshot, "response", 0, 8192)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
result = map[string]any{"binding": b, "request": req, "response": resp}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, result)
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTrafficBody(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bid <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
side := r.URL.Query().Get("side")
|
||||
store := s.evidenceStore()
|
||||
// Resolve the binding under the evidence lock, then read the blob without it:
|
||||
// both paths below are O(body size) and would otherwise stall every evidence
|
||||
// write for as long as the client takes to receive the data.
|
||||
b, err := store.Binding(r.Context(), id, bid)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("download") == "1" {
|
||||
f, length, err := store.OpenBody(b.Snapshot, side)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
defer f.Close()
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"evidence-%d-%s.bin\"", bid, side))
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(length, 10))
|
||||
_, _ = io.Copy(w, f)
|
||||
return
|
||||
}
|
||||
offset, length := int64(0), int64(8192)
|
||||
for name, dst := range map[string]*int64{"offset": &offset, "length": &length} {
|
||||
if raw := r.URL.Query().Get(name); raw != "" {
|
||||
v, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
*dst = v
|
||||
}
|
||||
}
|
||||
preview, err := readEvidencePreview(store, b.Snapshot, side, offset, length)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, preview)
|
||||
}
|
||||
|
||||
func (s *Server) toolGetFindingTraffic() actool.CoreTool {
|
||||
return roTool("get_finding_traffic", "读取漏洞已绑定的真实流量证据，不依赖捕获开关。finding_id 使用 report_finding JSON 返回的独立漏洞记录 ID（不是第一行的探索节点 ID）。先不传 binding_id 获取清单及 version；空清单是正常情况，TCP 等非 HTTP 漏洞或未采集时仍可依据文字/命令证据编写报告，不强制绑定。有绑定时按 binding_id、side(request/response)、offset 分段读取正文。写报告时将读取的 version 作为 evidence_version 传给 update_finding_report，后者 finding_id 仍使用探索节点 ID。",
|
||||
objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID"), "binding_id": strParam("清单里的绑定 ID，省略则返回清单"), "side": strParam("request 或 response，默认 response"), "offset": map[string]any{"type": "integer"}, "length": map[string]any{"type": "integer"}}, "finding_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
BindingID json.RawMessage `json:"binding_id"`
|
||||
Side string `json:"side"`
|
||||
Offset, Length int64
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
id, bid := parseProfileID(a.FindingID), parseProfileID(a.BindingID)
|
||||
if err := s.agentFindingTrafficAccess(ctx, id, false); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if len(a.BindingID) == 0 {
|
||||
list, err := s.m.pg.GetFindingTraffic(ctx, id)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(trafficSummary(list))
|
||||
}
|
||||
if a.Side == "" {
|
||||
a.Side = "response"
|
||||
}
|
||||
store := s.evidenceStore()
|
||||
var result any
|
||||
err := store.WithBinding(ctx, id, bid, func(b db.FindingTrafficBinding) error {
|
||||
preview, err := readEvidencePreview(store, b.Snapshot, a.Side, a.Offset, a.Length)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
head := b.Snapshot.ReqHead
|
||||
if a.Side == "response" {
|
||||
head = b.Snapshot.RespHead
|
||||
}
|
||||
result = map[string]any{"binding_id": strconv.FormatInt(bid, 10), "head": head, "body": preview}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(result)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/evidence"
|
||||
"github.com/Autumn-27/artex/report"
|
||||
)
|
||||
|
||||
func writeFindingsEvidenceZip(out io.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
zw := zip.NewWriter(out)
|
||||
if err := writeFindingsZipEntries(zw, findings, stage, now); err != nil {
|
||||
zw.Close()
|
||||
return err
|
||||
}
|
||||
return zw.Close()
|
||||
}
|
||||
|
||||
func writeFindingsZipEntries(zw *zip.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
store := evidence.New(nil, nil, stage) // private export copy; not subject to GC
|
||||
used := map[string]int{}
|
||||
for _, f := range findings {
|
||||
base := report.FindingFilename(f)
|
||||
name := base
|
||||
if n := used[base]; n > 0 {
|
||||
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
|
||||
}
|
||||
used[base]++
|
||||
w, err := zw.Create(name)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = io.WriteString(w, report.SingleFindingMarkdown(f, now)); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range f.TrafficBindings {
|
||||
prefix := fmt.Sprintf("evidence/%d/%d/", f.ID, b.ID)
|
||||
w, err := zw.Create(prefix + "manifest.json")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = json.NewEncoder(w).Encode(b); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, side := range []string{"request", "response"} {
|
||||
if err := func() error {
|
||||
body, _, err := store.OpenBody(b.Snapshot, side)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer body.Close()
|
||||
w, err := zw.Create(prefix + side + ".http")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
head := b.Snapshot.ReqHead
|
||||
if side == "response" {
|
||||
head = b.Snapshot.RespHead
|
||||
}
|
||||
if _, err = io.WriteString(w, strings.TrimRight(head, "\r\n")+"\n\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = io.Copy(w, body); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = body.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
w, err = zw.Create(prefix + side + ".bin")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = io.Copy(w, body)
|
||||
return err
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildFindingsEvidenceZip(path string, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
writeErr := writeFindingsEvidenceZip(f, findings, stage, now)
|
||||
syncErr := f.Sync()
|
||||
closeErr := f.Close()
|
||||
if err := errors.Join(writeErr, syncErr, closeErr); err != nil {
|
||||
return err
|
||||
}
|
||||
// Re-read every compressed entry to validate the finished ZIP, including CRCs,
|
||||
// before the caller starts an HTTP download response.
|
||||
archive, err := zip.OpenReader(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer archive.Close()
|
||||
for _, entry := range archive.File {
|
||||
r, err := entry.Open()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, readErr := io.Copy(io.Discard, r)
|
||||
if err := errors.Join(readErr, r.Close()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func trafficEvidenceServer(t *testing.T) (*Server, *db.RecordedFinding, func(string, string, string) *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { m.Close() })
|
||||
m.traffic, err = traffic.Open(filepath.Join(m.dir, "traffic"), ":0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := m.CreateTask("traffic evidence API", "local evidence", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tid, _ := strconv.ParseInt(task.ID, 10, 64)
|
||||
t.Cleanup(func() { m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, tid); m.pg.DeleteTask(tid) })
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
s.archiveWG.Wait()
|
||||
// The archive test below replaces the cancelled service context. Wait for
|
||||
// the side-question snapshot writer too before reusing this fixture.
|
||||
<-s.side.done
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := func(method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
r := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+token)
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
input := db.RecordFindingInput{TaskID: tid, ExplorationID: task.ExpID, Worker: "test", VulnClass: "TEST", Summary: "bound evidence", Name: "Evidence API", Severity: "low"}
|
||||
f, err := s.evidenceStore().Record(context.Background(), input, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return s, f, request
|
||||
}
|
||||
|
||||
func seedServerEvidenceFlow(t *testing.T, s *Server, id string, body []byte) {
|
||||
t.Helper()
|
||||
_, err := s.m.traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,1,'evidence.local','POST','/','http://evidence.local/',200,'application/octet-stream',0,?,'')`, id, len(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.m.traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,resp_head,req_body,resp_body) VALUES(?,'POST / HTTP/1.1','HTTP/1.1 200 OK',?,?)`, id, []byte{}, body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficToolUpgradePreservesCustomization(t *testing.T) {
|
||||
s, _, _ := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
worker, err := pg.GetAgentByKey("worker")
|
||||
if err != nil || worker == nil {
|
||||
t.Fatal("missing worker", err)
|
||||
}
|
||||
oldPrompt, _ := pg.CurrentPrompt(worker.ID)
|
||||
t.Cleanup(func() { pg.SavePrompt(worker.ID, oldPrompt, "restore test fixture", "test") })
|
||||
if _, err := pg.SavePrompt(worker.ID, "USER CUSTOM PROMPT", "test", "test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
custom := json.RawMessage(`{"type":"object","properties":{"evidence":{"type":"string","description":"user evidence instructions"}},"required":["evidence"]}`)
|
||||
for _, key := range []string{"report_finding", "update_finding_report", "get_finding_traffic"} {
|
||||
old, err := pg.GetTool(key)
|
||||
if err != nil || old == nil {
|
||||
t.Fatal("missing tool", key, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
bindings, _ := json.Marshal(old.Agents)
|
||||
pg.UpdateTool(old.Key, old.Description, old.Schema, bindings, old.Enabled)
|
||||
})
|
||||
if err := pg.UpdateTool(key, "user description", custom, json.RawMessage(`["worker"]`), false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := pg.SetSetting("finding_traffic_tools_v1", "false"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingTrafficTools()
|
||||
for key, property := range map[string]string{"report_finding": "traffic_refs", "update_finding_report": "evidence_version"} {
|
||||
tool, err := pg.GetTool(key)
|
||||
if err != nil || tool == nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var schema struct {
|
||||
Properties map[string]json.RawMessage `json:"properties"`
|
||||
Required []string `json:"required"`
|
||||
}
|
||||
if err := json.Unmarshal(tool.Schema, &schema); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.Enabled || tool.Description != "user description" || len(tool.Agents) != 1 || tool.Agents[0] != "worker" || len(schema.Required) != 1 || schema.Required[0] != "evidence" || len(schema.Properties[property]) == 0 || !strings.Contains(string(schema.Properties["evidence"]), "user evidence instructions") {
|
||||
t.Fatalf("custom configuration overwritten: %+v", tool)
|
||||
}
|
||||
}
|
||||
reader, _ := pg.GetTool("get_finding_traffic")
|
||||
if reader.Enabled || !contains(reader.Agents, "reporter") {
|
||||
t.Fatal("reader should be bound without enabling it")
|
||||
}
|
||||
if err := pg.RemoveAgentFromTool("reporter", "get_finding_traffic"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingTrafficTools()
|
||||
reader, _ = pg.GetTool("get_finding_traffic")
|
||||
if contains(reader.Agents, "reporter") {
|
||||
t.Fatal("restart overwrote user's unbinding")
|
||||
}
|
||||
if got, _ := pg.CurrentPrompt(worker.ID); got != "USER CUSTOM PROMPT" {
|
||||
t.Fatal("custom prompt replaced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficAPIAndExport(t *testing.T) {
|
||||
s, f, req := trafficEvidenceServer(t)
|
||||
body := bytes.Repeat([]byte{0, 255, 65, 66}, 5000)
|
||||
for _, id := range []string{"a", "b", "c"} {
|
||||
seedServerEvidenceFlow(t, s, id, body)
|
||||
}
|
||||
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
|
||||
w := req("POST", base, `{"traffic_refs":[{"traffic_id":"a","role":"baseline"},{"traffic_id":"b","role":"proof","note":"proof note"},{"traffic_id":"c"}]}`)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("bind: %d %s", w.Code, w.Body)
|
||||
}
|
||||
var list db.FindingTraffic
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list.Bindings) != 3 || list.Bindings[0].Snapshot.ReqHead != "" {
|
||||
t.Fatal("not a bounded summary")
|
||||
}
|
||||
first, second, third := list.Bindings[0].ID, list.Bindings[1].ID, list.Bindings[2].ID
|
||||
if w = req("PATCH", fmt.Sprintf("%s/%d", base, first), `{"version":0,"note":"stale"}`); w.Code != 409 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":1,"binding_ids":["%d","%d","%d"]}`, third, first, second)); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":2,"binding_ids":["%d"]}`, first)); w.Code != 409 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d", base, first), ""); w.Code != 200 || !strings.Contains(w.Body.String(), `"binary":true`) {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&offset=20001", base, first), ""); w.Code != 422 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if _, err := s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&download=1", base, first), ""); w.Code != 200 || !bytes.Equal(w.Body.Bytes(), body) {
|
||||
t.Fatal("download after original removal failed")
|
||||
}
|
||||
// IDs belong to the requested finding; cross-finding detail cannot be read.
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := s.m.pg.GetTask(*finding.TaskID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
other, err := s.evidenceStore().Record(t.Context(), db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "other", Severity: "low"}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("/api/exploration/findings/%d/traffic/%d", other.FindingID, first), ""); w.Code != 404 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
exportURL := fmt.Sprintf("/api/exploration/findings/export?format=md-zip&scope=selected&ids=%d", f.FindingID)
|
||||
w = req("GET", exportURL, "")
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
archive, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
binaries := 0
|
||||
linked := false
|
||||
for _, entry := range archive.File {
|
||||
r, err := entry.Open()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := io.ReadAll(r)
|
||||
r.Close()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.HasSuffix(entry.Name, "/response.bin") {
|
||||
binaries++
|
||||
if !bytes.Equal(raw, body) {
|
||||
t.Fatal("export changed binary")
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(entry.Name, ".md") && bytes.Contains(raw, []byte(fmt.Sprintf("evidence/%d/%d/", f.FindingID, third))) {
|
||||
linked = true
|
||||
}
|
||||
}
|
||||
if binaries != 3 || !linked {
|
||||
t.Fatalf("missing attachments/links: %d %v", binaries, linked)
|
||||
}
|
||||
// Reports can use snapshots with capture stopped and original rows gone.
|
||||
result, err := s.toolGetFindingTraffic().Call(context.Background(), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
|
||||
if err != nil || !strings.Contains(result.Flatten(), "proof note") {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
// Corrupt body must fail before an attachment response is sent.
|
||||
snap := list.Bindings[0].Snapshot
|
||||
path := filepath.Join(s.m.dir, "evidence", "blobs", snap.RespHash[:2], snap.RespHash+".bin")
|
||||
if err = os.WriteFile(path, []byte("corrupt"), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", exportURL, ""); w.Code < 400 || strings.Contains(w.Header().Get("Content-Type"), "zip") {
|
||||
t.Fatal("corrupt export returned an attachment", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficArchiveV3RoundTripAndRetry(t *testing.T) {
|
||||
s, f, _ := trafficEvidenceServer(t)
|
||||
ctx := context.Background()
|
||||
s.ctx = ctx
|
||||
body := []byte("evidence without original traffic\x00\xff")
|
||||
seedServerEvidenceFlow(t, s, "original", body)
|
||||
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "original", Role: "proof", Note: "keep note"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tid := *finding.TaskID
|
||||
sharedTask, err := s.m.CreateTask("shared evidence", "fixture", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sid, _ := strconv.ParseInt(sharedTask.ID, 10, 64)
|
||||
defer s.m.pg.DeleteTask(sid)
|
||||
shared, err := s.evidenceStore().Record(ctx, db.RecordFindingInput{TaskID: sid, ExplorationID: sharedTask.ExpID, Summary: "shared", Severity: "low"}, []db.TrafficRef{{TrafficID: "original"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.SetTaskPaused(strconv.FormatInt(tid, 10), true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
job, err := s.m.pg.QueueTaskArchive(tid)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.pg.EditFindingTraffic(ctx, f.FindingID, list.Bindings[0].ID, 1, nil, nil, true, nil); err == nil {
|
||||
t.Fatal("write allowed after archive queued")
|
||||
}
|
||||
if err = s.runOneTaskArchiveJob(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
|
||||
t.Fatal("archive retained hot finding")
|
||||
}
|
||||
if err = s.evidenceStore().Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.evidenceStore().WithBinding(ctx, shared.FindingID, shared.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
|
||||
if r != nil {
|
||||
r.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal("shared body deleted", err)
|
||||
}
|
||||
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Force a database error after validated body installation; then retry normally.
|
||||
if _, err = s.m.pg.Exec(`CREATE FUNCTION fail_evidence_restore_test() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'restore fixture failure'; END $$`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.pg.Exec(`CREATE TRIGGER fail_evidence_restore_test BEFORE INSERT ON finding_traffic_bindings FOR EACH ROW EXECUTE FUNCTION fail_evidence_restore_test()`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.m.pg.Exec(`DROP FUNCTION IF EXISTS fail_evidence_restore_test() CASCADE`)
|
||||
if err = s.runOneTaskArchiveJob(); err == nil {
|
||||
t.Fatal("expected restore failure")
|
||||
}
|
||||
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
|
||||
t.Fatal("partial restore")
|
||||
}
|
||||
if _, err = s.m.pg.Exec(`DROP FUNCTION fail_evidence_restore_test() CASCADE`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.runOneTaskArchiveJob(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
restored, err := s.m.pg.GetFindingTraffic(ctx, f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := restored.Bindings[0]
|
||||
if restored.Version != list.Version || b.ID != list.Bindings[0].ID || b.Note != "keep note" || b.Role != "proof" || b.Position != 0 {
|
||||
t.Fatalf("restored: %+v", restored)
|
||||
}
|
||||
if err = s.evidenceStore().WithBinding(ctx, f.FindingID, b.ID, func(b db.FindingTrafficBinding) error {
|
||||
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer r.Close()
|
||||
got, err := io.ReadAll(r)
|
||||
if !bytes.Equal(got, body) {
|
||||
t.Fatal("restore changed body")
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The read tool enforces task visibility too.
|
||||
result, err := s.toolGetFindingTraffic().Call(agent.WithRunInfo(ctx, agent.RunInfo{TaskID: sid}), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
|
||||
if err != nil || !strings.Contains(result.Flatten(), "不可读取") {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficFailedReportDoesNotTrigger(t *testing.T) {
|
||||
s, f, _ := trafficEvidenceServer(t)
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := s.m.pg.GetTask(*finding.TaskID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var last int64
|
||||
if err = s.m.pg.QueryRow(`SELECT COALESCE(max(id),0) FROM activity`).Scan(&last); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.pg.SetSchedState(schedKeyLastToolCall, strconv.FormatInt(last, 10)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, err := s.m.pg.Exploration(task.ExplorationID).AppendActivity(db.Activity{Kind: "tool_result", Tool: "report_finding", ToolUseID: "failure-fixture", IsError: true, Detail: "body missing"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// A nil server makes any attempted trigger fail this test immediately.
|
||||
scheduler := &Scheduler{pg: s.m.pg}
|
||||
scheduler.fireToolCalls([]*db.AgentTrigger{{OnToolCall: true, ToolNames: []string{"report_finding"}, AgentKey: "reporter"}})
|
||||
if got := scheduler.mustState(schedKeyLastToolCall); got != strconv.FormatInt(id, 10) {
|
||||
t.Fatal("failed finding must still advance watermark", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficUTF8SegmentsAndInheritedWrites(t *testing.T) {
|
||||
s, f, req := trafficEvidenceServer(t)
|
||||
ctx := context.Background()
|
||||
text := strings.Repeat("中文证据🙂", 2000)
|
||||
seedServerEvidenceFlow(t, s, "unicode", []byte(text))
|
||||
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "unicode"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var joined strings.Builder
|
||||
offset := int64(0)
|
||||
for {
|
||||
var preview evidencePreview
|
||||
err = s.evidenceStore().WithBinding(ctx, f.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
var e error
|
||||
preview, e = readEvidencePreview(s.evidenceStore(), b.Snapshot, "response", offset, 8192)
|
||||
return e
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if preview.Binary {
|
||||
t.Fatal("valid UTF-8 classified binary")
|
||||
}
|
||||
joined.WriteString(preview.Content)
|
||||
if !preview.Truncated {
|
||||
break
|
||||
}
|
||||
if preview.NextOffset <= offset {
|
||||
t.Fatal("preview made no progress")
|
||||
}
|
||||
offset = preview.NextOffset
|
||||
}
|
||||
if joined.String() != text {
|
||||
t.Fatal("UTF-8 character split damaged preview")
|
||||
}
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
child, err := s.m.CreateTaskWithOptions("inherited evidence API", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*finding.TaskID}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cid, _ := strconv.ParseInt(child.ID, 10, 64)
|
||||
defer s.m.pg.DeleteTask(cid)
|
||||
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
|
||||
suffix := "?context_task=" + child.ID
|
||||
if w := req("GET", base+suffix, ""); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
for _, method := range []string{"POST", "PATCH", "PUT", "DELETE"} {
|
||||
path := base
|
||||
if method == "PATCH" || method == "DELETE" {
|
||||
path += fmt.Sprintf("/%d", list.Bindings[0].ID)
|
||||
}
|
||||
if method == "PUT" {
|
||||
path += "/order"
|
||||
}
|
||||
w := req(method, path+suffix, `{"version":1,"traffic_refs":[{"traffic_id":"unicode"}]}`)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(method, w.Code, w.Body)
|
||||
}
|
||||
}
|
||||
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, err := s.toolUpdateFindingReport().Call(ctx, json.RawMessage(fmt.Sprintf(`{"finding_id":%d,"evidence_version":1,"report":"## 证据报告\n\n证据 #%d：已验证完整请求响应"}`, f.NodeID, list.Bindings[0].ID)), nil)
|
||||
if err != nil || result.IsError {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
updated, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.ReportEvidenceVersion != 1 || updated.EvidenceVersion != 1 {
|
||||
t.Fatal("report did not record the read evidence version")
|
||||
}
|
||||
for _, format := range []string{"md-single", "json", "csv"} {
|
||||
w := req("GET", fmt.Sprintf("/api/exploration/findings/export?scope=selected&ids=%d&format=%s", f.FindingID, format), "")
|
||||
if w.Code != 200 {
|
||||
t.Fatal(format, w.Code, w.Body)
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), fmt.Sprint(list.Bindings[0].ID)) {
|
||||
t.Fatal("export omitted evidence ID", format)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func (s *Server) seedFindingWorkflowTools() {
|
||||
const flag = "finding_workflow_tools_v2_reporter"
|
||||
if value, _, _ := s.m.pg.GetSetting(flag); value == "true" {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
row, err := s.m.pg.GetTool(key)
|
||||
if err != nil {
|
||||
log.Printf("[evidence] load %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if row == nil || !row.System {
|
||||
continue
|
||||
}
|
||||
var schema map[string]any
|
||||
if err := json.Unmarshal(row.Schema, &schema); err != nil {
|
||||
log.Printf("[evidence] invalid schema for %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if schema == nil {
|
||||
log.Printf("[evidence] missing object schema for %s", key)
|
||||
return
|
||||
}
|
||||
props := objectProperty(schema, "properties")
|
||||
if key == "report_finding" {
|
||||
if _, exists := props["evidence_hint_id"]; !exists {
|
||||
props["evidence_hint_id"] = map[string]any{"type": "integer", "description": "可选：本任务中对应此漏洞的 hint ID；读取该提示保存的 traffic_refs 一并绑定，无提示时省略"}
|
||||
}
|
||||
} else {
|
||||
if _, exists := props["traffic_refs"]; !exists {
|
||||
props["traffic_refs"] = agent.HintTrafficSchema()
|
||||
}
|
||||
hints := objectProperty(props, "hints")
|
||||
if _, exists := hints["type"]; !exists {
|
||||
hints["type"] = "array"
|
||||
}
|
||||
items := objectProperty(hints, "items")
|
||||
if _, ok := items["type"]; !ok {
|
||||
items["type"] = "object"
|
||||
}
|
||||
itemProps := objectProperty(items, "properties")
|
||||
for name, value := range map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()} {
|
||||
if _, exists := itemProps[name]; !exists {
|
||||
itemProps[name] = value
|
||||
}
|
||||
}
|
||||
}
|
||||
raw, _ := json.Marshal(schema)
|
||||
result, err := s.m.pg.Exec(`UPDATE tools SET schema=$2::jsonb,updated_at=now() WHERE key=$1 AND system AND schema=$3::jsonb`, key, string(raw), string(row.Schema))
|
||||
if err != nil {
|
||||
log.Printf("[evidence] upgrade %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if n, _ := result.RowsAffected(); n != 1 {
|
||||
return
|
||||
} // preserve concurrent user edits
|
||||
}
|
||||
// Upgrade only the original default binding. Customized lists and enabled
|
||||
// flags survive; the one-time flag also preserves future user unbinding.
|
||||
readers := `["worker","reporter"]`
|
||||
for _, key := range []string{"traffic_search", "traffic_get", "traffic_blob"} {
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$2::jsonb WHERE key=$1 AND system AND (agents='["worker"]'::jsonb OR (agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5))`, key, readers); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$1::jsonb WHERE key='get_finding_traffic' AND system AND agents @> '["auto","reporter"]'::jsonb AND jsonb_array_length(agents)=2`, `["auto","reporter","worker","planner","mainagent","pentest"]`); err != nil {
|
||||
return
|
||||
}
|
||||
// Replace the previous code default only; preserve customized binding lists.
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents='["reporter"]'::jsonb WHERE key='bind_finding_traffic' AND system AND agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5`); err != nil {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"bind_finding_traffic"}); err != nil {
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
func objectProperty(parent map[string]any, key string) map[string]any {
|
||||
value, ok := parent[key].(map[string]any)
|
||||
if !ok {
|
||||
value = map[string]any{}
|
||||
parent[key] = value
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func (s *Server) agentFindingTrafficAccess(ctx context.Context, id int64, write bool) error {
|
||||
if id <= 0 {
|
||||
return errors.New("finding_id 必须为独立漏洞记录 ID；不是探索节点 ID")
|
||||
}
|
||||
f, err := s.m.pg.GetFinding(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if f == nil {
|
||||
return fmt.Errorf("%w：finding_id=%d。证据工具使用独立漏洞记录 ID，请从 list_task_findings / get_task_node_detail 的 finding_id 字段读取；不要传 id / finding_node_id", db.ErrFindingNotFound, id)
|
||||
}
|
||||
if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 {
|
||||
task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10))
|
||||
if task == nil {
|
||||
return errors.New("任务不存在")
|
||||
}
|
||||
_, inherited, allowed := findingProvenanceInTask(task, f.TaskID)
|
||||
if !allowed {
|
||||
return errors.New("当前任务不可读取该漏洞")
|
||||
}
|
||||
if write && inherited {
|
||||
return errors.New("继承漏洞的流量证据只读，请到来源任务修改")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) toolBindFindingTraffic() actool.CoreTool {
|
||||
return wrTool("bind_finding_traffic", "为已登记漏洞补绑经核实的真实 HTTP 流量。finding_id 使用独立漏洞记录 ID；不要传探索节点 ID。同批引用全部成功或全部失败，重复引用不覆盖已有说明。补绑会使已有报告标记待更新；不要为补包重新探测或重复创建漏洞。",
|
||||
objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID，从 list_task_findings / get_task_node_detail 的 finding_id 字段读取"), "traffic_refs": agent.HintTrafficSchema()}, "finding_id", "traffic_refs"),
|
||||
func(ctx context.Context, raw json.RawMessage) (actool.Result, error) {
|
||||
if !s.m.pg.GetBool(settingAgentTrafficBinding, false) {
|
||||
return actool.Errorf("Agent 自动绑定流量已关闭；请在系统设置开启，或使用页面人工绑定。"), nil
|
||||
}
|
||||
var args struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Refs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &args); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
id := parseProfileID(args.FindingID)
|
||||
if err := s.agentFindingTrafficAccess(ctx, id, true); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if len(args.Refs) == 0 {
|
||||
return actool.Errorf("补绑需要至少一条已核实的 traffic_refs；无流量无需调用此工具"), nil
|
||||
}
|
||||
list, err := s.evidenceStore().Bind(ctx, id, args.Refs)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(trafficSummary(list))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func workflowCall(t *testing.T, ctx context.Context, tool actool.CoreTool, input any, wantError bool) string {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r, err := tool.Call(ctx, raw, nil)
|
||||
if err != nil || r.IsError != wantError {
|
||||
t.Fatalf("%s: error=%v result=%s", tool.Name(), err, r.Flatten())
|
||||
}
|
||||
return r.Flatten()
|
||||
}
|
||||
|
||||
func workflowTool(t *testing.T, tools []actool.CoreTool, name string) actool.CoreTool {
|
||||
t.Helper()
|
||||
for _, tool := range tools {
|
||||
if tool.Name() == name {
|
||||
return tool
|
||||
}
|
||||
}
|
||||
t.Fatalf("missing tool: %s", name)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestFindingWorkflowAutoHintToPlannerAndSetting(t *testing.T) {
|
||||
s, initial, request := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
ctx := context.Background()
|
||||
if _, err := pg.Exec(`DELETE FROM settings WHERE key=$1`, settingAgentTrafficBinding); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false) })
|
||||
if s.settingsPayload()[settingAgentTrafficBinding] != false {
|
||||
t.Fatal("missing setting must default off")
|
||||
}
|
||||
seedServerEvidenceFlow(t, s, "handoff-proof", []byte("verified local proof"))
|
||||
seedServerEvidenceFlow(t, s, "handoff-baseline", []byte("local baseline"))
|
||||
seedServerEvidenceFlow(t, s, "handoff-verification", []byte{0, 1, 255})
|
||||
// Manual binding is independent of the Agent switch.
|
||||
if r := request("POST", fmt.Sprintf("/api/exploration/findings/%d/traffic", initial.FindingID), `{"traffic_refs":[{"traffic_id":"handoff-proof"}]}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
bind := s.toolBindFindingTraffic()
|
||||
input := map[string]any{"finding_id": initial.FindingID, "traffic_refs": []db.TrafficRef{{TrafficID: "handoff-baseline"}}}
|
||||
workflowCall(t, ctx, bind, input, true)
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":true}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
if !pg.GetBool(settingAgentTrafficBinding, false) {
|
||||
t.Fatal("switch was not persisted")
|
||||
}
|
||||
f, _ := pg.GetFinding(initial.FindingID)
|
||||
task := s.m.ResolveTask(fmt.Sprint(*f.TaskID))
|
||||
refs := []db.TrafficRef{{TrafficID: "handoff-baseline", Role: "baseline", Note: "normal response"}, {TrafficID: "handoff-proof", Role: "proof", Note: "proves the finding"}, {TrafficID: "handoff-verification", Role: "verification", Note: "binary verification"}}
|
||||
// Real cross-task Auto hint handler -> persisted graph -> Planner tool.
|
||||
hintResult := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "hints": []any{map[string]any{"text": "Report the confirmed local finding with its verified evidence", "traffic_refs": refs}}}, false)
|
||||
var hints struct {
|
||||
IDs []int64 `json:"ids"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(hintResult), &hints); err != nil || len(hints.IDs) != 1 || hints.IDs[0] == 0 {
|
||||
t.Fatal(hintResult, err)
|
||||
}
|
||||
graph := workflowCall(t, ctx, s.toolGetTaskGraph(), map[string]any{"task_id": task.ID}, false)
|
||||
for _, ref := range refs {
|
||||
if !strings.Contains(graph, ref.TrafficID) {
|
||||
t.Fatal("handoff lost reference", ref.TrafficID)
|
||||
}
|
||||
}
|
||||
for range 6 {
|
||||
if _, err := task.Store.AddNode(db.KindFact, map[string]any{"summary": "ID separation"}, 0, "confirmed", "test", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ts := agent.NewToolSet(task.Store, "planner")
|
||||
ts.SetTaskID(*f.TaskID)
|
||||
ts.SetFindingRecorder(s.evidenceStore())
|
||||
notices := 0
|
||||
ts.SetNotifyFinding(func(int64, string) { notices++ })
|
||||
tools, def, cleanup := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
|
||||
defer cleanup()
|
||||
if !strings.Contains(def.FindingGuidance, "evidence_hint_id") || !strings.Contains(def.FindingGuidance, "取消 Worker") {
|
||||
t.Fatal("Planner missed runtime guidance")
|
||||
}
|
||||
report := workflowTool(t, tools, "report_finding")
|
||||
result := workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "Planner handoff fixture", "evidence_hint_id": hints.IDs[0]}, false)
|
||||
var recorded db.RecordedFinding
|
||||
if err := json.Unmarshal([]byte(strings.SplitN(result, "\n", 2)[1]), &recorded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if recorded.FindingID == recorded.NodeID || len(recorded.Traffic.Bindings) != 3 || notices != 1 {
|
||||
t.Fatal("bad finding/traffic result", result)
|
||||
}
|
||||
for i, b := range recorded.Traffic.Bindings {
|
||||
if b.Snapshot.SourceTrafficID != refs[i].TrafficID || b.Role != refs[i].Role || b.Note != refs[i].Note {
|
||||
t.Fatalf("handoff mismatch: %+v", b)
|
||||
}
|
||||
}
|
||||
list := workflowCall(t, ctx, s.toolListTaskFindings(), map[string]any{"task_id": task.ID}, false)
|
||||
var nodes []struct {
|
||||
ID int64 `json:"id"`
|
||||
FindingID int64 `json:"finding_id"`
|
||||
FindingNodeID int64 `json:"finding_node_id"`
|
||||
Count int `json:"traffic_count"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(list), &nodes); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, n := range nodes {
|
||||
if n.ID == recorded.NodeID {
|
||||
found = n.FindingID == recorded.FindingID && n.FindingNodeID == recorded.NodeID && n.Count == 3
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("canonical IDs/count missing", list)
|
||||
}
|
||||
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
wrong := workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": int64(900000000000000000)}, true)
|
||||
if !strings.Contains(wrong, "独立漏洞记录 ID") {
|
||||
t.Fatal("ambiguous ID error", wrong)
|
||||
}
|
||||
input["finding_id"] = recorded.FindingID
|
||||
// Duplicate append preserves metadata and version.
|
||||
workflowCall(t, ctx, bind, input, false)
|
||||
after, err := pg.GetFindingTraffic(ctx, recorded.FindingID)
|
||||
if err != nil || len(after.Bindings) != 3 || after.Version != recorded.Traffic.Version || after.Bindings[0].Note != refs[0].Note {
|
||||
t.Fatal("duplicate changed evidence", after, err)
|
||||
}
|
||||
// A malformed handoff must not create a partial finding or notify Planner.
|
||||
badHint := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "text": "missing packet", "traffic_refs": []db.TrafficRef{{TrafficID: "missing"}}}, false)
|
||||
var badID int64
|
||||
fmt.Sscanf(badHint, "hint added: %d", &badID)
|
||||
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "must fail", "evidence_hint_id": badID}, true)
|
||||
if notices != 1 {
|
||||
t.Fatal("failed handoff notified before commit")
|
||||
}
|
||||
child, err := s.m.CreateTaskWithOptions("inherited handoff", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*f.TaskID}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
childID, _ := strconv.ParseInt(child.ID, 10, 64)
|
||||
childRow, _ := pg.GetTask(childID)
|
||||
t.Cleanup(func() { pg.DeleteTask(childRow.ID) })
|
||||
childCtx := agent.WithRunInfo(ctx, agent.RunInfo{TaskID: childRow.ID})
|
||||
workflowCall(t, childCtx, bind, input, true)
|
||||
workflowCall(t, childCtx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
// Existing, already assembled tools must honor a later switch-off.
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
workflowCall(t, ctx, bind, input, true)
|
||||
// Turning binding off must still save a confirmed finding from an already
|
||||
// assembled tool, even if the model sends the old optional evidence fields.
|
||||
offResult := workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "off", "evidence_hint_id": hints.IDs[0]}, false)
|
||||
var offRecord struct {
|
||||
db.RecordedFinding
|
||||
EvidenceStatus string `json:"evidence_status"`
|
||||
EvidenceNote string `json:"evidence_note"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(strings.SplitN(offResult, "\n", 2)[1]), &offRecord); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if offRecord.FindingID <= 0 || len(offRecord.Traffic.Bindings) != 0 || offRecord.EvidenceStatus != "not_bound" || !strings.Contains(offRecord.EvidenceNote, "已关闭") {
|
||||
t.Fatal("disabled binding discarded finding or bound evidence", offResult)
|
||||
}
|
||||
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TCP", "severity": "low", "summary": "no packet needed"}, false)
|
||||
if notices != 3 {
|
||||
t.Fatal("optional no-packet report failed")
|
||||
}
|
||||
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
tools, def, closeTools := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
|
||||
defer closeTools()
|
||||
if def.FindingGuidance != "" {
|
||||
t.Fatal("disabled feature still injects guidance")
|
||||
}
|
||||
for _, tool := range tools {
|
||||
if tool.Name() == "bind_finding_traffic" {
|
||||
t.Fatal("disabled binding tool exposed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingWorkflowMigrationPreservesUserConfiguration(t *testing.T) {
|
||||
s, _, _ := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint", "traffic_search", "traffic_get", "traffic_blob", "bind_finding_traffic", "get_finding_traffic"} {
|
||||
old, err := pg.GetTool(key)
|
||||
if err != nil || old == nil {
|
||||
t.Fatal("missing tool", key, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
bindings, _ := json.Marshal(old.Agents)
|
||||
pg.UpdateTool(key, old.Description, old.Schema, bindings, old.Enabled)
|
||||
})
|
||||
}
|
||||
custom := json.RawMessage(`{"type":"object","properties":{"text":{"type":"string","description":"USER TEXT"},"hints":{"type":"array","items":{"type":"object","properties":{"text":{"type":"string","description":"USER ITEM"}}}}},"required":["text"]}`)
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
if err := pg.UpdateTool(key, "USER DESCRIPTION", custom, json.RawMessage(`["custom-agent"]`), false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"traffic_search", "traffic_get"} {
|
||||
row, _ := pg.GetTool(key)
|
||||
bindings := json.RawMessage(`["custom-agent"]`)
|
||||
if key == "traffic_search" {
|
||||
bindings = json.RawMessage(`["worker"]`)
|
||||
}
|
||||
if err := pg.UpdateTool(key, row.Description, row.Schema, bindings, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := pg.SetSetting("finding_workflow_tools_v2_reporter", "false"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingWorkflowTools()
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
row, _ := pg.GetTool(key)
|
||||
if row.Enabled || row.Description != "USER DESCRIPTION" || len(row.Agents) != 1 || row.Agents[0] != "custom-agent" {
|
||||
t.Fatal("changed user tool configuration", key)
|
||||
}
|
||||
if !strings.Contains(string(row.Schema), "USER TEXT") || !strings.Contains(string(row.Schema), "USER ITEM") {
|
||||
t.Fatal("changed user schema", key)
|
||||
}
|
||||
field := "traffic_refs"
|
||||
if key == "report_finding" {
|
||||
field = "evidence_hint_id"
|
||||
}
|
||||
if !strings.Contains(string(row.Schema), field) {
|
||||
t.Fatal("missing optional field", key)
|
||||
}
|
||||
}
|
||||
search, _ := pg.GetTool("traffic_search")
|
||||
get, _ := pg.GetTool("traffic_get")
|
||||
if search.Enabled || !contains(search.Agents, "reporter") || get.Enabled || len(get.Agents) != 1 || get.Agents[0] != "custom-agent" {
|
||||
t.Fatal("default/custom reader binding migration incorrect")
|
||||
}
|
||||
if err := pg.RemoveAgentFromTool("reporter", "traffic_search"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingWorkflowTools()
|
||||
search, _ = pg.GetTool("traffic_search")
|
||||
if contains(search.Agents, "reporter") {
|
||||
t.Fatal("one-time migration undid later unbinding")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingWorkflowReporterBindsBeforeWritingReport(t *testing.T) {
|
||||
s, f, request := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
ctx := context.Background()
|
||||
if err := pg.SetBool(settingAgentTrafficBinding, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false); s.m.SetTrafficEnabled(false) })
|
||||
if err := s.m.SetTrafficEnabled(true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
seedServerEvidenceFlow(t, s, "reporter-proof", []byte("local proof payload"))
|
||||
seedServerEvidenceFlow(t, s, "reporter-baseline", []byte("local normal response"))
|
||||
tools, def, cleanup := agent.AugmentTools(ctx, "reporter", nil)
|
||||
defer cleanup()
|
||||
if !strings.Contains(def.FindingGuidance, "报告前自动关联流量") || !strings.Contains(def.FindingGuidance, "绑定成功后重新调用") {
|
||||
t.Fatal("reporter did not receive binding workflow")
|
||||
}
|
||||
for _, name := range []string{"traffic_search", "traffic_get", "get_task_worker_trace", "get_task_node_detail", "bind_finding_traffic", "get_finding_traffic", "update_finding_report"} {
|
||||
workflowTool(t, tools, name)
|
||||
}
|
||||
finding, _ := pg.GetFinding(f.FindingID)
|
||||
detail := workflowCall(t, ctx, workflowTool(t, tools, "get_task_node_detail"), map[string]any{"task_id": fmt.Sprint(*finding.TaskID), "id": f.NodeID}, false)
|
||||
if !strings.Contains(detail, `"finding_id"`) || !strings.Contains(detail, `"finding_node_id"`) {
|
||||
t.Fatal("reporter lacks explicit ID mapping")
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "traffic_search"), map[string]any{"host": "evidence.local"}, false)
|
||||
for _, id := range []string{"reporter-baseline", "reporter-proof"} {
|
||||
packet := workflowCall(t, ctx, workflowTool(t, tools, "traffic_get"), map[string]any{"id": id}, false)
|
||||
if !strings.Contains(packet, "local") {
|
||||
t.Fatal("reporter cannot inspect packet", packet)
|
||||
}
|
||||
}
|
||||
refs := []db.TrafficRef{{TrafficID: "reporter-baseline", Role: "baseline"}, {TrafficID: "reporter-proof", Role: "proof", Note: "confirmed from recorded response"}}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "bind_finding_traffic"), map[string]any{"finding_id": f.FindingID, "traffic_refs": refs}, false)
|
||||
list := workflowCall(t, ctx, workflowTool(t, tools, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
|
||||
var summary struct {
|
||||
Version int64 `json:"version"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(list), &summary); err != nil || summary.Version != 1 {
|
||||
t.Fatal(list, err)
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "update_finding_report"), map[string]any{"finding_id": f.NodeID, "evidence_version": summary.Version, "report": "## Local report\n\nVerified baseline and proof using saved evidence."}, false)
|
||||
updated, err := pg.GetFinding(f.FindingID)
|
||||
if err != nil || updated.ReportEvidenceVersion != summary.Version || updated.EvidenceVersion != summary.Version {
|
||||
t.Fatal("report did not cover post-binding version", updated, err)
|
||||
}
|
||||
// Disabling automatic binding still lets Reporter read manually bound snapshots.
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
off, offDef, closeOff := agent.AugmentTools(ctx, "reporter", nil)
|
||||
defer closeOff()
|
||||
if offDef.FindingGuidance != "" {
|
||||
t.Fatal("off reporter still receives auto-binding guidance")
|
||||
}
|
||||
for _, tool := range off {
|
||||
if tool.Name() == "traffic_search" || tool.Name() == "traffic_get" || tool.Name() == "traffic_blob" || tool.Name() == "bind_finding_traffic" {
|
||||
t.Fatal("auto-binding tool exposed while off", tool.Name())
|
||||
}
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, off, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
|
||||
workflowTool(t, off, "update_finding_report")
|
||||
}
|
||||
+30
-3
@@ -2,6 +2,7 @@ package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
@@ -35,15 +36,15 @@ func (s *Server) wireInterceptReviewer() {
|
||||
}
|
||||
prov, _, ok := s.providerForProfile(profileID)
|
||||
if !ok {
|
||||
return intercept.Decision{}, fmt.Errorf("裁判模型 profile %d 不可用", profileID)
|
||||
return intercept.Decision{ProfileID: profileID}, fmt.Errorf("裁判模型 profile %d 不可用", profileID)
|
||||
}
|
||||
user := fmt.Sprintf("Tool: %s\nArguments:\n%s", tool, command)
|
||||
text, err := streamCollectText(ctx, prov, prompt, user)
|
||||
if err != nil {
|
||||
return intercept.Decision{}, err
|
||||
return intercept.Decision{ProfileID: profileID}, err
|
||||
}
|
||||
v := intercept.ParseVerdict(text)
|
||||
return intercept.Decision{Action: v.Action, Message: v.Reason}, nil
|
||||
return intercept.Decision{Action: v.Action, Message: v.Reason, ProfileID: profileID}, nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -278,6 +279,10 @@ func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
if err := s.m.interceptor.Decide(id, req.Decision == "allowed"); err != nil {
|
||||
if errors.Is(err, intercept.ErrAlreadyDecided) {
|
||||
writeErr(w, 409, err.Error())
|
||||
return
|
||||
}
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
@@ -396,3 +401,25 @@ func validateInterceptRuleReq(req interceptRuleReq) error {
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) interceptDetail(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok || id <= 0 {
|
||||
writeErr(w, 400, "bad approval id")
|
||||
return
|
||||
}
|
||||
detail, err := pg.GetInterceptDetail(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if detail == nil {
|
||||
writeErr(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, detail)
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
func TestInterceptDetailHTTP(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skip("no test database configured")
|
||||
}
|
||||
d, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
// Exercise the authenticated HTTP surface without starting unrelated task
|
||||
// schedulers or mutating the process-global tool assembly used by other tests.
|
||||
m := &Manager{pg: d, interceptor: intercept.New(d)}
|
||||
s := &Server{m: m, jwtKey: []byte("approval-http-test-signing-key")}
|
||||
h := s.Handler()
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, err := m.pg.CreateInterceptPending(0, 0, "approval-http", "test", "Write", []byte(`{}`), "[模型] 确认", &db.InterceptAudit{InitialAction: "ask", UserMessage: "snapshot", ExecutionStatus: "not_started"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = m.pg.Exec(`DELETE FROM intercept_pending WHERE id=$1`, id) })
|
||||
do := func(method, path, body string, authenticated bool) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
if authenticated {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
path := fmt.Sprintf("/api/intercept/history/%d", id)
|
||||
if r := do(http.MethodGet, path, "", false); r.Code != 401 {
|
||||
t.Fatalf("unprotected detail: %d", r.Code)
|
||||
}
|
||||
r := do(http.MethodGet, path, "", true)
|
||||
var detail db.InterceptDetail
|
||||
if r.Code != 200 || json.Unmarshal(r.Body.Bytes(), &detail) != nil || detail.Audit == nil || detail.Audit.UserMessage != "snapshot" {
|
||||
t.Fatalf("detail: %d %s", r.Code, r.Body.String())
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
path string
|
||||
code int
|
||||
}{{"/api/intercept/history/not-a-number", 400}, {"/api/intercept/history/0", 400}, {"/api/intercept/history/9223372036854775807", 404}} {
|
||||
if r := do(http.MethodGet, tc.path, "", true); r.Code != tc.code {
|
||||
t.Fatalf("%s: %d", tc.path, r.Code)
|
||||
}
|
||||
}
|
||||
decisionPath := fmt.Sprintf("/api/intercept/pending/%d/decide", id)
|
||||
if r := do(http.MethodPost, decisionPath, `{"decision":"denied"}`, true); r.Code != 200 {
|
||||
t.Fatalf("decide: %d %s", r.Code, r.Body.String())
|
||||
}
|
||||
if r := do(http.MethodPost, decisionPath, `{"decision":"allowed"}`, true); r.Code != 409 {
|
||||
t.Fatalf("duplicate decision: %d", r.Code)
|
||||
}
|
||||
}
|
||||
@@ -78,3 +78,29 @@ func (e *Engine) modelErrorRetryPolicy() (retries int, backoff time.Duration) {
|
||||
}
|
||||
return retries, backoff
|
||||
}
|
||||
|
||||
// emptyTurnNudgeLimit resolves how many empty-turn continuations one work may
|
||||
// inject (see steerHooks.Stop). It deliberately reuses layer ②'s knob —— 「空响应
|
||||
// 重试次数」:两者是同一件事的两种手段。SDK 那层管「一个内容块都没有」，手段是把
|
||||
// 同一个请求原样重发;这里管「只有思考、既无正文也无工具」，手段是追加一条指令让
|
||||
// 模型带着已有思考接着做(原样重发对这种由上下文形状决定的空转没有意义)。判空口径
|
||||
// 不同是因为 SDK 以「有没有 yield 过事件」为准，而思考增量本身就是事件——但用户配
|
||||
// 「空响应重试几次」时想表达的是「模型没产出实质内容就再来一次」，两层共用一个次数
|
||||
// 才对得上这个心智。
|
||||
//
|
||||
// 读全局策略而不是某个 profile 的覆盖:一个 run 中途可能因故障转移换 profile，而这
|
||||
// 是整条意图的总量闸，不该跟着换端点而变。语义与 SDK 的 emptyRetries() 同构:
|
||||
// 0 = 默认 defaultEmptyTurnNudges;-1(负) = 关闭空转续跑;>0 = 用该值。
|
||||
func (e *Engine) emptyTurnNudgeLimit() int {
|
||||
if e == nil || e.m == nil || e.m.pg == nil {
|
||||
return defaultEmptyTurnNudges
|
||||
}
|
||||
switch n := e.m.pg.LLMRetryPolicy().Empty.Attempts; {
|
||||
case n == 0:
|
||||
return defaultEmptyTurnNudges
|
||||
case n < 0:
|
||||
return 0
|
||||
default:
|
||||
return n
|
||||
}
|
||||
}
|
||||
+11
-6
@@ -249,12 +249,13 @@ type Manager struct {
|
||||
|
||||
// Settings keys the UI toggles at runtime.
|
||||
const (
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingAgentTrafficBinding = "agent_traffic_binding"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
// settingGlobalProxy is the global egress proxy for all target traffic
|
||||
// (http/https/socks5). Empty = direct. Distinct from web_search_proxy (which
|
||||
// only routes the search backend) and the per-profile LLM proxy.
|
||||
@@ -355,6 +356,10 @@ func NewManager(dir, proxyAddr string) (*Manager, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := pg.RecoverFindingRetests(); err != nil {
|
||||
pg.Close()
|
||||
return nil, fmt.Errorf("recover finding retests: %w", err)
|
||||
}
|
||||
if err := pg.EnsureLLMRecordsTable(); err != nil {
|
||||
log.Printf("[llmrec] create table: %v", err)
|
||||
}
|
||||
|
||||
+78
-23
@@ -39,6 +39,7 @@ func jsonResult(v any) (actool.Result, error) {
|
||||
//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()...) // 平台操作工具(建改 skill/工具/MCP，给 Auto 用)
|
||||
custom, err := s.customTools()
|
||||
if err != nil {
|
||||
@@ -75,6 +76,8 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
|
||||
s.toolSearchWorkerTraces(),
|
||||
s.toolGetTaskNodeDetail(),
|
||||
s.toolUpdateFindingReport(),
|
||||
s.toolGetFindingTraffic(),
|
||||
s.toolBindFindingTraffic(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -333,10 +336,11 @@ 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": strParam("任务 id"),
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组，每个元素字段同顶层（text/asset_ids/traffic_refs）。", "items": objSchema(map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()})},
|
||||
"text": strParam("[单条] 提示内容"),
|
||||
"traffic_refs": agent.HintTrafficSchema(),
|
||||
"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)
|
||||
@@ -392,20 +396,22 @@ 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": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||||
"report": strParam("详细报告全文,Markdown 格式"),
|
||||
"evidence_version": map[string]any{"type": "integer", "description": "get_finding_traffic 返回的证据 version；用于防止报告覆盖新的证据变更"},
|
||||
}, "finding_id", "report"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Report string `json:"report"`
|
||||
EvidenceVersion *int64 `json:"evidence_version"`
|
||||
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)
|
||||
n, err := s.m.pg.SetFindingReportVersionByNodeID(ctx, nodeID, a.Report, a.EvidenceVersion)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
@@ -439,7 +445,11 @@ func (s *Server) seedOrchestrationTools() {
|
||||
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)
|
||||
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())
|
||||
@@ -453,11 +463,14 @@ func (s *Server) seedOrchestrationTools() {
|
||||
s.seedWorkerReadToolsUnbind() // list_facts/node_detail/list_companies/跨 work 检索从 worker 默认解绑(一次性)
|
||||
s.seedAutoReportFindingBinding()
|
||||
s.unbindGoalMetDefault()
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.upgradeReporterTriggerMessage() // 老库补迁移:让 reporter 回传 evidence_version(一次性)
|
||||
s.seedFindingTrafficTools() // 增加可选证据参数及只读证据工具，保留用户配置
|
||||
s.seedFindingWorkflowTools()
|
||||
// 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||||
// pentest 一起 seed 好了（项目尚无旧库，不做迁移）。
|
||||
@@ -639,6 +652,50 @@ func (s *Server) reseedWorkerPrompt() {
|
||||
log.Printf("[prompts] worker 提示词已追加新默认版本(查上下文段收敛为 list_assets/list_findings,去掉 list_facts/node_detail/asset_neighbors,一次性)")
|
||||
}
|
||||
|
||||
// reporterToolCallMessage 必须无条件要求先读一次 get_finding_traffic 再写报告。
|
||||
// 该工具是只读的、「不依赖捕获开关」,自动绑定关不关都能读到人工绑定的证据。若这里
|
||||
// 写成「启用自动绑定才读」,默认关闭配置下 reporter 就不会传 evidence_version,
|
||||
// SetFindingReportVersionByNodeID 便按 legacy 语义写 -1,漏洞详情与 Markdown 导出
|
||||
// 从此常驻「证据已变更，报告待更新」,而 UI 上没有任何入口能把它清掉。
|
||||
const reporterToolCallMessage = "上面刚有一个漏洞被 report_finding 登记。请读取返回 JSON 的 finding_id（独立漏洞记录 ID）与 finding_node_id（探索节点 ID），" +
|
||||
"先用 get_finding_traffic(finding_id) 读取当前证据清单及其 version（空清单是正常情况，照常写报告）；" +
|
||||
"若运行指引启用自动绑定，在读取前先核实并关联本次漏洞的流量。节点详情使用 finding_node_id。" +
|
||||
"最后调用 update_finding_report(finding_id=finding_node_id, report, evidence_version=实际读取版本) 保存，" +
|
||||
"evidence_version 必须传，否则报告会被永久标记为待更新。不要混用两种编号。"
|
||||
|
||||
// 旧版触发消息(0.3.8 及更早)。只有仍与它逐字相同的记录才会被迁移覆盖，用户改过的保持原样。
|
||||
const reporterToolCallMessageV1 = "上面刚有一个漏洞被 report_finding 登记。请从触发上下文里取出 finding_id" +
|
||||
"（工具返回 \"finding recorded: <id>\" 里的数字）与任务 id，按你的职责撰写该漏洞的详细报告，" +
|
||||
"最后调用 update_finding_report(finding_id, report) 保存。"
|
||||
|
||||
// upgradeReporterTriggerMessage 把老库里仍是默认文案的 reporter 触发消息刷成新版本。
|
||||
// seedReporterAgent 受 reporter_agent_seed_v1 守卫且只在新建 agent 时写触发器，所以
|
||||
// 升级上来的库拿不到新文案 —— 工具 schema 由 seedFindingTrafficTools 补齐了
|
||||
// evidence_version，但没有任何东西告诉 reporter 去用它。一次性，且只覆盖未被改动的文案。
|
||||
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") }() // 只尝试一次
|
||||
triggers, err := s.m.pg.ListTriggersFor("reporter")
|
||||
if err != nil {
|
||||
log.Printf("[reporter] 读取触发器失败: %v", err)
|
||||
return
|
||||
}
|
||||
for _, t := range triggers {
|
||||
if !t.OnToolCall || t.ToolCallMessage != reporterToolCallMessageV1 {
|
||||
continue // 用户改过或不是 finding 触发器，不动。
|
||||
}
|
||||
t.ToolCallMessage = reporterToolCallMessage
|
||||
if err := s.m.pg.UpdateTrigger(t); err != nil {
|
||||
log.Printf("[reporter] 升级触发消息失败: %v", err)
|
||||
return
|
||||
}
|
||||
log.Printf("[reporter] 触发消息已升级为读取并回传 evidence_version")
|
||||
}
|
||||
}
|
||||
|
||||
// seedReporterAgent 预置一个「报告撰写」自定义 agent(builtin=false，可在 UI 编辑/删除)：
|
||||
// 绑定 update_finding_report + 任务查询工具，并挂一个「report_finding 被调用即触发」的
|
||||
// 触发器 —— 每登记一个漏洞就唤起它写详细报告。一次性(settings flag 守卫)：用户删掉后不再重建。
|
||||
@@ -679,13 +736,11 @@ func (s *Server) seedReporterAgent() {
|
||||
// 触发器：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) 保存。",
|
||||
AgentKey: "reporter",
|
||||
Enabled: true,
|
||||
OnToolCall: true,
|
||||
ToolNames: []string{"report_finding"},
|
||||
ToolCallMessage: reporterToolCallMessage,
|
||||
}); err != nil {
|
||||
log.Printf("[reporter] 创建触发器失败: %v", err)
|
||||
}
|
||||
|
||||
@@ -305,6 +305,11 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) {
|
||||
if len(want) == 0 {
|
||||
continue
|
||||
}
|
||||
// A failed finding write has no committed finding to report. Keep other
|
||||
// tool-error triggers available for user-defined automation.
|
||||
if e.ToolIsErr && e.Tool == "report_finding" {
|
||||
continue
|
||||
}
|
||||
errTag := ""
|
||||
if e.ToolIsErr {
|
||||
errTag = "[error] "
|
||||
|
||||
+52
-28
@@ -1,7 +1,6 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -11,6 +10,7 @@ import (
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
@@ -25,7 +25,6 @@ import (
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
"github.com/Autumn-27/artex/report"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/memory"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
@@ -220,9 +219,13 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
}
|
||||
wireAgentAugment(m.pg, s.skillDir, s.hostTools) // 可见 skills/MCP + 流量/编排 host 工具装配进 agent 工具集
|
||||
domainReg := buildDomainReg(m.Assets())
|
||||
wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值
|
||||
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
|
||||
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
|
||||
wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值
|
||||
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
|
||||
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
|
||||
if err := s.seedFindingRetester(); err != nil {
|
||||
log.Printf("[retester] seed: %v", err)
|
||||
}
|
||||
go s.evidenceStore().RunGC(s.ctx)
|
||||
s.seedPythonInterpreter() // 自定义脚本工具:开机检测 python 解释器入库(仅空时)
|
||||
go newScheduler(s).Run(s.ctx) // P3 触发器调度(定时/finding/目标事件),仅自定义 agent
|
||||
// Fill the tool cache for any enabled MCP that has none yet (notably the
|
||||
@@ -415,15 +418,16 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
|
||||
// the tools-table binding (default = worker), so worker behavior is unchanged.
|
||||
wProv, wCfg := s.providerForAgent("worker", pinID, gProv, gCfg)
|
||||
wk := agent.NewWorker(wProv, wCfg.Model, s.m.dir, tx, wCfg.CompactionWindow(), s.agentMaxTurns("worker"))
|
||||
wk.SetFindingRecorder(s.evidenceStore())
|
||||
wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second)
|
||||
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
|
||||
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
|
||||
wk.SetWebSearch(s.webSearchFor("worker"))
|
||||
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开;每轮读)
|
||||
wk.SetNonStreaming(nonStreamingResolver(wCfg)) // 该 profile 选非流式时走 Provider.Complete
|
||||
wk.SetMaxTokens(maxTokensResolver(wCfg)) // 单次回复输出上限(0 = 不发)
|
||||
pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg)
|
||||
pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner"))
|
||||
pl.SetFindingRecorder(s.evidenceStore())
|
||||
pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work
|
||||
pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction
|
||||
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
@@ -496,6 +500,7 @@ func (s *Server) applyLLM(cfg agent.Config) error {
|
||||
mProv, mCfg := s.providerForAgent("mainagent", nil, prov, cfg)
|
||||
s.cfgMu.Lock()
|
||||
s.mainAgent = agent.NewMainAgent(mProv, mCfg.Model, s.m.dir, tx, mCfg.CompactionWindow(), s.agentMaxTurns("mainagent"))
|
||||
s.mainAgent.SetFindingRecorder(s.evidenceStore())
|
||||
s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
s.mainAgent.SetWebSearch(s.webSearchFor("mainagent"))
|
||||
s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏
|
||||
@@ -827,9 +832,13 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /api/exploration/findings/asset-tree", s.findingAssetTree)
|
||||
mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats)
|
||||
mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport)
|
||||
s.registerFindingTraffic(mux)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage)
|
||||
mux.HandleFunc("POST /api/exploration/findings/{id}/deepen", s.deepenFinding)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}/retests", s.listFindingRetests)
|
||||
mux.HandleFunc("GET /api/exploration/findings/retests/active", s.listActiveFindingRetests)
|
||||
mux.HandleFunc("POST /api/exploration/findings/{id}/retests", s.startFindingRetest)
|
||||
mux.HandleFunc("PATCH /api/exploration/findings/{id}", s.patchFinding)
|
||||
mux.HandleFunc("DELETE /api/exploration/findings/{id}", s.deleteFinding)
|
||||
mux.HandleFunc("GET /api/exploration/intents", s.intents)
|
||||
@@ -965,6 +974,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /api/intercept/pending/{id}", s.interceptGetOne)
|
||||
mux.HandleFunc("POST /api/intercept/pending/{id}/decide", s.interceptDecide)
|
||||
mux.HandleFunc("GET /api/intercept/history", s.interceptHistory)
|
||||
mux.HandleFunc("GET /api/intercept/history/{id}", s.interceptDetail)
|
||||
mux.HandleFunc("GET /api/intercept/task/{taskID}", s.interceptListTaskItems)
|
||||
mux.HandleFunc("GET /api/intercept/tool-config", s.interceptGetToolConfig)
|
||||
mux.HandleFunc("PUT /api/intercept/tool-config", s.interceptSetToolConfig)
|
||||
@@ -2151,6 +2161,16 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
stage, err := os.MkdirTemp("", "artex-finding-export-")
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer os.RemoveAll(stage)
|
||||
if err = s.evidenceStore().StageFindingsExport(r.Context(), fs, stage, format == "md-zip"); err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
stamp := now.Format("20060102-150405")
|
||||
setDownload := func(contentType, filename string) {
|
||||
@@ -2163,27 +2183,19 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
setDownload("text/markdown; charset=utf-8", "findings-"+stamp+".md")
|
||||
_, _ = w.Write([]byte(report.FindingsMarkdown(fs, now)))
|
||||
case "md-zip":
|
||||
path := filepath.Join(stage, "findings.zip")
|
||||
if err := buildFindingsEvidenceZip(path, fs, stage, now); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
setDownload("application/zip", "findings-"+stamp+".zip")
|
||||
zw := zip.NewWriter(w)
|
||||
used := map[string]int{}
|
||||
for _, f := range fs {
|
||||
base := report.FindingFilename(f)
|
||||
name := base
|
||||
// 去重:同名文件追加 -2、-3……
|
||||
if n := used[base]; n > 0 {
|
||||
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
|
||||
}
|
||||
used[base]++
|
||||
fw, werr := zw.Create(name)
|
||||
if werr != nil {
|
||||
log.Printf("[findings-export] zip create %s: %v", name, werr)
|
||||
continue
|
||||
}
|
||||
_, _ = fw.Write([]byte(report.SingleFindingMarkdown(f, now)))
|
||||
}
|
||||
if cerr := zw.Close(); cerr != nil {
|
||||
log.Printf("[findings-export] zip close: %v", cerr)
|
||||
}
|
||||
http.ServeContent(w, r, "findings.zip", now, file)
|
||||
case "csv":
|
||||
setDownload("text/csv; charset=utf-8", "findings-"+stamp+".csv")
|
||||
_, _ = w.Write(report.FindingsCSV(fs))
|
||||
@@ -2191,6 +2203,10 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
assets := s.resolveFindingAssets(fs)
|
||||
out := make([]FindingDTO, 0, len(fs))
|
||||
for _, f := range fs {
|
||||
for i := range f.TrafficBindings {
|
||||
f.TrafficBindings[i].Snapshot.ReqHead = ""
|
||||
f.TrafficBindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
out = append(out, findingFromDB(f, assets))
|
||||
}
|
||||
setDownload("application/json; charset=utf-8", "findings-"+stamp+".json")
|
||||
@@ -3138,6 +3154,7 @@ func (s *Server) settingsPayload() map[string]any {
|
||||
}
|
||||
return map[string]any{
|
||||
"traffic_capture": s.m.TrafficEnabled(),
|
||||
"agent_traffic_binding": s.m.pg.GetBool(settingAgentTrafficBinding, false),
|
||||
"llm_record": s.m.LLMRecordEnabled(),
|
||||
"web_search_enabled": on,
|
||||
"web_search_backend": backend,
|
||||
@@ -3178,8 +3195,9 @@ func (s *Server) pgDetectPython(w http.ResponseWriter, r *http.Request) {
|
||||
// off, agents get no proxy config, no traffic tools, and no proxy prompt content.
|
||||
func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
TrafficCapture *bool `json:"traffic_capture"`
|
||||
LLMRecord *bool `json:"llm_record"` // LLM 录制开关（默认关）；即时生效，无需重建 agent
|
||||
TrafficCapture *bool `json:"traffic_capture"`
|
||||
AgentTrafficBinding *bool `json:"agent_traffic_binding"`
|
||||
LLMRecord *bool `json:"llm_record"` // LLM 录制开关（默认关）；即时生效，无需重建 agent
|
||||
// Web search. WebSearchEnabled/Backend toggle the tool + backend; BraveKey/TavilyKey
|
||||
// are optional — omit (null) to leave a stored key untouched, send "" to clear.
|
||||
WebSearchEnabled *bool `json:"web_search_enabled"`
|
||||
@@ -3276,6 +3294,12 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
// built under the OLD switch state — drop them so they pick up the new one.
|
||||
s.invalidateProfileAgents()
|
||||
}
|
||||
if req.AgentTrafficBinding != nil {
|
||||
if err := s.m.pg.SetBool(settingAgentTrafficBinding, *req.AgentTrafficBinding); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
if req.TrafficCapture != nil {
|
||||
if err := s.m.SetTrafficEnabled(*req.TrafficCapture); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
|
||||
+20
-2
@@ -181,6 +181,13 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
snapshot.DataCounts["traffic"] = trafficCount
|
||||
}
|
||||
archivePath := taskArchivePath(s.m.dir, job.ID, taskID)
|
||||
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = s.evidenceStore().CopySnapshots(s.ctx, evidenceSnapshots, filepath.Join(fileStage.payload, "evidence")); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50)
|
||||
originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot)
|
||||
if err != nil {
|
||||
@@ -273,6 +280,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
|
||||
return pgdb.ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(&snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.evidenceStore().WithInstalledSnapshots(s.ctx, evidenceSnapshots, filepath.Join(extracted, "evidence"), func() error {
|
||||
return s.restoreTaskArchivePayload(job, &snapshot, extracted)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) restoreTaskArchivePayload(job *pgdb.TaskArchive, snapshot *pgdb.TaskArchiveSnapshot, extracted string) (runErr error) {
|
||||
var err error
|
||||
var llmRecordsFile *os.File
|
||||
relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"]
|
||||
if len(snapshot.StreamedTables) > 0 &&
|
||||
@@ -312,13 +330,13 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
var warnings []string
|
||||
if llmRecordsFile != nil {
|
||||
warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords(
|
||||
job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
|
||||
job.ID, snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
|
||||
)
|
||||
closeErr := llmRecordsFile.Close()
|
||||
llmRecordsFile = nil
|
||||
err = errors.Join(err, closeErr)
|
||||
} else {
|
||||
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
|
||||
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, snapshot, job.RemainingTimeoutSeconds)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
+3
-2
@@ -16,7 +16,6 @@ import (
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/memory"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
|
||||
@@ -563,15 +562,16 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
|
||||
tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts"))
|
||||
window := workerRuntime.CompactionWindow()
|
||||
wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker"))
|
||||
wk.SetFindingRecorder(s.evidenceStore())
|
||||
wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow)
|
||||
wk.SetNonStreaming(workerRuntime.nonStreaming) // 按任务当前激活 profile 的流式开关(每轮读)
|
||||
wk.SetMaxTokens(workerRuntime.maxTokens) // 同上,输出上限也跟随当前激活 profile
|
||||
wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second)
|
||||
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
|
||||
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
|
||||
wk.SetWebSearch(s.webSearchFor("worker"))
|
||||
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开)
|
||||
pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner"))
|
||||
pl.SetFindingRecorder(s.evidenceStore())
|
||||
pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow)
|
||||
pl.SetNonStreaming(plannerRuntime.nonStreaming)
|
||||
pl.SetMaxTokens(plannerRuntime.maxTokens)
|
||||
@@ -581,6 +581,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
|
||||
pl.SetWebSearch(s.webSearchFor("planner"))
|
||||
pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开)
|
||||
main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent"))
|
||||
main.SetFindingRecorder(s.evidenceStore())
|
||||
main.SetCompactionWindowResolver(mainRuntime.CompactionWindow)
|
||||
main.SetNonStreaming(mainRuntime.nonStreaming)
|
||||
main.SetMaxTokens(mainRuntime.maxTokens)
|
||||
|
||||
Reference in new issue
Block a user