From 8ad507e31f1c22afd31aaf85339cf62c3941c201 Mon Sep 17 00:00:00 2001 From: RuoJi6 Date: Tue, 27 Oct 2026 17:56:50 +0800 Subject: [PATCH] fix(findings): address grok review on traffic step_index UX, ref merge, tool schema migration --- agent/finding_recorder_test.go | 48 ++++ agent/tools.go | 28 ++- server/finding_workflow.go | 46 ++++ server/finding_workflow_test.go | 10 + .../finding-reproduction-section.tsx | 206 ++++++++++++++-- web/src/components/finding-traffic-panel.tsx | 220 ++++++++++-------- 6 files changed, 433 insertions(+), 125 deletions(-) diff --git a/agent/finding_recorder_test.go b/agent/finding_recorder_test.go index 07e58fa..47b5c48 100644 --- a/agent/finding_recorder_test.go +++ b/agent/finding_recorder_test.go @@ -19,6 +19,15 @@ func (failingFindingRecorder) Record(context.Context, db.RecordFindingInput, []d return nil, errors.New("fixture persistence failed") } +type capturingFindingRecorder struct { + refs []db.TrafficRef +} + +func (r *capturingFindingRecorder) Record(_ context.Context, _ db.RecordFindingInput, refs []db.TrafficRef) (*db.RecordedFinding, error) { + r.refs = append([]db.TrafficRef(nil), refs...) + return &db.RecordedFinding{Traffic: &db.FindingTraffic{}}, nil +} + func TestReportFindingOptionalTrafficWithoutCapture(t *testing.T) { d := testDB(t) defer d.Close() @@ -90,6 +99,45 @@ func TestReportFindingPersistsStructuredReproduction(t *testing.T) { } } +func TestReportFindingStepTrafficPreservesExplicitReferenceMetadata(t *testing.T) { + old := FindingTrafficBindingEnabled + FindingTrafficBindingEnabled = func() bool { return true } + t.Cleanup(func() { FindingTrafficBindingEnabled = old }) + d := testDB(t) + defer d.Close() + task, err := d.CreateTask("step traffic merge", "fixture", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer d.DeleteTask(task.ID) + recorder := &capturingFindingRecorder{} + ts := NewToolSet(d.Exploration(task.ExplorationID), "fixture") + ts.taskID = task.ID + ts.SetFindingRecorder(recorder) + reproduction := `{"steps":[{"traffic_ids":["shared"]},{"traffic_ids":["new-step"]}]}` + input := `{"vulnclass":"TEST","severity":"low","summary":"step merge","traffic_refs":[{"traffic_id":"shared","role":"proof","note":"preserve this note"},{"traffic_id":"overview","role":"baseline","note":"overview note"}],"reproduction":` + reproduction + `}` + result, err := ts.addFinding().Call(t.Context(), json.RawMessage(input), nil) + if err != nil || result.IsError { + t.Fatalf("report failed: %v %s", err, result.Flatten()) + } + if len(recorder.refs) != 3 { + t.Fatalf("unexpected refs: %#v", recorder.refs) + } + byID := make(map[string]db.TrafficRef, len(recorder.refs)) + for _, ref := range recorder.refs { + byID[ref.TrafficID] = ref + } + if got := byID["shared"]; got.Role != "proof" || got.Note != "preserve this note" || got.StepIndex != 1 { + t.Fatalf("overlapping ref metadata changed: %#v", got) + } + if got := byID["overview"]; got.Role != "baseline" || got.Note != "overview note" || got.StepIndex != 0 { + t.Fatalf("overview ref changed: %#v", got) + } + if got := byID["new-step"]; got.Role != "supporting" || got.Note != "" || got.StepIndex != 2 { + t.Fatalf("new step ref not defaulted: %#v", got) + } +} + func TestReportFindingAtomicContract(t *testing.T) { old := FindingTrafficBindingEnabled FindingTrafficBindingEnabled = func() bool { return true } diff --git a/agent/tools.go b/agent/tools.go index 6dea34a..632b567 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -1270,8 +1270,9 @@ func (t *ToolSet) addFinding() actool.CoreTool { if err != nil { return actool.Errorf(err.Error()), nil } - // Step traffic_ids are resolved at report time. Explicit traffic_refs that - // are not listed by a step remain overview bindings (step_index=0). + // Step traffic_ids are resolved at report time. Keep metadata supplied by + // traffic_refs when an id is mentioned by a reproduction step; only ids + // that are absent from traffic_refs get the supporting default. if findingTrafficBindingEnabled() && len(a.Reproduction) > 0 && string(a.Reproduction) != "null" { var recipe struct { Steps []struct { @@ -1281,19 +1282,30 @@ func (t *ToolSet) addFinding() actool.CoreTool { if err := json.Unmarshal(a.Reproduction, &recipe); err != nil { return actool.Errorf(fmt.Sprintf("reproduction JSON 无效: %v", err)), nil } - seen := map[string]bool{} - stepRefs := make([]db.TrafficRef, 0) + refIndex := make(map[string]int, len(refs)) + for i := range refs { + id := strings.TrimSpace(refs[i].TrafficID) + if id != "" { + refs[i].TrafficID = id + if _, exists := refIndex[id]; !exists { + refIndex[id] = i + } + } + } for i, step := range recipe.Steps { for _, id := range step.TrafficIDs { id = strings.TrimSpace(id) - if id == "" || seen[id] { + if id == "" { continue } - seen[id] = true - stepRefs = append(stepRefs, db.TrafficRef{TrafficID: id, Role: "supporting", StepIndex: i + 1}) + if ref, exists := refIndex[id]; exists { + refs[ref].StepIndex = i + 1 + continue + } + refIndex[id] = len(refs) + refs = append(refs, db.TrafficRef{TrafficID: id, Role: "supporting", StepIndex: i + 1}) } } - refs = append(stepRefs, refs...) } input := db.RecordFindingInput{TaskID: t.taskID, ExplorationID: t.ts.ID(), IntentID: pid(a.IntentID), VulnClass: a.VulnClass, Name: a.Name, Severity: a.Severity, Summary: a.Summary, Evidence: a.Evidence, Worker: t.worker, AssetIDs: pidList(a.AssetIDs), Reproduction: a.Reproduction} var recorded *db.RecordedFinding diff --git a/server/finding_workflow.go b/server/finding_workflow.go index a5e1eb8..de4bb19 100644 --- a/server/finding_workflow.go +++ b/server/finding_workflow.go @@ -15,6 +15,7 @@ import ( ) func (s *Server) seedFindingWorkflowTools() { + s.migrateFindingWorkflowSchemas() const hostSearchDescriptionFlag = "finding_workflow_tools_v3_host_search_description" if value, _, _ := s.m.pg.GetSetting(hostSearchDescriptionFlag); value != "true" { // Only replace the original built-in text. A user-edited description is @@ -106,6 +107,51 @@ func (s *Server) seedFindingWorkflowTools() { _ = s.m.pg.SetSetting(flag, "true") } +// migrateFindingWorkflowSchemas adds schema fields introduced after the +// original finding workflow seed. It patches only missing properties so a +// user's description, agent bindings, enabled state, and other schema edits +// remain untouched. +func (s *Server) migrateFindingWorkflowSchemas() { + const flag = "finding_workflow_tools_v3_reproduction_step_index" + if value, _, _ := s.m.pg.GetSetting(flag); value == "true" { + return + } + properties := map[string]json.RawMessage{} + for _, seed := range agent.BuiltinToolSeeds() { + if seed.Key != "report_finding" { + continue + } + if raw, err := json.Marshal(objectProperty(seed.Schema, "properties")["reproduction"]); err == nil && string(raw) != "null" { + properties["report_finding"] = raw + } + break + } + for _, tool := range s.orchestrationTools() { + if tool.Name() != "bind_finding_traffic" { + continue + } + schema := tool.InputSchema() + if raw, err := json.Marshal(objectProperty(schema, "properties")["step_index"]); err == nil && string(raw) != "null" { + properties["bind_finding_traffic"] = raw + } + break + } + for key, property := range map[string]string{"report_finding": "reproduction", "bind_finding_traffic": "step_index"} { + raw := properties[key] + if len(raw) == 0 { + log.Printf("[evidence] missing code schema property %s.%s", key, property) + continue + } + if _, err := s.m.pg.Exec(`UPDATE tools +SET schema=jsonb_set(COALESCE(schema,'{}'::jsonb),ARRAY['properties',$2::text],$3::jsonb,true),updated_at=now() +WHERE key=$1 AND system AND NOT(COALESCE(schema->'properties','{}'::jsonb) ? $2)`, key, property, string(raw)); err != nil { + log.Printf("[evidence] upgrade %s.%s: %v", key, property, err) + 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 { diff --git a/server/finding_workflow_test.go b/server/finding_workflow_test.go index ada3b10..4b0cbc7 100644 --- a/server/finding_workflow_test.go +++ b/server/finding_workflow_test.go @@ -227,6 +227,9 @@ func TestFindingWorkflowMigrationPreservesUserConfiguration(t *testing.T) { if err := pg.SetSetting("finding_workflow_tools_v2_reporter", "false"); err != nil { t.Fatal(err) } + if err := pg.SetSetting("finding_workflow_tools_v3_reproduction_step_index", "false"); err != nil { + t.Fatal(err) + } s.seedFindingWorkflowTools() for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} { row, _ := pg.GetTool(key) @@ -243,6 +246,13 @@ func TestFindingWorkflowMigrationPreservesUserConfiguration(t *testing.T) { if !strings.Contains(string(row.Schema), field) { t.Fatal("missing optional field", key) } + if key == "report_finding" && !strings.Contains(string(row.Schema), "reproduction") { + t.Fatal("missing reproduction schema migration") + } + } + bindSchema, _ := pg.GetTool("bind_finding_traffic") + if !strings.Contains(string(bindSchema.Schema), "step_index") { + t.Fatal("missing bind_finding_traffic step_index migration") } search, _ := pg.GetTool("traffic_search") get, _ := pg.GetTool("traffic_get") diff --git a/web/src/components/finding-reproduction-section.tsx b/web/src/components/finding-reproduction-section.tsx index 74708db..34c0770 100644 --- a/web/src/components/finding-reproduction-section.tsx +++ b/web/src/components/finding-reproduction-section.tsx @@ -3,6 +3,7 @@ import * as React from "react"; import { ExternalLinkIcon, PlusIcon } from "lucide-react"; +import { toast } from "sonner"; import { HttpCodeBlock } from "@/components/http-code-block"; import { TrafficEvidenceViewer } from "@/components/traffic-evidence-viewer"; @@ -10,8 +11,27 @@ import { TrafficPickerDialog } from "@/components/traffic-picker-dialog"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; import { api } from "@/lib/api"; -import type { FindingTraffic, Reproduction } from "@/lib/types"; +import type { FindingTraffic, FindingTrafficBinding, Reproduction, TrafficEvidenceRole } from "@/lib/types"; + +const ROLES: Record = { + baseline: "正常对照", + proof: "漏洞证明", + verification: "补充验证", + supporting: "辅助证据", +}; export function hasStructuredReproduction(reproduction?: Reproduction | null) { return Boolean(reproduction && ((reproduction.steps?.length ?? 0) > 0 || reproduction.entry_point)); @@ -33,14 +53,75 @@ export function FindingReproductionSection({ const [traffic, setTraffic] = React.useState(null); const [pickerStep, setPickerStep] = React.useState(null); const [preview, setPreview] = React.useState(null); + const [editing, setEditing] = React.useState(null); + const [role, setRole] = React.useState("supporting"); + const [note, setNote] = React.useState(""); + const [editingStepIndex, setEditingStepIndex] = React.useState("0"); + const [busy, setBusy] = React.useState(false); const [reload, setReload] = React.useState(0); const load = React.useCallback(() => { - api.findingTraffic(findingId, contextTask).then(setTraffic).catch(() => setTraffic(null)); + api + .findingTraffic(findingId, contextTask) + .then(setTraffic) + .catch(() => setTraffic(null)); }, [findingId, contextTask]); + // biome-ignore lint/correctness/useExhaustiveDependencies: retry token intentionally refreshes this request. React.useEffect(() => load(), [load, reload]); const steps = reproduction.steps ?? []; const allBindings = traffic?.bindings ?? []; + async function mutate(action: () => Promise) { + setBusy(true); + try { + setTraffic(await action()); + setEditing(null); + onChanged?.(); + toast.success("流量证据已更新"); + } catch (error) { + toast.error((error as Error).message); + setReload((value) => value + 1); + } finally { + setBusy(false); + } + } + function editBinding(item: FindingTrafficBinding) { + setEditing(item); + setRole(item.role); + setNote(item.note); + setEditingStepIndex(String(item.step_index ?? 0)); + } + function bindingActions(item: FindingTrafficBinding) { + return ( +
+ + {ROLES[item.role]} + + {item.note ? {item.note} : null} + + {!readOnly ? ( + <> + + + + ) : null} +
+ ); + } return ( @@ -57,7 +138,9 @@ export function FindingReproductionSection({
前置条件
    - {reproduction.prerequisites?.map((item, index) =>
  • {item}
    • )} + {reproduction.prerequisites?.map((item, index) => ( +
  • {item}
    • + ))}
) : null} @@ -69,15 +152,16 @@ export function FindingReproductionSection({
步骤 {stepIndex} {step.expected ? 预期：{step.expected} : null} - {step.proves ? 证明：{step.proves} : null} + {step.proves ? 证明：{step.proves} : null}
- {bound.map((item) => ( - - ))} + {bound.map(bindingActions)} {!readOnly ? ( - ) : null} @@ -86,15 +170,30 @@ export function FindingReproductionSection({ {step.request?.trim() ? : null} {step.response?.trim() ? (
-
响应
+
响应
) : null} - {bound.length > 0 ?
已绑定 {bound.length} 条流量
: null} + {bound.length > 0 ? ( +
已绑定 {bound.length} 条流量
+ ) : null}
); })} - {reproduction.provenance ?

来源：{reproduction.provenance}

: null} + {allBindings.some((item) => (item.step_index ?? 0) > steps.length) ? ( +
+
+ 超出步骤范围 + 这些绑定仍可编辑、解绑或改回总览/有效步骤。 +
+
+ {allBindings.filter((item) => (item.step_index ?? 0) > steps.length).map(bindingActions)} +
+
+ ) : null} + {reproduction.provenance ? ( +

来源：{reproduction.provenance}

+ ) : null} {pickerStep !== null ? ( ) : null} - setPreview(null)} /> + setPreview(null)} + /> + { + if (!open && !busy) setEditing(null); + }} + > + + + 编辑流量证据 + 修改用途、说明或所属复现步骤。 + + + + 用途 + + + + 证据说明 +