diff --git a/agent/finding_cases.go b/agent/finding_cases.go index 7e69412..01ec1c3 100644 --- a/agent/finding_cases.go +++ b/agent/finding_cases.go @@ -45,13 +45,13 @@ func findingCaseWorkflow(key string, tools []actool.CoreTool) ([]actool.CoreTool } return tools, ` -【同一漏洞的多次上报整理】保留本次独立报告，不覆盖其他原报告。默认报告完成即停止的旧指引在此扩展：先保存本次报告，然后 search_finding_duplicates(finding_id=独立记录ID) 分页读取同任务候选摘要，按需 get_finding_record 查看证据。仅凭漏洞类型/标题相似或高中低等级不能归并；须核对目标、接口、参数、权限边界、根因及实际验证证据。 -【判断准则】接口、路径、后端、读/写影响、评级不同，都不能单独作为不重复的理由；同一登录入口、同一令牌或同一漏洞类型也不能单独证明同一根因。比较实际缺失的鉴权检查、对象权限边界、会话门或共同组件，并引用每条记录中的具体证据。读和写可能是同一权限检查缺失的不同影响；也可能是独立 handler 的独立缺陷。只有证据证明共同实际根因才自动归并，证据证明不同根因才判独立；存在关联但无法确认时必须 suggest_finding_merge，说明已知共同点、缺失证据和待确认问题，不用“路径不同/影响不同”代替根因判断。明确同一缺陷时 merge_finding_records(task_id,finding_ids,title,reason)，不确定时 suggest_finding_merge 留待人工确认，不能强行归并。不重复上报、不删除记录、不重新探测目标。人工选择整理时只能判断所选记录，不能加入未选记录（既有组成员可作为报告上下文）。 -归并成功或已有case_id时，get_finding_case 记录实际version，分页读取全部成员，按需get_finding_record、get_finding_traffic读取原报告、PoC及已保存流量。update_finding_case_report保存完整统一报告（包含有内容的概述、影响、前置条件、复现步骤、证据、PoC、根因分析、修复建议章节，不得半篇截断）和基于已验证影响的severity及severity_reason；不能直接取成员最高等级。集中给出前置条件、完整执行顺序、输入输出如何衔接、可运行的完整脚本及命令，不把必要步骤分散到原报告；缺失条件、未验证步骤如实标明。版本冲突重新读取材料并生成，不能仅换version重试。各成员仍有独立报告、评级和证据。没有明确重复则保留独立记录。工具被关闭时不绕过配置。` +【同一漏洞的多次上报整理】保留本次独立报告，不覆盖其他原报告。默认报告完成即停止的旧指引在此扩展：先保存本次报告，然后 search_finding_duplicates(finding_id=独立记录ID) 分页读取同任务候选摘要（按 score 排序；strong=同漏洞类型且资产重叠，优先核对；weak=未同时具备同漏洞类型和资产重叠，结合 hit_reasons 核对，仅作参考），按需 get_finding_record 查看证据。仅凭漏洞类型、标题、令牌、登录入口相似或高中低等级不能归并；须核对目标、接口、参数、权限边界、根因及实际验证证据。只能处理同一任务，尊重人工否决和 finding_case_blocks。 +【判断准则】若证据表明端点/路由、处理器或后端组件、易受影响参数，或缺失检查属于不同代码路径，可以据此判定 independent；必须说明各自证据。仅凭路径、接口、登录入口、令牌、标题或漏洞类型相同/不同，都不能单独证明同一根因或不同根因。只有根因明确相同才直接调用 merge_finding_records(task_id,finding_ids,title,reason)。只有存在具体共同点（两条记录都证据化地指向同一端点、参数、组件或缺失检查），但缺少确认共同根因的关键证据时，才调用 suggest_finding_merge；一次提交覆盖2..20条同任务记录的一个分组，不要拆成两两建议。仅共享漏洞类型或资产没有具体共同点，不得提交建议；没有具体共同点时保持 independent，不要因为不确定就 suggest。不得重复上报、删除记录或重新探测目标。人工选择整理时只能判断所选记录，不能加入未选记录（既有组成员可作为报告上下文）。 +归并成功或已有case_id时，get_finding_case 记录实际version，分页读取全部成员，按需 get_finding_record、get_finding_traffic 读取原报告、PoC及已保存流量。根因明确时归并不被统一报告阻塞，应随后尽力用 update_finding_case_report 保存统一报告；如果本次无法完成，文件夹显示“统一报告待更新”，由后续 reporter 运行或人工补齐。统一报告在保存时仍须包含有内容的概述、影响、前置条件、复现步骤、证据、PoC、根因分析、修复建议章节，不得半篇截断；评级及 severity_reason 必须基于已验证影响，不能直接取成员最高等级。集中给出前置条件、完整执行顺序、输入输出如何衔接、可运行的完整脚本及命令，不把必要步骤分散到原报告；缺失条件、未验证步骤如实标明。版本冲突重新读取材料并生成，不能仅换version重试。各成员仍有独立报告、评级和证据。工具被关闭时不绕过配置。` } const historicalFindingReviewGuidance = ` 【历史整理模式，优先于旧报告撰写及自动绑流量指引】本次只整理已经存在的上报。原始报告即使为空或待更新，也只记录缺失情况，绝不调用 update_finding_report、report_finding、bind_finding_traffic，不修改原始标题、评级、状态、证据或报告，不执行 Bash、写文件或重新探测目标。 -先逐条 get_finding_record 读取所选记录，已有原报告直接作为判断材料复用，不重复撰写。只在所选范围内检查同一实际缺陷，可按需读取已存在的证据。按上述根因准则形成明确归并/独立/疑似建议；无法确认时保存建议供人工确认。 -已有或新建文件夹必须读取全部成员并写入完整、最新的统一报告；统一报告是单独的材料，不覆盖原报告。最后必须调用 complete_finding_case_review，为每条所选 finding_id 提交 verdict（independent/merged/suggested）、reason 和 root_cause_evidence，并提供总结。根因证据应说明具体相同/不同的检查及证据来源；疑似建议明确待补证据。服务端验收通过后才可宣布完成。工具不可用或验收失败时如实报告，不能只用自然语言声称完成。` +先逐条 get_finding_record 读取所选记录，已有原报告直接作为判断材料复用，不重复撰写。只在所选范围内检查同一实际缺陷，可按需读取已存在的证据。结构化证据显示不同端点/路由、处理器或后端组件、参数或缺失检查属于不同代码路径时，可判 independent；仅类型、资产、标题、令牌或登录入口相似不够。只有具体共同点明确且关键根因证据缺失时才保存一个覆盖2..20条记录的疑似分组，没有具体共同点就保持 independent。历史整理只能处理 selected scope，必须遵守同一任务限制和人工否决。 +历史模式禁止调用 update_finding_report、report_finding、bind_finding_traffic，不修改原始标题、评级、状态、证据或报告，不执行 Bash、写文件或重新探测目标。已有或新建文件夹应读取全部成员并尽力用 update_finding_case_report 写入完整、最新的统一报告；归并不因报告暂缺而阻塞，未完成时文件夹显示“统一报告待更新”，可由后续运行或人工补齐。统一报告是单独的材料，不覆盖原报告。最后必须调用 complete_finding_case_review，为每条所选 finding_id 提交 verdict（independent/merged/suggested）、reason 和 root_cause_evidence，并提供总结。根因证据应说明具体相同/不同的检查及证据来源；疑似建议明确待补证据。服务端验收通过后才可宣布完成。人工整理只能判断 selected scope，不能加入未选记录。工具不可用或验收失败时如实报告，不能只用自然语言声称完成。` diff --git a/agent/finding_cases_test.go b/agent/finding_cases_test.go index a8177ea..0120f3d 100644 --- a/agent/finding_cases_test.go +++ b/agent/finding_cases_test.go @@ -46,7 +46,7 @@ func TestHistoricalFindingReviewFiltersMutationAndProbeTools(t *testing.T) { t.Fatal("historical mutation capability", tool.Name()) } } - if len(out) != 6 || def.FindingGuidance != "" || !strings.Contains(def.CaseGuidance, "历史整理模式") || !strings.Contains(def.CaseGuidance, "接口、路径、后端、读/写影响、评级不同") { + if len(out) != 6 || def.FindingGuidance != "" || !strings.Contains(def.CaseGuidance, "历史整理模式") || !strings.Contains(def.CaseGuidance, "端点/路由") || !strings.Contains(def.CaseGuidance, "2..20") { t.Fatal("missing historical isolation", def) } normal, _, cl := AugmentTools(t.Context(), "reporter", tools) diff --git a/agent/goals.go b/agent/goals.go index 9bf9e17..2dc57cf 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -82,7 +82,7 @@ type GoalSpec struct { // shares the rate limiter, gets recorded by llmrec, and participates in LLM // failover instead of quietly bypassing all three. // -// desc is the task's free-text description (背景：靶标范围/flag 数量/交战说明等). +// desc is the task's free-text description (背景：靶标范围/数量要求/交战说明等). // It is fed alongside the goal so the decomposer no longer splits blind — the // prompt still forbids inventing anything the two texts don't state. // @@ -140,7 +140,7 @@ func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, } userMsg := "任务目标：\n" + goalText if d := strings.TrimSpace(desc); d != "" { - userMsg += "\n\n任务描述（背景信息，可能含靶标范围/flag 数量/交战说明；仅供参考，不要臆造其中未提及的内容）：\n" + d + userMsg += "\n\n任务描述（背景信息，可能含靶标范围/数量要求/交战说明；仅供参考，不要臆造其中未提及的内容）：\n" + d } // Use captureRun so every LLM step is emitted as an activity record (visible in // the plan tab under the round-0 marker). Falls back gracefully when emit is nil. diff --git a/agent/planner.go b/agent/planner.go index 3a13830..1010231 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -313,7 +313,7 @@ const plannerDefaultTmpl = `你是一个网络安全平台授权渗透测试系 - **要更深细节才按需调**：list_facts（分页，最新在前，默认 20，可 q 过滤、before 翻页，带 total/has_more）、list_findings（全部漏洞）、node_detail(id)（完整证据/详情；列表/recent_facts 只给摘要）、list_assets（pull：q 搜索、type/company_id/task_id 过滤、分页，或 id/ids 直取）、asset_neighbors。资产全局共享，别默认拉全量。 2. **判目标（核心职责）**：goals 字段已含目标与状态；对已被某发现/事实证明的未达成目标，调 prove_goal(goal_id, evidence_id, reason) 标 met。**当你标记的恰是最后一个未完成目标时，系统自动判定整个任务完成**——收官只由逐个 prove_goal 驱动，没有别的"一键完成"手段。 - - ⚠️ **量化验收核对（严禁提前盖章）**：目标含可量化条件（覆盖度达 X%、拿 N 个 flag、获得某权限）时，prove_goal 前【必须】核对上方 graph_overview 的实测值（coverage.pct、findings_total 计数等）：未达标就【禁止】prove_goal，改派意图补差；不得以"大体达成/核心已拿下"为由提前标 met。例：要求覆盖度 100% 而实测 coverage.pct=40% → 未达成，继续派补测意图。 + - ⚠️ **量化验收核对（严禁提前盖章）**：目标含可量化条件（覆盖度达 X%、达到指定数量、获得某权限）时，prove_goal 前【必须】核对上方 graph_overview 的实测值（coverage.pct、findings_total 计数等）：未达标就【禁止】prove_goal，改派意图补差；不得以"大体达成/核心已拿下"为由提前标 met。例：要求覆盖度 100% 而实测 coverage.pct=40% → 未达成，继续派补测意图。 3. **（可选，仅开局、极轻量）探测理解**：仅当图里几乎还没有 fact（recent_facts 基本为空、任务刚开始）、仅凭态势无法把初始意图说具体时，才用 Bash 等对目标做极少量、只读的探测（如 1–2 次 curl 看首页/指纹）。**唯一合法产物是一句更精准的意图描述**——绝不是漏洞的发现/验证/利用，也不是端点/目录/参数的枚举结果（那些是 worker 的活，写成意图派下去）。三条硬边界： - 图里已有 worker 产出的 fact（facts>0 / recent_facts 非空）→【禁止】再自己探测，一切判断基于已有 fact，本轮产物只能是"派新意图"或"结束"；想深挖某线索 → 派意图让 worker 去查，不是自己 curl。 diff --git a/agent/promptcatalog.go b/agent/promptcatalog.go index 1f0bd1b..67f1fa7 100644 --- a/agent/promptcatalog.go +++ b/agent/promptcatalog.go @@ -15,7 +15,7 @@ package agent const autoDefaultTmpl = `你是 **Auto**，这个渗透测试平台的「操作助手」。你不亲自渗透，而是**用工具操作平台**、按用户指令把事情办好。 你能做的（取决于给你开放了哪些工具）： -1. **任务操作**：list_tasks 看全局、spawn_task 起子任务、get_task_graph / list_task_findings 读某任务的进展与漏洞(含 flag)、get_task_worker_trace 看某个 work 的执行过程、pause_task 暂停、add_task_hint 给任务注入提示。 +1. **任务操作**：list_tasks 看全局、spawn_task 起子任务、get_task_graph / list_task_findings 读某任务的进展与漏洞、get_task_worker_trace 看某个 work 的执行过程、pause_task 暂停、add_task_hint 给任务注入提示。 2. **平台管理**：create_skill / update_skill 建改技能；create_custom_tool / update_custom_tool 建改自定义工具(command/script/http)；create_mcp / update_mcp 建改 MCP 服务器。 原则： diff --git a/agent/tools.go b/agent/tools.go index adb09c4..52774b2 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -1208,7 +1208,7 @@ func (t *ToolSet) proveGoal() actool.CoreTool { } func (t *ToolSet) goalMet() actool.CoreTool { - return writeTool("goal_met", "【立即结束整个任务】——仅当你确认任务的【全部目标都已真正达成、整体收官】时才调（注意是任务【整体】完成；仅仅达成了其中某一个目标/某一个 flag/某一个漏洞【不算】——那种情况用 prove_goal 标记该目标即可）。⚠️它不是用来“结束本轮规划”的：本轮没有新意图要派、或在等 worker 产出，都【直接结束本轮即可，不要调本工具】（0 个意图是完全正常的）。正常判定优先用 prove_goal 逐个证明目标；goal_met 只是绕过逐个证明、直接从全局收官的手段。", + return writeTool("goal_met", "【立即结束整个任务】——仅当你确认任务的【全部目标都已真正达成、整体收官】时才调（注意是任务【整体】完成；仅仅达成了其中某一个目标/某一个漏洞【不算】——那种情况用 prove_goal 标记该目标即可）。⚠️它不是用来“结束本轮规划”的：本轮没有新意图要派、或在等 worker 产出，都【直接结束本轮即可，不要调本工具】（0 个意图是完全正常的）。正常判定优先用 prove_goal 逐个证明目标；goal_met 只是绕过逐个证明、直接从全局收官的手段。", obj(map[string]any{"reason": str("达成理由（必须是目标真正达成的证据，不能是“本轮无新方向”这类结束本轮的理由）")}, "reason"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct{ Reason string } diff --git a/db/finding_cases.go b/db/finding_cases.go index d7ec976..e469c54 100644 --- a/db/finding_cases.go +++ b/db/finding_cases.go @@ -41,13 +41,14 @@ type FindingCaseEvent struct { CreatedAt time.Time `json:"created_at"` } type FindingCaseSuggestion struct { - ID int64 `json:"id,string"` - TaskID int64 `json:"task_id,string"` - LeftID int64 `json:"left_id,string"` - RightID int64 `json:"right_id,string"` - Title string `json:"title"` - Reason string `json:"reason"` - State string `json:"state"` + ID int64 `json:"id,string"` + TaskID int64 `json:"task_id,string"` + LeftID int64 `json:"left_id,string"` + RightID int64 `json:"right_id,string"` + Title string `json:"title"` + Reason string `json:"reason"` + State string `json:"state"` + FindingIDs []string `json:"finding_ids"` } type FindingCaseRow struct { Case *FindingCase `json:"case,omitempty"` @@ -198,9 +199,9 @@ func lockCaseTask(tx *sql.Tx, taskID int64) error { return tx.QueryRow(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&id) } func (d *DB) MergeFindingCase(ctx context.Context, taskID int64, ids []int64, title, reason, actor string) (int64, error) { - return d.mergeFindingCase(ctx, taskID, ids, title, reason, actor, 0) + return d.mergeFindingCase(ctx, taskID, ids, title, reason, actor, 0, nil) } -func (d *DB) mergeFindingCase(ctx context.Context, taskID int64, ids []int64, title, reason, actor string, suggestionID int64) (int64, error) { +func (d *DB) mergeFindingCase(ctx context.Context, taskID int64, ids []int64, title, reason, actor string, suggestionID int64, dropped []int64) (int64, error) { ids, err := sortedFindingIDs(ids) if err != nil { return 0, err @@ -208,6 +209,7 @@ func (d *DB) mergeFindingCase(ctx context.Context, taskID int64, ids []int64, ti if strings.TrimSpace(title) == "" || strings.TrimSpace(reason) == "" { return 0, errors.New("必须提供漏洞名称及基于证据的归并理由") } + keptIDs := append([]int64(nil), ids...) var cid int64 err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { if err := lockCaseTask(tx, taskID); err != nil { @@ -308,11 +310,69 @@ func (d *DB) mergeFindingCase(ctx context.Context, taskID int64, ids []int64, ti return err } } - _, err := tx.Exec(`UPDATE finding_case_suggestions SET state='accepted' WHERE state='pending' AND left_id IN(SELECT value::bigint FROM jsonb_array_elements_text($1::jsonb)) AND right_id IN(SELECT value::bigint FROM jsonb_array_elements_text($1::jsonb))`, string(raw)) - return err + for _, excluded := range dropped { + if all[excluded] { + return fmt.Errorf("排除的记录 #%d 已与所选记录在同一文件夹，请先移出该文件夹后再确认", excluded) + } + for _, kept := range keptIDs { + if _, err := tx.Exec(`INSERT INTO finding_case_blocks(left_id,right_id,reason) + VALUES(LEAST($1::bigint,$2::bigint),GREATEST($1::bigint,$2::bigint),'人工确认分组时排除') + ON CONFLICT DO NOTHING`, excluded, kept); err != nil { + return err + } + } + } + if suggestionID > 0 { + if _, err := tx.Exec(`UPDATE finding_case_suggestions SET state='accepted' WHERE id=$1 AND state='pending'`, suggestionID); err != nil { + return err + } + } + return sweepFindingCaseSuggestions(tx, taskID) }) return cid, err } + +// sweepFindingCaseSuggestions reconciles pending cards after a merge or veto. +// Cross-task member rows are ignored so a malformed legacy row cannot leak into +// a task's suggestions. +func sweepFindingCaseSuggestions(tx *sql.Tx, taskID int64) error { + if _, err := tx.Exec(`UPDATE finding_case_suggestions s SET state='accepted' + WHERE s.task_id=$1 AND s.state='pending' + AND (SELECT count(*) FROM finding_case_suggestion_members sm + WHERE sm.suggestion_id=s.id + AND EXISTS (SELECT 1 FROM findings f WHERE f.id=sm.finding_id AND f.task_id=s.task_id)) >= 2 + AND (SELECT count(*) FROM finding_case_suggestion_members sm + JOIN findings f ON f.id=sm.finding_id AND f.task_id=s.task_id + JOIN finding_case_members cm ON cm.finding_id=sm.finding_id + JOIN finding_cases c ON c.id=cm.case_id AND c.active + WHERE sm.suggestion_id=s.id) >= 2 + AND (SELECT count(DISTINCT cm.case_id) FROM finding_case_suggestion_members sm + JOIN findings f ON f.id=sm.finding_id AND f.task_id=s.task_id + JOIN finding_case_members cm ON cm.finding_id=sm.finding_id + JOIN finding_cases c ON c.id=cm.case_id AND c.active + WHERE sm.suggestion_id=s.id) = 1 + AND NOT EXISTS (SELECT 1 FROM finding_case_suggestion_members sm + JOIN findings f ON f.id=sm.finding_id AND f.task_id=s.task_id + LEFT JOIN finding_case_members cm ON cm.finding_id=sm.finding_id + LEFT JOIN finding_cases c ON c.id=cm.case_id AND c.active + WHERE sm.suggestion_id=s.id AND c.id IS NULL)`, taskID); err != nil { + return err + } + _, err := tx.Exec(`UPDATE finding_case_suggestions s SET state='obsolete' + WHERE s.task_id=$1 AND s.state='pending' AND ( + (SELECT count(*) FROM finding_case_suggestion_members sm + JOIN findings f ON f.id=sm.finding_id AND f.task_id=s.task_id + WHERE sm.suggestion_id=s.id) < 2 + OR EXISTS (SELECT 1 + FROM finding_case_suggestion_members a + JOIN finding_case_suggestion_members b + ON b.suggestion_id=a.suggestion_id AND a.finding_id < b.finding_id + JOIN findings fa ON fa.id=a.finding_id AND fa.task_id=s.task_id + JOIN findings fb ON fb.id=b.finding_id AND fb.task_id=s.task_id + JOIN finding_case_blocks bl ON bl.left_id=a.finding_id AND bl.right_id=b.finding_id + WHERE a.suggestion_id=s.id))`, taskID) + return err +} func (d *DB) RemoveFindingCaseMember(ctx context.Context, cid, fid int64, reason string) error { c, err := d.GetFindingCase(cid) if err != nil { @@ -355,8 +415,10 @@ func (d *DB) RemoveFindingCaseMember(ctx context.Context, cid, fid int64, reason if _, err := tx.Exec(`INSERT INTO finding_case_events(case_id,action,actor,reason,finding_ids) VALUES($1,'detach','human',$2,jsonb_build_array($3::bigint))`, cid, reason, fid); err != nil { return err } - _, err := tx.Exec(`DELETE FROM finding_case_members WHERE case_id=$1 AND finding_id=$2`, cid, fid) - return err + if _, err := tx.Exec(`DELETE FROM finding_case_members WHERE case_id=$1 AND finding_id=$2`, cid, fid); err != nil { + return err + } + return sweepFindingCaseSuggestions(tx, origin) }) } func (d *DB) UpdateFindingCaseReport(ctx context.Context, cid, version int64, title, report, severity, reason string) error { @@ -382,81 +444,345 @@ func (d *DB) UpdateFindingCaseReport(ctx context.Context, cid, version int64, ti return err }) } -func (d *DB) SuggestFindingCase(ctx context.Context, taskID, left, right int64, title, reason string) (int64, error) { + +type suggestionQuery interface { + Query(string, ...any) (*sql.Rows, error) +} + +// suggestionMembers reads only surviving members in the suggestion's task. The +// anchor fallback keeps rows created by older binaries resolvable if a database +// was upgraded before the idempotent backfill ran. +func suggestionMembers(q suggestionQuery, suggestionID, taskID int64) ([]int64, error) { + rows, err := q.Query(`SELECT sm.finding_id + FROM finding_case_suggestion_members sm + JOIN findings f ON f.id=sm.finding_id AND f.task_id=$2 + WHERE sm.suggestion_id=$1 ORDER BY sm.finding_id`, suggestionID, taskID) + if err != nil { + return nil, err + } + ids := []int64{} + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + rows.Close() + return nil, err + } + ids = append(ids, id) + } + err = rows.Err() + rows.Close() + if err != nil { + return nil, err + } + if len(ids) > 0 { + return ids, nil + } + rows, err = q.Query(`SELECT x.finding_id FROM ( + SELECT s.left_id AS finding_id FROM finding_case_suggestions s WHERE s.id=$1 + UNION SELECT s.right_id FROM finding_case_suggestions s WHERE s.id=$1 + ) x JOIN findings f ON f.id=x.finding_id AND f.task_id=$2 ORDER BY x.finding_id`, suggestionID, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + ids = append(ids, id) + } + return ids, rows.Err() +} + +// SuggestFindingCase stores a pending group of 2..20 same-task findings. It +// refreshes an exact group, reuses a pending superset, and supersedes strict +// subsets while holding the task advisory lock. +func (d *DB) SuggestFindingCase(ctx context.Context, taskID int64, findingIDs []int64, title, reason string) (int64, error) { if strings.TrimSpace(title) == "" || strings.TrimSpace(reason) == "" { return 0, errors.New("建议需要名称和具体证据理由") } - ids, err := sortedFindingIDs([]int64{left, right}) + if len(findingIDs) > 20 { + return 0, errors.New("疑似重复分组最多20条记录") + } + seen := make(map[int64]bool, len(findingIDs)) + for _, fid := range findingIDs { + if seen[fid] { + return 0, errors.New("疑似重复分组中的记录编号必须唯一") + } + seen[fid] = true + } + ids, err := sortedFindingIDs(findingIDs) if err != nil { return 0, err } - var id int64 + var suggestionID int64 err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { if err := lockCaseTask(tx, taskID); err != nil { return err } + caseIDs := map[int64]bool{} + sameActiveCase := true + var firstCase int64 for _, fid := range ids { var tid sql.NullInt64 - if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, fid).Scan(&tid); err != nil { + var cid int64 + if err := tx.QueryRow(`SELECT f.task_id,COALESCE(CASE WHEN c.active THEN m.case_id ELSE 0 END,0) + FROM findings f LEFT JOIN finding_case_members m ON m.finding_id=f.id + LEFT JOIN finding_cases c ON c.id=m.case_id WHERE f.id=$1`, fid).Scan(&tid, &cid); err != nil { return err } if !tid.Valid || tid.Int64 != taskID { return errors.New("只能建议同一任务的归并") } + if cid == 0 { + sameActiveCase = false + } else if firstCase == 0 { + firstCase = cid + caseIDs[cid] = true + } else if firstCase != cid { + sameActiveCase = false + } } - var blocked bool - if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM finding_case_blocks WHERE left_id=$1 AND right_id=$2)`, ids[0], ids[1]).Scan(&blocked); err != nil { + if sameActiveCase && firstCase > 0 && len(caseIDs) == 1 { + return errors.New("已在同一文件夹") + } + + raw, _ := json.Marshal(ids) + rows, err := tx.Query(`SELECT left_id,right_id FROM finding_case_blocks + WHERE left_id IN (SELECT value::bigint FROM jsonb_array_elements_text($1::jsonb)) + AND right_id IN (SELECT value::bigint FROM jsonb_array_elements_text($1::jsonb)) + ORDER BY left_id,right_id`, string(raw)) + if err != nil { return err } - if blocked { - return errors.New("人工已否决该归并") + blocked := []string{} + for rows.Next() { + var left, right int64 + if err := rows.Scan(&left, &right); err != nil { + rows.Close() + return err + } + blocked = append(blocked, fmt.Sprintf("#%d/#%d", left, right)) } - return tx.QueryRow(`INSERT INTO finding_case_suggestions(task_id,left_id,right_id,title,reason) VALUES($1,$2,$3,$4,$5) ON CONFLICT(left_id,right_id) DO UPDATE SET reason=EXCLUDED.reason WHERE finding_case_suggestions.state='pending' RETURNING id`, taskID, ids[0], ids[1], title, reason).Scan(&id) + err = rows.Err() + rows.Close() + if err != nil { + return err + } + if len(blocked) > 0 { + return fmt.Errorf("人工已否决这些记录的归并：%s；请去掉其中一方后重新提交", strings.Join(blocked, "、")) + } + + type pendingGroup struct { + id int64 + members []int64 + title string + reason string + } + rows, err = tx.Query(`SELECT id,title,reason FROM finding_case_suggestions + WHERE task_id=$1 AND state='pending' ORDER BY id`, taskID) + if err != nil { + return err + } + groups := []pendingGroup{} + for rows.Next() { + var g pendingGroup + if err := rows.Scan(&g.id, &g.title, &g.reason); err != nil { + rows.Close() + return err + } + groups = append(groups, g) + } + err = rows.Err() + rows.Close() + if err != nil { + return err + } + for i := range groups { + groups[i].members, err = suggestionMembers(tx, groups[i].id, taskID) + if err != nil { + return err + } + } + want := map[int64]bool{} + for _, fid := range ids { + want[fid] = true + } + contains := func(group []int64, set map[int64]bool) bool { + for _, fid := range group { + if !set[fid] { + return false + } + } + return true + } + exact, superset := int64(0), int64(0) + supersede := []int64{} + for _, g := range groups { + gm := map[int64]bool{} + for _, fid := range g.members { + gm[fid] = true + } + switch { + case len(g.members) == len(ids) && contains(ids, gm) && contains(g.members, want): + if exact == 0 { + exact = g.id + } else { + supersede = append(supersede, g.id) + } + case len(g.members) > len(ids) && contains(ids, gm): + if superset == 0 { + superset = g.id + } + case len(g.members) < len(ids) && contains(g.members, want): + supersede = append(supersede, g.id) + } + } + for _, sid := range supersede { + if _, err := tx.Exec(`UPDATE finding_case_suggestions SET state='superseded' WHERE id=$1 AND state='pending'`, sid); err != nil { + return err + } + } + if exact > 0 { + suggestionID = exact + _, err = tx.Exec(`UPDATE finding_case_suggestions SET title=$2,reason=$3 WHERE id=$1`, exact, title, reason) + return err + } + if superset > 0 { + suggestionID = superset + return nil + } + if err := tx.QueryRow(`INSERT INTO finding_case_suggestions(task_id,left_id,right_id,title,reason) + VALUES($1,$2,$3,$4,$5) RETURNING id`, taskID, ids[0], ids[1], title, reason).Scan(&suggestionID); err != nil { + return err + } + for _, fid := range ids { + if _, err := tx.Exec(`INSERT INTO finding_case_suggestion_members(suggestion_id,finding_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, suggestionID, fid); err != nil { + return err + } + } + return nil }) - return id, err + return suggestionID, err } + func (d *DB) FindingCaseSuggestions(taskID int64) ([]FindingCaseSuggestion, error) { - rows, err := d.Query(`SELECT id,task_id,left_id,right_id,title,reason,state FROM finding_case_suggestions WHERE state='pending' AND ($1::bigint=0 OR task_id=$1) ORDER BY id DESC`, taskID) + rows, err := d.Query(`SELECT id,task_id,left_id,right_id,title,reason,state + FROM finding_case_suggestions WHERE state='pending' AND ($1::bigint=0 OR task_id=$1) ORDER BY id DESC`, taskID) if err != nil { return nil, err } - defer rows.Close() - out := []FindingCaseSuggestion{} + items := []FindingCaseSuggestion{} for rows.Next() { var s FindingCaseSuggestion if err := rows.Scan(&s.ID, &s.TaskID, &s.LeftID, &s.RightID, &s.Title, &s.Reason, &s.State); err != nil { + rows.Close() return nil, err } + items = append(items, s) + } + err = rows.Err() + rows.Close() + if err != nil { + return nil, err + } + out := []FindingCaseSuggestion{} + for _, s := range items { + ids, err := suggestionMembers(d, s.ID, s.TaskID) + if err != nil { + return nil, err + } + if len(ids) < 2 { + continue + } + s.LeftID, s.RightID = ids[0], ids[1] + for _, fid := range ids { + s.FindingIDs = append(s.FindingIDs, strconv.FormatInt(fid, 10)) + } out = append(out, s) } - return out, rows.Err() + return out, nil } -func (d *DB) ResolveFindingCaseSuggestion(ctx context.Context, id int64, accept bool) (int64, error) { + +func (d *DB) ResolveFindingCaseSuggestion(ctx context.Context, id int64, accept bool, exclude []int64) (int64, error) { var s FindingCaseSuggestion - if err := d.QueryRow(`SELECT task_id,left_id,right_id,title,reason,state FROM finding_case_suggestions WHERE id=$1`, id).Scan(&s.TaskID, &s.LeftID, &s.RightID, &s.Title, &s.Reason, &s.State); err != nil { + if err := d.QueryRow(`SELECT task_id,title,reason,state FROM finding_case_suggestions WHERE id=$1`, id).Scan(&s.TaskID, &s.Title, &s.Reason, &s.State); err != nil { return 0, err } if s.State != "pending" { return 0, ErrFindingCaseConflict } + members, err := suggestionMembers(d, id, s.TaskID) + if err != nil { + return 0, err + } + if len(members) < 2 { + return 0, errors.New("疑似重复分组成员不足") + } + memberSet := map[int64]bool{} + for _, fid := range members { + memberSet[fid] = true + } + excludeSet := map[int64]bool{} + for _, fid := range exclude { + if !memberSet[fid] { + return 0, fmt.Errorf("排除记录 #%d 不属于该分组", fid) + } + excludeSet[fid] = true + } + kept := []int64{} + for _, fid := range members { + if !excludeSet[fid] { + kept = append(kept, fid) + } + } if accept { - return d.mergeFindingCase(ctx, s.TaskID, []int64{s.LeftID, s.RightID}, s.Title, s.Reason, "human", id) + if len(kept) < 2 { + return 0, errors.New("至少保留两条记录") + } + return d.mergeFindingCase(ctx, s.TaskID, kept, s.Title, s.Reason, "human", id, exclude) + } + if len(exclude) > 0 { + return 0, errors.New("否决分组时不能排除记录") } return 0, d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { if err := lockCaseTask(tx, s.TaskID); err != nil { return err } - res, err := tx.Exec(`UPDATE finding_case_suggestions SET state='rejected' WHERE id=$1 AND state='pending'`, id) + var state string + if err := tx.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1 FOR UPDATE`, id).Scan(&state); err != nil { + return err + } + if state != "pending" { + return ErrFindingCaseConflict + } + members, err := suggestionMembers(tx, id, s.TaskID) if err != nil { return err } - if n, _ := res.RowsAffected(); n != 1 { - return ErrFindingCaseConflict + for i := 0; i < len(members); i++ { + for j := i + 1; j < len(members); j++ { + var leftCase, rightCase sql.NullInt64 + if err := tx.QueryRow(`SELECT m.case_id FROM finding_case_members m JOIN finding_cases c ON c.id=m.case_id AND c.active WHERE m.finding_id=$1`, members[i]).Scan(&leftCase); err != nil && !errors.Is(err, sql.ErrNoRows) { + return err + } + if err := tx.QueryRow(`SELECT m.case_id FROM finding_case_members m JOIN finding_cases c ON c.id=m.case_id AND c.active WHERE m.finding_id=$1`, members[j]).Scan(&rightCase); err != nil && !errors.Is(err, sql.ErrNoRows) { + return err + } + if leftCase.Valid && rightCase.Valid && leftCase.Int64 == rightCase.Int64 { + continue + } + if _, err := tx.Exec(`INSERT INTO finding_case_blocks(left_id,right_id,reason) + VALUES(LEAST($1::bigint,$2::bigint),GREATEST($1::bigint,$2::bigint),'人工否决') ON CONFLICT DO NOTHING`, members[i], members[j]); err != nil { + return err + } + } } - if _, err := tx.Exec(`INSERT INTO finding_case_blocks(left_id,right_id,reason) VALUES($1,$2,'人工否决') ON CONFLICT DO NOTHING`, s.LeftID, s.RightID); err != nil { + if _, err := tx.Exec(`UPDATE finding_case_suggestions SET state='rejected' WHERE id=$1`, id); err != nil { return err } - return nil + return sweepFindingCaseSuggestions(tx, s.TaskID) }) } @@ -479,13 +805,13 @@ func (d *DB) DistinctFindingStats(taskID string) (*FindingDistinctStats, error) } // Compact candidate pages exclude PoC/report bodies and do not treat severity as identity. -func (d *DB) FindingCaseCandidates(fid int64, page, size int, scope ...[]int64) ([]map[string]any, int, error) { +func (d *DB) FindingCaseCandidates(fid int64, page, size int, scope ...[]int64) ([]map[string]any, int, int, error) { f, err := d.FindingCaseSummary(fid) if err != nil { - return nil, 0, err + return nil, 0, 0, err } if f == nil || f.TaskID == nil { - return nil, 0, errors.New("漏洞任务不可用") + return nil, 0, 0, errors.New("漏洞任务不可用") } if page < 1 { page = 1 @@ -505,32 +831,58 @@ func (d *DB) FindingCaseCandidates(fid int64, page, size int, scope ...[]int64) if len(scope) > 0 { raw, err := json.Marshal(scope[0]) if err != nil { - return nil, 0, err + return nil, 0, 0, err } where += ` AND f.id IN (SELECT jsonb_array_elements_text($6::jsonb)::bigint)` args = append(args, string(raw)) } - var total int - if err := d.QueryRow(`SELECT count(*) FROM findings f`+where, args...).Scan(&total); err != nil { - return nil, 0, err + classHit := `(trim(f.vulnclass)<>'' AND lower(f.vulnclass)=lower($3))` + assetHit := `EXISTS(SELECT 1 FROM jsonb_array_elements(COALESCE(NULLIF(f.asset_ids,'null'::jsonb),'[]'::jsonb)) a WHERE a.value IN (SELECT value FROM jsonb_array_elements($5::jsonb)))` + nameHit := `(trim(f.name)<>'' AND lower(f.name)=lower($4))` + scoreExpr := fmt.Sprintf(`(CASE WHEN %s THEN 2 ELSE 0 END + CASE WHEN %s THEN 2 ELSE 0 END + CASE WHEN %s THEN 1 ELSE 0 END + CASE WHEN (%s AND %s) THEN 3 ELSE 0 END)`, classHit, assetHit, nameHit, classHit, assetHit) + strongExpr := fmt.Sprintf(`(%s AND %s)`, classHit, assetHit) + base := `WITH candidates AS (SELECT f.id,f.name,f.vulnclass,f.severity,left(f.summary,400) AS summary,f.asset_ids, + COALESCE(m.case_id,0) AS case_id, + ` + scoreExpr + ` AS score, + ` + strongExpr + ` AS strong, + ARRAY_REMOVE(ARRAY[ + CASE WHEN ` + classHit + ` THEN '同漏洞类型' END, + CASE WHEN ` + assetHit + ` THEN '资产重叠' END, + CASE WHEN ` + nameHit + ` THEN '同名称' END + ]::text[], NULL) AS hit_reasons + FROM findings f LEFT JOIN finding_case_members m ON m.finding_id=f.id` + where + `)` + var total, strongTotal int + if err := d.QueryRow(base+` SELECT count(*),count(*) FILTER(WHERE strong) FROM candidates`, args...).Scan(&total, &strongTotal); err != nil { + return nil, 0, 0, err } args = append(args, size, (page-1)*size) - rows, err := d.Query(`SELECT f.id,f.name,f.vulnclass,f.severity,left(f.summary,400),f.asset_ids,COALESCE(m.case_id,0) FROM findings f LEFT JOIN finding_case_members m ON m.finding_id=f.id`+where+fmt.Sprintf(` ORDER BY f.id DESC LIMIT $%d OFFSET $%d`, len(args)-1, len(args)), args...) + rows, err := d.Query(base+fmt.Sprintf(` SELECT id,name,vulnclass,severity,summary,asset_ids,case_id,score,strong,to_json(hit_reasons) + FROM candidates ORDER BY score DESC,id DESC LIMIT $%d OFFSET $%d`, len(args)-1, len(args)), args...) if err != nil { - return nil, 0, err + return nil, 0, 0, err } defer rows.Close() out := []map[string]any{} for rows.Next() { - var id, cid int64 + var id, cid, score int64 var name, class, severity, summary string var assetIDs json.RawMessage - if err := rows.Scan(&id, &name, &class, &severity, &summary, &assetIDs, &cid); err != nil { - return nil, 0, err + var strong bool + var hitReasonsRaw json.RawMessage + if err := rows.Scan(&id, &name, &class, &severity, &summary, &assetIDs, &cid, &score, &strong, &hitReasonsRaw); err != nil { + return nil, 0, 0, err } - out = append(out, map[string]any{"finding_id": strconv.FormatInt(id, 10), "name": name, "vulnclass": class, "severity": severity, "summary": summary, "asset_ids": assetIDs, "case_id": strconv.FormatInt(cid, 10)}) + var hitReasons []string + if err := json.Unmarshal(hitReasonsRaw, &hitReasons); err != nil { + return nil, 0, 0, err + } + tier := "weak" + if strong { + tier = "strong" + } + out = append(out, map[string]any{"finding_id": strconv.FormatInt(id, 10), "name": name, "vulnclass": class, "severity": severity, "summary": summary, "asset_ids": assetIDs, "case_id": strconv.FormatInt(cid, 10), "score": int(score), "hit_reasons": hitReasons, "tier": tier}) } - return out, total, rows.Err() + return out, total, strongTotal, rows.Err() } // Batch task-list counters avoid an additional database round trip per task. diff --git a/db/finding_cases_test.go b/db/finding_cases_test.go index 48398a5..f54c337 100644 --- a/db/finding_cases_test.go +++ b/db/finding_cases_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "reflect" + "strings" "sync" "testing" ) @@ -118,11 +119,14 @@ func TestFindingCaseScopeSuggestionsAndDeletion(t *testing.T) { if _, err := d.MergeFindingCase(ctx, task.ID, []int64{ids[0], fid}, "x", "x", "reporter"); err == nil { t.Fatal("cross task accepted") } - sid, err := d.SuggestFindingCase(ctx, task.ID, ids[0], ids[1], "order", "uncertain") + if _, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], fid}, "x", "x"); err == nil { + t.Fatal("cross task suggestion accepted") + } + sid, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1]}, "order", "uncertain") if err != nil { t.Fatal(err) } - if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, false); err != nil { + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, false, nil); err != nil { t.Fatal(err) } if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "x", "x", "reporter"); err == nil { @@ -233,7 +237,7 @@ func TestFindingCaseTaskDeletionAndRepeatedDetach(t *testing.T) { func TestFindingCaseSuggestionResolutionIsAtomic(t *testing.T) { d, task, ids := caseTestDB(t) ctx := context.Background() - sid, err := d.SuggestFindingCase(ctx, task.ID, ids[0], ids[1], "orders", "uncertain") + sid, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1]}, "orders", "uncertain") if err != nil { t.Fatal(err) } @@ -241,7 +245,7 @@ func TestFindingCaseSuggestionResolutionIsAtomic(t *testing.T) { var wg sync.WaitGroup for _, accept := range []bool{true, false} { wg.Add(1) - go func(a bool) { defer wg.Done(); _, e := d.ResolveFindingCaseSuggestion(ctx, sid, a); results <- e }(accept) + go func(a bool) { defer wg.Done(); _, e := d.ResolveFindingCaseSuggestion(ctx, sid, a, nil); results <- e }(accept) } wg.Wait() close(results) @@ -269,7 +273,7 @@ func TestFindingCaseCandidatesAreCompactAndRelated(t *testing.T) { if _, err := d.Exec(`UPDATE findings SET vulnclass='SQLi',name='search injection' WHERE id=$1`, ids[3]); err != nil { t.Fatal(err) } - items, total, err := d.FindingCaseCandidates(ids[0], 1, 1) + items, total, _, err := d.FindingCaseCandidates(ids[0], 1, 1) if err != nil || total != 2 || len(items) != 1 { t.Fatalf("candidate pagination %d %+v %v", total, items, err) } @@ -286,7 +290,7 @@ func TestFindingCaseCandidatesAreCompactAndRelated(t *testing.T) { if err := d.RemoveFindingCaseMember(context.Background(), cid, ids[1], "different cause"); err != nil { t.Fatal(err) } - _, total, err = d.FindingCaseCandidates(ids[0], 1, 20) + _, total, _, err = d.FindingCaseCandidates(ids[0], 1, 20) if err != nil || total != 1 { t.Fatalf("veto excluded %d %v", total, err) } @@ -355,3 +359,242 @@ func TestFindingCaseViewContextAndAssetFiltering(t *testing.T) { t.Fatalf("original severity order %+v %v", rows, err) } } + +func TestFindingCaseSuggestionClusterLifecycle(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same handler evidence") + if err != nil { + t.Fatal(err) + } + items, err := d.FindingCaseSuggestions(task.ID) + if err != nil || len(items) != 1 || len(items[0].FindingIDs) != 3 { + t.Fatalf("cluster listing %+v %v", items, err) + } + if reused, err := d.SuggestFindingCase(ctx, task.ID, ids[:2], "weaker", "subset"); err != nil || reused != sid { + t.Fatalf("pending superset was not reused: %d %v", reused, err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, nil); err != nil { + t.Fatal(err) + } + var state string + if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, sid).Scan(&state); err != nil || state != "accepted" { + t.Fatalf("cluster state %q %v", state, err) + } + for _, fid := range ids[:3] { + if cid, err := d.FindingCaseID(fid); err != nil || cid == 0 { + t.Fatalf("member %d not merged: %d %v", fid, cid, err) + } + } +} + +func TestFindingCaseSuggestionSweepScopesEachCard(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + first, err := d.SuggestFindingCase(ctx, task.ID, ids[:2], "first", "first group") + if err != nil { + t.Fatal(err) + } + second, err := d.SuggestFindingCase(ctx, task.ID, ids[2:], "second", "second group") + if err != nil { + t.Fatal(err) + } + if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "first", "same root cause", "human"); err != nil { + t.Fatal(err) + } + var firstState, secondState string + if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, first).Scan(&firstState); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, second).Scan(&secondState); err != nil { + t.Fatal(err) + } + if firstState != "accepted" || secondState != "pending" { + t.Fatalf("suggestion states first=%q second=%q", firstState, secondState) + } +} + +func TestFindingCaseSuggestionSubsetSuperseded(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + subset, err := d.SuggestFindingCase(ctx, task.ID, ids[:2], "orders", "partial evidence") + if err != nil { + t.Fatal(err) + } + superset, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "complete group evidence") + if err != nil { + t.Fatal(err) + } + if subset == superset { + t.Fatal("subset suggestion was reused instead of superseded") + } + var state string + if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, subset).Scan(&state); err != nil || state != "superseded" { + t.Fatalf("subset state %q %v", state, err) + } +} + +func TestFindingCaseSuggestionAcceptExclusionAndReject(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same component") + if err != nil { + t.Fatal(err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, []int64{ids[2]}); err != nil { + t.Fatal(err) + } + var blocks int + if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=$1 AND right_id=$2`, ids[0], ids[2]).Scan(&blocks); err != nil || blocks != 1 { + t.Fatalf("excluded member was not vetoed: %d %v", blocks, err) + } + if cid, _ := d.FindingCaseID(ids[2]); cid != 0 { + t.Fatalf("excluded member entered case %d", cid) + } +} + +func TestFindingCaseSuggestionExclusionBlocksOnlyKeptSuggestionMembers(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "existing folder", "same root cause", "human"); err != nil { + t.Fatal(err) + } + sid, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[2], ids[3]}, "orders", "same component") + if err != nil { + t.Fatal(err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, []int64{ids[3]}); err != nil { + t.Fatal(err) + } + var folderMateBlock, keptBlock int + if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=$1 AND right_id=$2`, ids[1], ids[3]).Scan(&folderMateBlock); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=$1 AND right_id=$2`, ids[0], ids[3]).Scan(&keptBlock); err != nil { + t.Fatal(err) + } + if folderMateBlock != 0 || keptBlock != 1 { + t.Fatalf("exclusion blocks folder mate=%d kept=%d", folderMateBlock, keptBlock) + } +} + +func TestFindingCaseSuggestionRejectVetoesAllPairs(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same component") + if err != nil { + t.Fatal(err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, false, nil); err != nil { + t.Fatal(err) + } + var count int + if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=ANY($1::bigint[]) AND right_id=ANY($1::bigint[])`, ids[:3]).Scan(&count); err != nil || count != 3 { + t.Fatalf("reject did not veto all pairs: %d %v", count, err) + } +} + +func TestFindingCaseSuggestionRejectSkipsExistingFolderAndBlockedSuggest(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "orders", "same handler", "human"); err != nil { + t.Fatal(err) + } + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same handler") + if err != nil { + t.Fatal(err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, false, nil); err != nil { + t.Fatal(err) + } + var blocked int + if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=$1 AND right_id=$2`, ids[0], ids[1]).Scan(&blocked); err != nil || blocked != 0 { + t.Fatalf("same-folder pair was vetoed: %d %v", blocked, err) + } + if _, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[2]}, "orders", "retry"); err == nil || !strings.Contains(err.Error(), fmt.Sprintf("#%d/#%d", ids[0], ids[2])) { + t.Fatalf("blocked pair was accepted or unnamed: %v", err) + } +} + +func TestFindingCaseSuggestionExclusionCannotPullExistingFolder(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "orders", "same handler", "human"); err != nil { + t.Fatal(err) + } + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same handler") + if err != nil { + t.Fatal(err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, []int64{ids[1]}); err == nil { + t.Fatal("excluded member was pulled through its existing folder") + } + if cid, _ := d.FindingCaseID(ids[2]); cid != 0 { + t.Fatalf("failed exclusion merged remaining record into case %d", cid) + } + var state string + if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, sid).Scan(&state); err != nil || state != "pending" { + t.Fatalf("failed exclusion changed suggestion state: %q %v", state, err) + } +} + +func TestFindingCaseSuggestionAnchorDeletionAndLegacyPair(t *testing.T) { + d, task, ids := caseTestDB(t) + ctx := context.Background() + sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "same handler") + if err != nil { + t.Fatal(err) + } + if _, err := d.DeleteFinding(ids[0]); err != nil { + t.Fatal(err) + } + items, err := d.FindingCaseSuggestions(task.ID) + if err != nil || len(items) != 1 || items[0].ID != sid || items[0].LeftID != ids[1] || items[0].RightID != ids[2] { + t.Fatalf("anchor deletion lost group: %+v %v", items, err) + } + if got := items[0].FindingIDs; len(got) != 2 || got[0] != fmt.Sprint(ids[1]) || got[1] != fmt.Sprint(ids[2]) { + t.Fatalf("surviving members %+v", got) + } + if _, err := d.Exec(`DELETE FROM finding_case_suggestion_members WHERE suggestion_id=$1`, sid); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`UPDATE finding_case_suggestions SET left_id=$2,right_id=$3 WHERE id=$1`, sid, ids[1], ids[3]); err != nil { + t.Fatal(err) + } + legacy, err := d.FindingCaseSuggestions(task.ID) + if err != nil || len(legacy) != 1 || len(legacy[0].FindingIDs) != 2 { + t.Fatalf("legacy pair was not listed: %+v %v", legacy, err) + } + if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, nil); err != nil { + t.Fatal(err) + } +} + +func TestFindingCaseCandidatesScoreAndTier(t *testing.T) { + d, _, ids := caseTestDB(t) + if _, err := d.Exec(`UPDATE findings SET vulnclass='IDOR',name='order access',asset_ids='[101]'::jsonb WHERE id=$1`, ids[0]); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`UPDATE findings SET vulnclass='idor',name='different',asset_ids='[101]'::jsonb WHERE id=$1`, ids[1]); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`UPDATE findings SET vulnclass='SQLi',name='different',asset_ids='[101]'::jsonb WHERE id=$1`, ids[2]); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`UPDATE findings SET vulnclass='SQLi',name='order access',asset_ids='[]'::jsonb WHERE id=$1`, ids[3]); err != nil { + t.Fatal(err) + } + items, total, strong, err := d.FindingCaseCandidates(ids[0], 1, 20) + if err != nil || total != 3 || strong != 1 || len(items) != 3 { + t.Fatalf("candidate totals %d %d %+v %v", total, strong, items, err) + } + if items[0]["finding_id"] != fmt.Sprint(ids[1]) || items[0]["score"] != 7 || items[0]["tier"] != "strong" || !reflect.DeepEqual(items[0]["hit_reasons"], []string{"同漏洞类型", "资产重叠"}) { + t.Fatalf("strong candidate %+v", items[0]) + } + if items[1]["finding_id"] != fmt.Sprint(ids[2]) || items[1]["score"] != 2 || items[1]["tier"] != "weak" { + t.Fatalf("asset candidate %+v", items[1]) + } + if items[2]["finding_id"] != fmt.Sprint(ids[3]) || items[2]["score"] != 1 { + t.Fatalf("name candidate %+v", items[2]) + } +} diff --git a/db/intelligence.go b/db/intelligence.go index a9fc6e9..a559c50 100644 --- a/db/intelligence.go +++ b/db/intelligence.go @@ -260,8 +260,8 @@ func (d *DB) ListIntelligence(ctx context.Context, f IntelligenceFilter) (*Intel if err != nil { return nil, err } - // Search titles and reporters; return the recorded content directly in each row. - where := scope + ` AND ($3='' OR kind=$3) AND ($4='' OR i.status=$4) AND ($5='' OR i.title ILIKE '%'||$5||'%' OR i.worker ILIKE '%'||$5||'%')` + // Search titles, content, evidence, and reporters; return the recorded content directly in each row. + where := scope + ` AND ($3='' OR kind=$3) AND ($4='' OR i.status=$4) AND ($5='' OR i.title ILIKE '%'||$5||'%' OR i.content ILIKE '%'||$5||'%' OR i.evidence ILIKE '%'||$5||'%' OR i.worker ILIKE '%'||$5||'%')` args = append(args, f.Kind, f.Status, f.Query) if err = d.QueryRowContext(ctx, `SELECT count(*)`+where, args...).Scan(&page.Total); err != nil { return nil, err diff --git a/db/schema.sql b/db/schema.sql index 1fe2311..e855921 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -524,6 +524,7 @@ CREATE TABLE IF NOT EXISTS tasks ( timeout_seconds INTEGER NOT NULL DEFAULT 0, plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300, coverage_enabled BOOLEAN NOT NULL DEFAULT true, + worker_count INTEGER NOT NULL DEFAULT 0, -- 0=用全局默认 workers pinned_at TIMESTAMPTZ, first_run_at TIMESTAMPTZ, deadline_at TIMESTAMPTZ, @@ -546,6 +547,8 @@ ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false -- 给 agent 开放 add_task_scope/list_untested_assets;false=全部关闭(见 task_scope.go)。 -- 存量任务默认 true 保持原行为;company 关联(task_scope kind=company)不受此开关影响。 ALTER TABLE tasks ADD COLUMN IF NOT EXISTS coverage_enabled BOOLEAN NOT NULL DEFAULT true; +-- 任务级 Worker 运行数量;0=用全局默认(settings.workers)。补旧库。 +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS worker_count INTEGER NOT NULL DEFAULT 0; -- queued_at makes admission FIFO reflect the actual enqueue order rather than the -- task creation order. queue_mode distinguishes first bootstrap from resuming an -- exploration that already owns goals/history. @@ -617,16 +620,21 @@ CREATE TABLE IF NOT EXISTS task_templates ( nkey TEXT NOT NULL UNIQUE, description TEXT NOT NULL, goal TEXT NOT NULL, + remark TEXT NOT NULL DEFAULT '', -- 预设的任务分类；分类删除时置空（与 tasks.category_id 一致，不阻断）。 category_id BIGINT REFERENCES task_categories(id) ON DELETE SET NULL, -- 预设的任务级拦截/允许规则快照(AssetInterceptRuleInput 数组)；应用模板时灌进新任务。 intercept_rules JSONB NOT NULL DEFAULT '[]', + -- 预设的任务级 Worker 运行数量;0=用全局默认。 + worker_count INTEGER NOT NULL DEFAULT 0, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); -- 补旧库(已发版,加列带 IF NOT EXISTS)。 ALTER TABLE task_templates ADD COLUMN IF NOT EXISTS category_id BIGINT REFERENCES task_categories(id) ON DELETE SET NULL; ALTER TABLE task_templates ADD COLUMN IF NOT EXISTS intercept_rules JSONB NOT NULL DEFAULT '[]'; +ALTER TABLE task_templates ADD COLUMN IF NOT EXISTS remark TEXT NOT NULL DEFAULT ''; +ALTER TABLE task_templates ADD COLUMN IF NOT EXISTS worker_count INTEGER NOT NULL DEFAULT 0; CREATE INDEX IF NOT EXISTS idx_task_templates_updated ON task_templates(updated_at DESC, id DESC); DROP TRIGGER IF EXISTS trg_task_templates_upd ON task_templates; CREATE TRIGGER trg_task_templates_upd BEFORE UPDATE ON task_templates @@ -1488,6 +1496,54 @@ CREATE TABLE IF NOT EXISTS finding_case_suggestions ( UNIQUE(left_id,right_id), CHECK(left_id < right_id) ); +-- A suggestion is a group. left_id/right_id remain compatibility anchors for +-- old clients and are kept at the two smallest surviving member ids. +ALTER TABLE finding_case_suggestions DROP CONSTRAINT IF EXISTS finding_case_suggestions_left_id_right_id_key; +CREATE TABLE IF NOT EXISTS finding_case_suggestion_members ( + suggestion_id BIGINT NOT NULL REFERENCES finding_case_suggestions(id) ON DELETE CASCADE, + finding_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE, + PRIMARY KEY (suggestion_id, finding_id) +); +CREATE INDEX IF NOT EXISTS idx_finding_case_suggestion_members_finding + ON finding_case_suggestion_members(finding_id); +INSERT INTO finding_case_suggestion_members(suggestion_id, finding_id) + SELECT id, left_id FROM finding_case_suggestions ON CONFLICT DO NOTHING; +INSERT INTO finding_case_suggestion_members(suggestion_id, finding_id) + SELECT id, right_id FROM finding_case_suggestions ON CONFLICT DO NOTHING; +-- Anchors have cascading foreign keys. Re-anchor a group before an anchor +-- finding is deleted so a 3+ member suggestion survives the cascade. +CREATE OR REPLACE FUNCTION finding_case_suggestion_reanchor() RETURNS trigger +LANGUAGE plpgsql AS $$ +DECLARE + sid BIGINT; + new_left BIGINT; + new_right BIGINT; +BEGIN + FOR sid IN + SELECT id FROM finding_case_suggestions + WHERE left_id = OLD.id OR right_id = OLD.id + LOOP + SELECT min(m.finding_id) INTO new_left + FROM finding_case_suggestion_members m + JOIN findings f ON f.id = m.finding_id + WHERE m.suggestion_id = sid AND m.finding_id <> OLD.id; + SELECT min(m.finding_id) INTO new_right + FROM finding_case_suggestion_members m + JOIN findings f ON f.id = m.finding_id + WHERE m.suggestion_id = sid AND m.finding_id <> OLD.id + AND m.finding_id > new_left; + IF new_left IS NOT NULL AND new_right IS NOT NULL THEN + UPDATE finding_case_suggestions + SET left_id = new_left, right_id = new_right + WHERE id = sid; + END IF; + END LOOP; + RETURN OLD; +END $$; +DROP TRIGGER IF EXISTS finding_case_suggestion_reanchor ON findings; +CREATE TRIGGER finding_case_suggestion_reanchor +BEFORE DELETE ON findings FOR EACH ROW +EXECUTE FUNCTION finding_case_suggestion_reanchor(); -- Rejected/detached pairs are a durable human veto, also checked on group unions. CREATE TABLE IF NOT EXISTS finding_case_blocks ( left_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE, diff --git a/db/settings.go b/db/settings.go index f8fafcc..5c3727c 100644 --- a/db/settings.go +++ b/db/settings.go @@ -3,6 +3,7 @@ package db import ( "context" "database/sql" + "sort" ) // Settings is a tiny key-value store for global app config the UI toggles at @@ -33,6 +34,32 @@ ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key return err } +// SetSettingsAtomic upserts several settings in one transaction. Keys are +// sorted so concurrent callers acquire row locks in a deterministic order. +func (d *DB) SetSettingsAtomic(ctx context.Context, values map[string]string) error { + if len(values) == 0 { + return nil + } + tx, err := d.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck // commit makes this a no-op + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + if _, err := tx.ExecContext(ctx, ` +INSERT INTO settings(key, value) VALUES ($1, $2) +ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key, values[key]); err != nil { + return err + } + } + return tx.Commit() +} + // InsertSettingIfAbsent 只在 key 尚未存在时写入，已存在则原样保留并返回 inserted=false。 // 给 auth.password_hash 这类"只允许首次设置"的键用：判定落在数据库的主键约束上， // 调用方先 GetSetting 再写的那种检查就只是快速失败路径——读出错或并发撞车时， diff --git a/db/settings_test.go b/db/settings_test.go new file mode 100644 index 0000000..9de68a9 --- /dev/null +++ b/db/settings_test.go @@ -0,0 +1,28 @@ +package db + +import ( + "context" + "fmt" + "testing" + "time" +) + +func TestSetSettingsAtomic(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v)", err) + } + defer d.Close() + prefix := fmt.Sprintf("test.atomic_settings.%d", time.Now().UnixNano()) + defer func() { _, _ = d.Exec(`DELETE FROM settings WHERE key LIKE $1`, prefix+".%") }() + values := map[string]string{prefix + ".worker_pools": "[]", prefix + ".workers": "3"} + if err := d.SetSettingsAtomic(context.Background(), values); err != nil { + t.Fatal(err) + } + for key, want := range values { + got, ok, err := d.GetSetting(key) + if err != nil || !ok || got != want { + t.Fatalf("key %q: value=%q ok=%v err=%v, want %q", key, got, ok, err, want) + } + } +} diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go index 01094ca..56f102b 100644 --- a/db/task_archives_restore.go +++ b/db/task_archives_restore.go @@ -122,7 +122,7 @@ AND NOT EXISTS ( if _, err := tx.Exec(`UPDATE tasks SET name='',remark='',category_id=NULL,description='',goal='',paused=true,queued=false,queued_at=NULL,queue_mode='', llm_profile_id=NULL,active_llm_profile_id=NULL,llm_chain_revision=llm_chain_revision+1, - company_id=NULL,parent_ref=NULL,timeout_seconds=0,coverage_enabled=true,pinned_at=NULL, + company_id=NULL,parent_ref=NULL,timeout_seconds=0,coverage_enabled=true,worker_count=0,pinned_at=NULL, first_run_at=NULL,deadline_at=NULL,archived_at=now(),deleted_at=now() WHERE id=$1`, taskID); err != nil { return err @@ -322,7 +322,7 @@ func restoreTaskStub(tx *sql.Tx, row map[string]any, taskID, remaining int64) er llm_profile_id=archived.llm_profile_id,active_llm_profile_id=archived.active_llm_profile_id, llm_chain_revision=archived.llm_chain_revision,company_id=archived.company_id,parent_ref=archived.parent_ref, timeout_seconds=archived.timeout_seconds,plan_heartbeat_seconds=archived.plan_heartbeat_seconds, - coverage_enabled=archived.coverage_enabled,pinned_at=archived.pinned_at,first_run_at=archived.first_run_at, + coverage_enabled=archived.coverage_enabled,worker_count=COALESCE(archived.worker_count,0),pinned_at=archived.pinned_at,first_run_at=archived.first_run_at, deadline_at=CASE WHEN archived.paused AND $3>0 THEN now()+make_interval(secs=>$3::double precision) ELSE archived.deadline_at END, deleted_at=NULL,archived_at=NULL,completed_at=archived.completed_at, diff --git a/db/task_archives_test.go b/db/task_archives_test.go index 93a1364..bc4274a 100644 --- a/db/task_archives_test.go +++ b/db/task_archives_test.go @@ -132,7 +132,9 @@ func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { if err := d.EnsureLLMUsageTable(); err != nil { t.Fatal(err) } - task, err := d.CreateTaskWithOptions("archive database roundtrip", "restore exact graph", TaskCreateOptions{Name: "cold task"}) + task, err := d.CreateTaskWithOptions("archive database roundtrip", "restore exact graph", TaskCreateOptions{ + Name: "cold task", WorkerCount: 7, + }) if err != nil { t.Fatal(err) } @@ -297,7 +299,7 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro if err != nil || live == nil { t.Fatalf("restored task = %+v, %v", live, err) } - if live.Name != "cold task" || live.Remark != remark || !live.Paused { + if live.Name != "cold task" || live.Remark != remark || live.WorkerCount != 7 || !live.Paused { t.Fatalf("restored task metadata mismatch: %+v", live) } if err := d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodes); err != nil { diff --git a/db/task_options_test.go b/db/task_options_test.go new file mode 100644 index 0000000..100bda7 --- /dev/null +++ b/db/task_options_test.go @@ -0,0 +1,102 @@ +package db + +import ( + "fmt" + "strings" + "testing" + "time" +) + +func TestTaskAndTemplateRemarkLimitIsShared(t *testing.T) { + boundary := strings.Repeat("界", MaxTaskRemarkRunes) + if _, _, err := normalizeTaskTemplateInput(TaskTemplateInput{Name: "n", Description: "d", Goal: "g", Remark: boundary}); err != nil { + t.Fatalf("boundary remark rejected: %v", err) + } + if _, _, err := normalizeTaskTemplateInput(TaskTemplateInput{Name: "n", Description: "d", Goal: "g", Remark: boundary + "界"}); err == nil { + t.Fatal("oversized template remark accepted") + } +} + +func TestWorkerCountNormalizationAndLimit(t *testing.T) { + if got := normalizeWorkerCount(-1); got != 0 { + t.Fatalf("negative worker count = %d, want 0", got) + } + if _, _, err := normalizeTaskTemplateInput(TaskTemplateInput{Name: "n", Description: "d", Goal: "g", WorkerCount: MaxTaskWorkerCount}); err != nil { + t.Fatalf("maximum worker count rejected: %v", err) + } + if _, _, err := normalizeTaskTemplateInput(TaskTemplateInput{Name: "n", Description: "d", Goal: "g", WorkerCount: MaxTaskWorkerCount + 1}); err == nil { + t.Fatal("oversized worker count accepted") + } +} + +func TestTaskOptionsPersistAndValidateLimits(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + if _, err := d.CreateTaskWithOptions("oversized remark", "goal", TaskCreateOptions{ + Remark: strings.Repeat("界", MaxTaskRemarkRunes+1), + }); err == nil { + t.Fatal("2001-rune task remark accepted") + } + if _, err := d.CreateTaskWithOptions("oversized workers", "goal", TaskCreateOptions{ + WorkerCount: MaxTaskWorkerCount + 1, + }); err == nil { + t.Fatal("51-worker task accepted") + } + + remark := strings.Repeat("界", MaxTaskRemarkRunes) + task, err := d.CreateTaskWithOptions("task options boundary", "goal", TaskCreateOptions{ + Remark: remark, WorkerCount: MaxTaskWorkerCount, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + gotTask, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if gotTask == nil || gotTask.Remark != remark || gotTask.WorkerCount != MaxTaskWorkerCount { + t.Fatalf("task options round trip = %+v", gotTask) + } + + name := fmt.Sprintf("options template %d", time.Now().UnixNano()) + templateRemark := "模板备注 📝" + template, err := d.CreateTaskTemplate(TaskTemplateInput{ + Name: name, Description: "template description", Goal: "template goal", + Remark: templateRemark, WorkerCount: MaxTaskWorkerCount, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(template.ID) }) + gotTemplate, err := d.GetTaskTemplate(template.ID) + if err != nil { + t.Fatal(err) + } + if gotTemplate == nil || gotTemplate.Remark != templateRemark || gotTemplate.WorkerCount != MaxTaskWorkerCount { + t.Fatalf("created template options round trip = %+v", gotTemplate) + } + + updatedRemark := strings.Repeat("改", MaxTaskRemarkRunes) + updatedWorkers := MaxTaskWorkerCount - 1 + updated, err := d.PatchTaskTemplate(template.ID, TaskTemplatePatch{ + Remark: &updatedRemark, WorkerCount: &updatedWorkers, + }) + if err != nil { + t.Fatal(err) + } + if updated.Remark != updatedRemark || updated.WorkerCount != updatedWorkers { + t.Fatalf("updated template options = %+v", updated) + } + gotTemplate, err = d.GetTaskTemplate(template.ID) + if err != nil { + t.Fatal(err) + } + if gotTemplate == nil || gotTemplate.Remark != updatedRemark || gotTemplate.WorkerCount != updatedWorkers { + t.Fatalf("updated template options round trip = %+v", gotTemplate) + } +} diff --git a/db/task_templates.go b/db/task_templates.go index 80e3116..5401139 100644 --- a/db/task_templates.go +++ b/db/task_templates.go @@ -23,16 +23,18 @@ var ( ErrTaskTemplateNotFound = errors.New("task template not found") ) -// TaskTemplate is a reusable task preset (description/goal + optional category -// and task-level intercept/allow rules). +// TaskTemplate is a reusable task preset (description/goal + optional remark, +// worker count, category, and task-level intercept/allow rules). type TaskTemplate struct { ID int64 `json:"id"` Name string `json:"name"` NKey string `json:"-"` Description string `json:"description"` Goal string `json:"goal"` + Remark string `json:"remark"` CategoryID *int64 `json:"category_id"` InterceptRules []TaskInterceptRuleInput `json:"intercept_rules"` + WorkerCount int `json:"worker_count"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` } @@ -42,8 +44,10 @@ type TaskTemplateInput struct { Name string Description string Goal string + Remark string CategoryID *int64 InterceptRules []TaskInterceptRuleInput + WorkerCount int } // TaskTemplatePatch changes only fields flagged as set. Name/Description/Goal use @@ -53,18 +57,20 @@ type TaskTemplatePatch struct { Name *string Description *string Goal *string + Remark *string CategoryID *int64 SetCategoryID bool InterceptRules []TaskInterceptRuleInput SetInterceptRules bool + WorkerCount *int } -const taskTemplateCols = `id, name, nkey, description, goal, category_id, intercept_rules, created_at, updated_at` +const taskTemplateCols = `id, name, nkey, description, goal, COALESCE(remark,''), category_id, intercept_rules, COALESCE(worker_count,0), created_at, updated_at` func scanTaskTemplate(row interface{ Scan(...any) error }) (TaskTemplate, error) { var t TaskTemplate var rulesRaw []byte - if err := row.Scan(&t.ID, &t.Name, &t.NKey, &t.Description, &t.Goal, &t.CategoryID, &rulesRaw, &t.CreatedAt, &t.UpdatedAt); err != nil { + if err := row.Scan(&t.ID, &t.Name, &t.NKey, &t.Description, &t.Goal, &t.Remark, &t.CategoryID, &rulesRaw, &t.WorkerCount, &t.CreatedAt, &t.UpdatedAt); err != nil { return t, err } t.InterceptRules = []TaskInterceptRuleInput{} @@ -112,11 +118,21 @@ func normalizeTaskTemplateInput(in TaskTemplateInput) (TaskTemplateInput, string case utf8.RuneCountInString(in.Goal) > MaxTaskTemplateTextRunes: return in, "", fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes) } + in.Remark = strings.TrimSpace(in.Remark) + // Keep template remarks within the same shared limit as task remarks so + // applying a template cannot produce a rejected create-task payload. + if utf8.RuneCountInString(in.Remark) > MaxTaskRemarkRunes { + return in, "", fmt.Errorf("%w: remark exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskRemarkRunes) + } + in.WorkerCount = normalizeWorkerCount(in.WorkerCount) + if err := validateWorkerCount(in.WorkerCount); err != nil { + return in, "", fmt.Errorf("%w: %v", ErrTaskTemplateInvalid, err) + } return in, taskTemplateNKey(in.Name), nil } func normalizeTaskTemplatePatch(patch TaskTemplatePatch) (TaskTemplatePatch, *string, error) { - if patch.Name == nil && patch.Description == nil && patch.Goal == nil && !patch.SetCategoryID && !patch.SetInterceptRules { + if patch.Name == nil && patch.Description == nil && patch.Goal == nil && patch.Remark == nil && !patch.SetCategoryID && !patch.SetInterceptRules && patch.WorkerCount == nil { return patch, nil, fmt.Errorf("%w: no fields supplied", ErrTaskTemplateInvalid) } var nkey *string @@ -152,6 +168,20 @@ func normalizeTaskTemplatePatch(patch TaskTemplatePatch) (TaskTemplatePatch, *st } patch.Goal = &goal } + if patch.Remark != nil { + remark := strings.TrimSpace(*patch.Remark) + if utf8.RuneCountInString(remark) > MaxTaskRemarkRunes { + return patch, nil, fmt.Errorf("%w: remark exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskRemarkRunes) + } + patch.Remark = &remark + } + if patch.WorkerCount != nil { + n := normalizeWorkerCount(*patch.WorkerCount) + if err := validateWorkerCount(n); err != nil { + return patch, nil, fmt.Errorf("%w: %v", ErrTaskTemplateInvalid, err) + } + patch.WorkerCount = &n + } return patch, nkey, nil } @@ -171,10 +201,10 @@ func (d *DB) CreateTaskTemplate(in TaskTemplateInput) (*TaskTemplate, error) { return nil, err } t, err := scanTaskTemplate(d.QueryRow(` -INSERT INTO task_templates(name, nkey, description, goal, category_id, intercept_rules) -VALUES ($1,$2,$3,$4,$5,$6) +INSERT INTO task_templates(name, nkey, description, goal, remark, category_id, intercept_rules, worker_count) +VALUES ($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT (nkey) DO NOTHING -RETURNING `+taskTemplateCols, in.Name, nkey, in.Description, in.Goal, in.CategoryID, rulesJSON)) +RETURNING `+taskTemplateCols, in.Name, nkey, in.Description, in.Goal, in.Remark, in.CategoryID, rulesJSON, in.WorkerCount)) if err == sql.ErrNoRows { return nil, ErrTaskTemplateNameConflict } @@ -221,7 +251,8 @@ func (d *DB) UpdateTaskTemplate(id int64, in TaskTemplateInput) (*TaskTemplate, return nil, err } return d.PatchTaskTemplate(id, TaskTemplatePatch{ - Name: &in.Name, Description: &in.Description, Goal: &in.Goal, + Name: &in.Name, Description: &in.Description, Goal: &in.Goal, Remark: &in.Remark, + WorkerCount: &in.WorkerCount, }) } @@ -242,16 +273,20 @@ SET name=CASE WHEN $2 THEN $3::text ELSE name END, nkey=CASE WHEN $2 THEN $4::text ELSE nkey END, description=CASE WHEN $5 THEN $6::text ELSE description END, goal=CASE WHEN $7 THEN $8::text ELSE goal END, - category_id=CASE WHEN $9 THEN $10::bigint ELSE category_id END, - intercept_rules=CASE WHEN $11 THEN $12::jsonb ELSE intercept_rules END + remark=CASE WHEN $9 THEN $10::text ELSE remark END, + category_id=CASE WHEN $11 THEN $12::bigint ELSE category_id END, + intercept_rules=CASE WHEN $13 THEN $14::jsonb ELSE intercept_rules END, + worker_count=CASE WHEN $15 THEN $16::integer ELSE worker_count END WHERE id=$1 RETURNING `+taskTemplateCols, id, patch.Name != nil, patch.Name, nkey, patch.Description != nil, patch.Description, patch.Goal != nil, patch.Goal, + patch.Remark != nil, patch.Remark, patch.SetCategoryID, patch.CategoryID, patch.SetInterceptRules, rulesJSON, + patch.WorkerCount != nil, patch.WorkerCount, )) if err == sql.ErrNoRows { return nil, ErrTaskTemplateNotFound diff --git a/db/tasks.go b/db/tasks.go index b675bea..c17f1b6 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -6,7 +6,9 @@ import ( "errors" "fmt" "strconv" + "strings" "time" + "unicode/utf8" ) // Task is a row in the task registry (1:1 with an exploration). @@ -51,6 +53,9 @@ type Task struct { // list_untested_assets、态势里不注入 coverage 块(scope 字段仍保留)。company 关联 // (task_scope kind=company)与此开关无关，永不受影响。见 db/task_scope.go。 CoverageEnabled bool `json:"coverage_enabled"` + // WorkerCount is the per-task concurrent work-agent count. 0 means use the + // global settings.workers default when the engine starts this task. + WorkerCount int `json:"worker_count"` } // TaskDeleteResult reports optional related-data cleanup performed in the same @@ -132,6 +137,29 @@ func normalizeHeartbeat(sec int) int { return sec } +// MaxTaskWorkerCount is the hard upper bound for a per-task worker override. +const MaxTaskWorkerCount = 50 + +// MaxTaskRemarkRunes is the shared limit for task and template remarks. +const MaxTaskRemarkRunes = 2000 + +// normalizeWorkerCount converts negative per-task worker overrides to 0. +// The upper bound is enforced separately by validateWorkerCount; 0 means +// "use the global settings.workers default". +func normalizeWorkerCount(n int) int { + if n < 0 { + return 0 + } + return n +} + +func validateWorkerCount(n int) error { + if n > MaxTaskWorkerCount { + return fmt.Errorf("Worker 运行数量最多 %d", MaxTaskWorkerCount) + } + return nil +} + func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { var ids []int64 if llmProfileID != nil { @@ -156,6 +184,10 @@ type TaskCreateOptions struct { // CoverageEnabled 是「资产覆盖度功能」开关;nil=默认开(true)，让不关心该开关的创建 // 路径(编排 spawn、老 API)沿用原行为。仅 web 创建任务时可显式传 false 关闭。 CoverageEnabled *bool + // WorkerCount 是本任务并发运行的工作 Agent 数;0/省略=用全局默认。上限 50。 + WorkerCount int + // Remark 是创建时可带的任务备注(最多 2000 字);空=无备注。 + Remark string // InterceptRules 是任务级资产拦截规则,创建时随任务在同一事务内写入 task_intercept_rules。 InterceptRules []TaskInterceptRuleInput } @@ -200,6 +232,14 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); } opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds) coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled + opts.WorkerCount = normalizeWorkerCount(opts.WorkerCount) + if err := validateWorkerCount(opts.WorkerCount); err != nil { + return nil, err + } + opts.Remark = strings.TrimSpace(opts.Remark) + if utf8.RuneCountInString(opts.Remark) > MaxTaskRemarkRunes { + return nil, fmt.Errorf("备注最多 %d 个字符", MaxTaskRemarkRunes) + } var categoryName string if opts.CategoryID != nil { if *opts.CategoryID <= 0 { @@ -225,12 +265,12 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...), CompanyIDs: append([]int64(nil), opts.CompanyIDs...), TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds, - CoverageEnabled: coverageEnabled, + CoverageEnabled: coverageEnabled, WorkerCount: opts.WorkerCount, Remark: opts.Remark, } if err := tx.QueryRow(` -INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled) -VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9) -RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(name, remark, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled, worker_count) +VALUES ($1,$2,$3,$4,$5,$6,$7,$7,$8,$9,$10,$11) +RETURNING id, status, paused, created_at`, opts.Name, opts.Remark, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled, opts.WorkerCount).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { @@ -363,11 +403,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { const taskCols = `id, COALESCE(name,''), remark, category_id, COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''), -description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at` +description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), COALESCE(worker_count,0), first_run_at, deadline_at` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Name, &t.Remark, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Name, &t.Remark, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.WorkerCount, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } t.Pinned = t.PinnedAt != nil diff --git a/go.mod b/go.mod index 35f1a14..ee5f10a 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/Autumn-27/artex -go 1.26.3 +go 1.26.10 require ( mvdan.cc/sh/v3 v3.10.0 diff --git a/server/assembly_test.go b/server/assembly_test.go index 1eabeba..1ea94bf 100644 --- a/server/assembly_test.go +++ b/server/assembly_test.go @@ -6,9 +6,9 @@ import ( "path/filepath" "testing" - actool "github.com/Autumn-27/norma/tool" "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" + actool "github.com/Autumn-27/norma/tool" ) // hasSkillTool reports whether the packed tool set contains the Skill meta-tool. diff --git a/server/customtool.go b/server/customtool.go index 4d9e30e..ad68517 100644 --- a/server/customtool.go +++ b/server/customtool.go @@ -462,9 +462,16 @@ func (s *Server) httpProxyTransport(spec httpExec) *http.Transport { proxyStr := strings.TrimSpace(spec.Proxy) var caFile string if spec.UseRecordingProxy { - if addr := s.m.ProxyAddr(); addr != "" { - proxyStr = "http://" + addr - caFile = s.m.ProxyCACert() + // ProxyAddr is already a complete URL while capture is enabled, and + // returns the configured global proxy URL when capture is off. Keep both + // forms intact (including socks5:// and credentials). + if s.m.TrafficEnabled() { + if addr := s.m.ProxyAddr(); addr != "" { + proxyStr = addr + caFile = s.m.ProxyCACert() + } + } else if global := strings.TrimSpace(s.m.GlobalProxy()); global != "" { + proxyStr = global } } if proxyStr == "" { @@ -474,6 +481,9 @@ func (s *Server) httpProxyTransport(spec httpExec) *http.Transport { if err != nil { return nil } + // Go 1.26+ supports http/https/socks5/socks5h in ProxyURL directly, + // including proxy userinfo. Keep one code path so existing HTTP tool + // behavior and proxy authentication remain unchanged. tr := &http.Transport{Proxy: http.ProxyURL(pu)} if caFile != "" { if pem, err := os.ReadFile(caFile); err == nil { diff --git a/server/dto.go b/server/dto.go index 60785a5..339c1f2 100644 --- a/server/dto.go +++ b/server/dto.go @@ -64,6 +64,7 @@ type TaskDTO struct { ArchiveBlockedBy string `json:"archive_blocked_by_task_id,omitempty"` CompanyIDs []int64 `json:"company_ids"` CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定) + WorkerCount int `json:"worker_count"` // 本任务 Worker 数;0=用全局默认 } // FindingSeverityDTO 是任务列表里按严重度分档的漏洞计数（严重/高/中/低）。 @@ -138,6 +139,7 @@ func taskDTO(t *Task, status string) TaskDTO { SourceTaskIDs: sourceIDs, CompanyIDs: lifecycle.CompanyIDs, CoverageEnabled: t.CoverageEnabled, + WorkerCount: t.WorkerCount, } } diff --git a/server/engine.go b/server/engine.go index b5643b2..7c61b74 100644 --- a/server/engine.go +++ b/server/engine.go @@ -148,11 +148,19 @@ type Engine struct { // (for example an exhausted failover chain) — there is no global-pair fallback. resolve func(t *Task) (*agent.Planner, *agent.Worker) resolveAuthoritative bool + workerResolve func(t *Task, slot WorkerSlot) (*agent.Worker, error) + slotState sync.Map // task/slot -> slotAvailability + slotPlans sync.Map // task -> map[slot name]profile ID; immutable for a run // readiness reports whether a global LLM provider is configured — the signal behind // Ready()/the llm_configured indicator. Wired once at startup; nil → not ready. readiness func() bool } +type slotAvailability struct { + available bool + lastEmit time.Time +} + type taskRuntime struct { ctx context.Context cancel context.CancelFunc @@ -285,6 +293,14 @@ func (e *Engine) StopTask(taskID string) { e.stamped.Delete(taskID) e.inflight.Delete(taskID) e.coordStarted.Delete(taskID) + e.slotPlans.Delete(taskID) + e.slotState.Range(func(key, _ any) bool { + keyText, ok := key.(string) + if ok && strings.HasPrefix(keyText, taskID+"/") { + e.slotState.Delete(key) + } + return true + }) e.deleteMu.Lock() e.deleting.Delete(taskID) e.deleteMu.Unlock() @@ -715,6 +731,37 @@ func (e *Engine) SetAuthoritativeAgentResolver(fn func(t *Task) (*agent.Planner, e.resolveAuthoritative = true } +func (e *Engine) SetWorkerResolver(fn func(t *Task, slot WorkerSlot) (*agent.Worker, error)) { + e.workerResolve = fn +} + +func (e *Engine) reportSlotAvailability(t *Task, slot WorkerSlot, err error) { + key := t.ID + "/" + slot.Name + now := time.Now() + value, loaded := e.slotState.LoadOrStore(key, slotAvailability{available: err == nil, lastEmit: now}) + state := value.(slotAvailability) + if err == nil { + if loaded && !state.available { + e.slotState.Store(key, slotAvailability{available: true, lastEmit: now}) + e.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: fmt.Sprintf("%s 绑定的 LLM 配置 #%d 已恢复，可以继续领取任务", slot.Name, slot.ProfileID)}) + } + return + } + if loaded && state.available { + state.available = false + state.lastEmit = now + e.slotState.Store(key, state) + } else if loaded && now.Sub(state.lastEmit) < 5*time.Minute { + return + } else { + state = slotAvailability{available: false, lastEmit: now} + e.slotState.Store(key, state) + } + summary := fmt.Sprintf("%s 绑定的 LLM 配置 #%d 不可用（已删除或配置无效），该会话暂停领取任务，不会回退到默认模型", slot.Name, slot.ProfileID) + e.emitActivity(t, db.Activity{Worker: "system", Kind: "text", IsError: true, Summary: summary, Detail: err.Error()}) + log.Printf("[worker %s] task %s pinned profile #%d unavailable: %v", slot.Name, t.ID, slot.ProfileID, err) +} + // snapshotFor returns the planner/worker a task should run on, from the task-router // resolver. nil,nil means the task is deliberately unavailable (e.g. an exhausted // failover chain); there is no global-pair fallback. @@ -746,7 +793,11 @@ func (e *Engine) ReadyFor(t *Task) bool { // but no-op until an LLM is configured (so a task created while idle picks up // automatically once LLM is set from the UI). func (e *Engine) Run(ctx context.Context, t *Task) { - workers := e.m.Workers() + slots, err := workerSlotsForTask(t, e.m.WorkerPools(), e.m.Workers()) + if err != nil { + log.Printf("[engine] task %s worker pool invalid, using legacy workers: %v", t.ID, err) + slots, _ = ExpandWorkerSlots(nil, workerCountForTask(t, e.m.Workers())) + } e.deleteMu.RLock() if e.IsDeleting(t.ID) { e.deleteMu.RUnlock() @@ -757,13 +808,18 @@ func (e *Engine) Run(ctx context.Context, t *Task) { t.Notify() // already running — just nudge a planning round return } - rt := e.registerTaskRoutines(ctx, t.ID, 1+workers) + rt := e.registerTaskRoutines(ctx, t.ID, 1+len(slots)) + plan := make(map[string]int64, len(slots)) + for _, slot := range slots { + plan[slot.Name] = slot.ProfileID + } + e.slotPlans.Store(t.ID, plan) e.deleteMu.RUnlock() e.touch(t.ID) runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) }) - for i := 0; i < workers; i++ { - name := fmt.Sprintf("work#%d", i+1) - runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, name) }) + for _, slot := range slots { + slot := slot + runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, slot) }) } e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重) // 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已 @@ -776,6 +832,27 @@ func (e *Engine) Run(ctx context.Context, t *Task) { } } +// workerSlotsForTask preserves an explicit task concurrency override. Tasks +// using the global default inherit its model pools and their pinned profiles. +func workerSlotsForTask(t *Task, pools []WorkerPool, global int) ([]WorkerSlot, error) { + if t != nil && t.WorkerCount > 0 { + return ExpandWorkerSlots(nil, workerCountForTask(t, global)) + } + return ExpandWorkerSlots(pools, global) +} + +func workerCountForTask(t *Task, global int) int { + if t != nil && t.WorkerCount > 0 { + // Direct database edits can bypass validation; cap those values so the + // engine never starts more than the supported per-task worker count. + if t.WorkerCount > db.MaxTaskWorkerCount { + return db.MaxTaskWorkerCount + } + return t.WorkerCount + } + return global +} + // plannerHeartbeatInterval 解析任务的 planner 心跳间隔。db.CreateTask 已归一 // (低于 600 一律抬到 600);这里再兜一次底,防内存态异常值。 func plannerHeartbeatInterval(t *Task) time.Duration { @@ -913,14 +990,32 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { } } -func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { +func (e *Engine) workerLoop(ctx context.Context, t *Task, slot WorkerSlot) { + name := slot.Name for { select { case <-ctx.Done(): return default: } - _, worker := e.snapshotFor(t) + var worker *agent.Worker + var resolveErr error + if slot.ProfileID > 0 { + if e.workerResolve == nil { + resolveErr = errors.New("pinned worker resolver unavailable") + } else { + worker, resolveErr = e.workerResolve(t, slot) + if worker == nil && resolveErr == nil { + resolveErr = errors.New("pinned worker resolver returned no worker") + } + } + e.reportSlotAvailability(t, slot, resolveErr) + if resolveErr != nil { + worker = nil + } + } else { + _, worker = e.snapshotFor(t) + } if worker == nil { if sleepCtx(ctx, 1500*time.Millisecond) { return @@ -1178,10 +1273,6 @@ func (e *Engine) runDetachedIntent(ctx context.Context, t *Task, intentID int64, e.decInflight(t.ID) } }() - _, worker := e.snapshotFor(t) - if worker == nil { - return fmt.Errorf("worker 尚未就绪") - } node, err := t.Store.GetNode(intentID) if err != nil { return err @@ -1189,6 +1280,30 @@ func (e *Engine) runDetachedIntent(ctx context.Context, t *Task, intentID int64, if node == nil || node.Kind != db.KindIntent { return fmt.Errorf("intent not found") } + owner := node.Owner + var worker *agent.Worker + if value, ok := e.slotPlans.Load(t.ID); ok { + if plan, ok := value.(map[string]int64); ok { + if profileID, pinned := plan[owner]; pinned && profileID > 0 { + if e.workerResolve == nil { + return fmt.Errorf("意图所属的 LLM 配置 #%d 不可用：pinned worker resolver unavailable", profileID) + } + worker, err = e.workerResolve(t, WorkerSlot{Name: owner, ProfileID: profileID}) + if err != nil || worker == nil { + if err == nil { + err = errors.New("worker resolver returned no worker") + } + return fmt.Errorf("意图所属的 LLM 配置 #%d 不可用：%w", profileID, err) + } + } + } + } + if worker == nil { + _, worker = e.snapshotFor(t) + if worker == nil { + return fmt.Errorf("worker 尚未就绪") + } + } changed, err := t.Store.CompareAndSetIntentState(intentID, "paused", "running") if err != nil { return err diff --git a/server/engine_worker_test.go b/server/engine_worker_test.go new file mode 100644 index 0000000..25b1f4b --- /dev/null +++ b/server/engine_worker_test.go @@ -0,0 +1,59 @@ +package server + +import ( + "fmt" + "testing" + + "github.com/Autumn-27/artex/db" +) + +func TestWorkerCountForTaskUsesOverrideOrGlobalDefault(t *testing.T) { + if got := workerCountForTask(&Task{WorkerCount: 5}, 3); got != 5 { + t.Fatalf("override = %d, want 5", got) + } + if got := workerCountForTask(&Task{WorkerCount: 0}, 3); got != 3 { + t.Fatalf("zero override = %d, want global 3", got) + } + if got := workerCountForTask(&Task{WorkerCount: -1}, 3); got != 3 { + t.Fatalf("negative override = %d, want global 3", got) + } + if got := workerCountForTask(&Task{WorkerCount: db.MaxTaskWorkerCount + 1}, 3); got != db.MaxTaskWorkerCount { + t.Fatalf("oversized override = %d, want capped %d", got, db.MaxTaskWorkerCount) + } +} + +func TestWorkerSlotsForTaskCombinesOverridesAndModelPools(t *testing.T) { + pools := []WorkerPool{{ProfileID: 7, Workers: 2}, {ProfileID: 12, Workers: 1}} + for _, tc := range []struct { + name string + task *Task + pools []WorkerPool + global int + profiles []int64 + }{ + {"task inherits model pools", &Task{}, pools, 9, []int64{7, 7, 12}}, + {"task override replaces global pools", &Task{WorkerCount: 2}, pools, 9, []int64{0, 0}}, + {"task inherits legacy count", &Task{}, nil, 4, []int64{0, 0, 0, 0}}, + {"task override without pools", &Task{WorkerCount: 1}, nil, 4, []int64{0}}, + {"invalid task value uses pools", &Task{WorkerCount: -1}, pools, 9, []int64{7, 7, 12}}, + } { + t.Run(tc.name, func(t *testing.T) { + slots, err := workerSlotsForTask(tc.task, tc.pools, tc.global) + if err != nil || len(slots) != len(tc.profiles) { + t.Fatalf("slots=%+v err=%v, want profiles %v", slots, err, tc.profiles) + } + for i, profile := range tc.profiles { + if slots[i].ProfileID != profile || slots[i].Name != fmt.Sprintf("work#%d", i+1) { + t.Fatalf("slot[%d]=%+v, want profile %d", i, slots[i], profile) + } + } + }) + } + slots, err := workerSlotsForTask(&Task{WorkerCount: db.MaxTaskWorkerCount + 1}, pools, 3) + if err != nil || len(slots) != db.MaxTaskWorkerCount { + t.Fatalf("oversized task override: slots=%d err=%v", len(slots), err) + } + if _, err := workerSlotsForTask(&Task{}, []WorkerPool{{ProfileID: 0, Workers: 1}}, 3); err == nil { + t.Fatal("invalid inherited pool accepted") + } +} diff --git a/server/finding_case_reviews.go b/server/finding_case_reviews.go index 16a9936..e0ea191 100644 --- a/server/finding_case_reviews.go +++ b/server/finding_case_reviews.go @@ -33,9 +33,9 @@ func (s *Server) toolCompleteFindingCaseReview() actool.CoreTool { decision := objSchema(map[string]any{ "finding_id": idParam("所选独立漏洞记录 ID"), "verdict": map[string]any{"type": "string", "enum": []string{"independent", "merged", "suggested"}}, - "reason": strParam("判断理由，接口或读写影响不同不能单独证明不同根因"), - "root_cause_evidence": strParam("原始报告/证据中支持实际共同或不同根因的具体材料；不确定则说明缺失证据及已保存建议")}, "finding_id", "verdict", "reason", "root_cause_evidence") - return wrTool("complete_finding_case_review", "历史整理结束前提交每条所选记录的根因判断并验收：必须已读取全部记录，疑似必须已有建议，归并必须已有完整最新统一报告。验收失败不能宣布完成。", objSchema(map[string]any{"summary": strParam("整理结果总结"), "decisions": map[string]any{"type": "array", "items": decision}}, "summary", "decisions"), func(ctx context.Context, raw json.RawMessage) (actool.Result, error) { + "reason": strParam("判断理由，引用已证据化的共同点或不同代码路径；仅类型、标题或资产相同不足以证明共同根因"), + "root_cause_evidence": strParam("原始报告/证据中支持实际共同或不同根因的具体材料；suggested 须说明具体共同点、缺失的关键证据及已保存分组；没有具体共同点应判 independent")}, "finding_id", "verdict", "reason", "root_cause_evidence") + return wrTool("complete_finding_case_review", "历史整理结束前提交每条所选记录的根因判断并验收：必须已读取全部记录，疑似必须已有待确认分组建议；归并只需已在活动文件夹中，统一报告可稍后补齐。验收失败不能宣布完成。", objSchema(map[string]any{"summary": strParam("整理结果总结"), "decisions": map[string]any{"type": "array", "items": decision}}, "summary", "decisions"), func(ctx context.Context, raw json.RawMessage) (actool.Result, error) { var a findingReviewConclusion if err := json.Unmarshal(raw, &a); err != nil { return actool.Errorf(err.Error()), nil @@ -101,7 +101,9 @@ func (s *Server) saveFindingReviewConclusion(ctx context.Context, conv int64, a } case "suggested": var pending bool - if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM finding_case_suggestions WHERE state='pending' AND (left_id=$1 OR right_id=$1))`, id).Scan(&pending); err != nil { + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM finding_case_suggestions s + JOIN finding_case_suggestion_members sm ON sm.suggestion_id=s.id + WHERE s.state='pending' AND sm.finding_id=$1)`, id).Scan(&pending); err != nil { return err } if !pending { @@ -115,17 +117,12 @@ func (s *Server) saveFindingReviewConclusion(ctx context.Context, conv int64, a return errors.New("无效判断类型") } if cid.Valid { - var report string - var version, rv int64 var active bool - if err := tx.QueryRow(`SELECT report,version,report_version,active FROM finding_cases WHERE id=$1 FOR SHARE`, cid.Int64).Scan(&report, &version, &rv, &active); err != nil { + if err := tx.QueryRow(`SELECT active FROM finding_cases WHERE id=$1 FOR SHARE`, cid.Int64).Scan(&active); err != nil { return err } - if !active || rv != version { - return fmt.Errorf("文件夹 #%d 的统一报告待更新，请重新读取后更新", cid.Int64) - } - if err := validateUnifiedFindingReport(report); err != nil { - return err + if !active { + return fmt.Errorf("文件夹 #%d 已关闭", cid.Int64) } } } diff --git a/server/finding_case_reviews_test.go b/server/finding_case_reviews_test.go index 06e1b58..2b15441 100644 --- a/server/finding_case_reviews_test.go +++ b/server/finding_case_reviews_test.go @@ -4,11 +4,46 @@ import ( "context" "encoding/json" "fmt" - "github.com/Autumn-27/artex/intercept" "strings" "testing" + + "github.com/Autumn-27/artex/intercept" ) +func TestFindingCaseReviewTaskMessagePolicy(t *testing.T) { + message := fmt.Sprintf(findingCaseReviewTaskMessage, `[101,102,103]`) + for _, want := range []string{ + "finding_ids=[101,102,103]", + "根因明确相同则直接归并", + "相同端点、处理器/后端组件、参数或缺失检查", + "缺少确认共同根因的关键证据", + "一个覆盖2..20条同任务疑似成员的分组", + "不要拆成两两建议", + "仅类型或资产相同不足以建议", + "没有具体共同点时保持 independent", + "不同代码路径时可判 independent", + "尽力更新统一报告", + "统一报告暂缺或待更新不阻塞 complete_finding_case_review", + "不归并未选记录", + "不探测目标", + "遵守人工否决", + } { + if !strings.Contains(message, want) { + t.Errorf("task message missing policy: %s", want) + } + } + if strings.Contains(message, "不确定提交建议") { + t.Fatal("task message suggests solely because of uncertainty") + } + schema, err := json.Marshal((&Server{}).toolCompleteFindingCaseReview().InputSchema()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(schema), "没有具体共同点应判 independent") || strings.Contains(string(schema), "不确定则说明缺失证据及已保存建议") { + t.Fatal("completion schema must not require suggestions solely because of uncertainty") + } +} + const completeFindingReviewReport = `# 同一权限缺陷 ## 概述 同一对象权限检查缺失。 @@ -43,7 +78,7 @@ func TestUnifiedFindingReportRejectsIncompleteAndTruncated(t *testing.T) { } } -func TestFindingCaseReviewCompletionRequiresReadsAndFreshReport(t *testing.T) { +func TestFindingCaseReviewCompletionRequiresReadsNotFreshReport(t *testing.T) { s, task, ids := testCaseServer(t) runs, err := s.m.pg.CreateFindingCaseReviews(t.Context(), map[int64][]int64{task: ids}) if err != nil { @@ -71,13 +106,6 @@ func TestFindingCaseReviewCompletionRequiresReadsAndFreshReport(t *testing.T) { if err != nil { t.Fatal(err) } - if err := s.saveFindingReviewConclusion(ctx, conv, a); err == nil { - t.Fatal("missing unified report accepted") - } - c, _ := s.m.pg.GetFindingCase(cid) - if err := s.m.pg.UpdateFindingCaseReport(ctx, cid, c.Version, "orders", completeFindingReviewReport, "high", "verified impact"); err != nil { - t.Fatal(err) - } if err := s.saveFindingReviewConclusion(ctx, conv, a); err != nil { t.Fatal(err) } @@ -90,8 +118,8 @@ func TestFindingCaseReviewCompletionRequiresReadsAndFreshReport(t *testing.T) { if _, err := s.m.pg.Exec(`UPDATE findings SET evidence='new proof' WHERE id=$1`, ids[0]); err != nil { t.Fatal(err) } - if err := s.validateFindingReviewCompletion(conv); err == nil { - t.Fatal("stale group reported completed") + if err := s.validateFindingReviewCompletion(conv); err != nil { + t.Fatalf("stale unified report should not block merged conclusion: %v", err) } } @@ -138,7 +166,7 @@ func TestFindingCaseReviewSuggestionNeedsSavedSuggestion(t *testing.T) { if err := s.saveFindingReviewConclusion(ctx, conv, a); err == nil { t.Fatal("unpersisted suggestion accepted") } - if _, err := s.m.pg.SuggestFindingCase(ctx, task, ids[0], ids[1], "orders", "common component unknown"); err != nil { + if _, err := s.m.pg.SuggestFindingCase(ctx, task, []int64{ids[0], ids[1]}, "orders", "common component unknown"); err != nil { t.Fatal(err) } if err := s.saveFindingReviewConclusion(ctx, conv, a); err != nil { diff --git a/server/finding_cases.go b/server/finding_cases.go index d8afaf8..c4466e6 100644 --- a/server/finding_cases.go +++ b/server/finding_cases.go @@ -6,7 +6,9 @@ import ( "encoding/json" "errors" "fmt" + "log" "net/http" + "strconv" "strings" "github.com/Autumn-27/artex/db" @@ -203,12 +205,22 @@ func (s *Server) caseSuggestions(w http.ResponseWriter, r *http.Request) { } func (s *Server) resolveCaseSuggestion(w http.ResponseWriter, r *http.Request) { var a struct { - Accept bool `json:"accept"` + Accept bool `json:"accept"` + ExcludeIDs []string `json:"exclude_ids"` } if !decodeCaseRequest(w, r, &a) { return } - id, err := s.m.pg.ResolveFindingCaseSuggestion(r.Context(), int64(atoiDefault(r.PathValue("id"), 0)), a.Accept) + exclude := make([]int64, 0, len(a.ExcludeIDs)) + for _, raw := range a.ExcludeIDs { + id, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64) + if err != nil || id <= 0 { + caseHTTPError(w, errors.New("无效的排除记录编号")) + return + } + exclude = append(exclude, id) + } + id, err := s.m.pg.ResolveFindingCaseSuggestion(r.Context(), int64(atoiDefault(r.PathValue("id"), 0)), a.Accept, exclude) if err != nil { caseHTTPError(w, err) return @@ -253,10 +265,10 @@ func (s *Server) findingCaseTools() []actool.CoreTool { name, desc string required []string }{ - {"search_finding_duplicates", "只读检索同任务候选，返回摘要与资产；等级不同也可能重复。支持分页，按需读取详情。", []string{"finding_id"}}, + {"search_finding_duplicates", "只读检索同任务候选摘要与资产，按分数排序并返回 score、tier、hit_reasons：strong=同漏洞类型且资产重叠，weak=未同时具备同漏洞类型和资产重叠，请结合 hit_reasons 核对，仅作参考。等级不同也可能重复。支持分页，按需读取详情。", []string{"finding_id"}}, {"get_finding_record", "读取独立漏洞完整原始证据和报告，返回探索节点编号以便追溯执行记录。", []string{"finding_id"}}, {"merge_finding_records", "证据明确证明同一缺陷时归并；保留所有原始记录。不跨任务、不改变原评级，不可绕过人工否决。", []string{"task_id", "finding_ids", "title", "reason"}}, - {"suggest_finding_merge", "证据不足时提交疑似重复建议，待人工确认；不是自动合并。只允许两条同任务记录。", []string{"task_id", "finding_ids", "title", "reason"}}, + {"suggest_finding_merge", "证据存在具体共同点但缺少关键证据时，提交一个疑似重复分组（2..20 条同任务记录，一次覆盖全部疑似成员，不要拆成两两建议），待人工确认；不是自动合并。仅类型/资产相同不应提交。", []string{"task_id", "finding_ids", "title", "reason"}}, {"get_finding_case", "读取归并组和成员分页。先记录version，再读取成员的原报告、证据及流量；版本变化必须重新读取。", []string{"case_id"}}, {"update_finding_case_report", "保存统一报告及评级理由，version必须来自实际读取。保留各条原报告，不新建漏洞。", []string{"case_id", "version", "title", "report", "severity", "severity_reason"}}, } @@ -312,8 +324,9 @@ func (s *Server) performFindingCaseTool(ctx context.Context, name string, a find if scope != nil { scopes = append(scopes, scope) } - items, total, err := s.m.pg.FindingCaseCandidates(fid, page, 20, scopes...) - return map[string]any{"items": items, "total": total, "page": page, "page_size": 20}, err + items, total, strongTotal, err := s.m.pg.FindingCaseCandidates(fid, page, 20, scopes...) + return map[string]any{"items": items, "total": total, "strong_total": strongTotal, "page": page, "page_size": 20, + "note": "按分数排序；weak 表示未同时具备同漏洞类型和资产重叠，请结合 hit_reasons 核对，仅作参考，通常不应据此建议归并"}, err case "get_finding_record": f, err := s.m.pg.GetFinding(fid) if err != nil { @@ -341,11 +354,11 @@ func (s *Server) performFindingCaseTool(ctx context.Context, name string, a find return map[string]string{"case_id": i64s(id)}, err case "suggest_finding_merge": ids := a.ids() - if len(ids) != 2 { - return nil, errors.New("疑似重复建议需要两条记录") + if len(ids) < 2 || len(ids) > 20 { + return nil, errors.New("疑似重复分组需要2至20条记录") } - id, err := s.m.pg.SuggestFindingCase(ctx, parseProfileID(a.TaskID), ids[0], ids[1], a.Title, a.Reason) - return map[string]string{"suggestion_id": i64s(id)}, err + id, err := s.m.pg.SuggestFindingCase(ctx, parseProfileID(a.TaskID), ids, a.Title, a.Reason) + return map[string]any{"suggestion_id": i64s(id), "finding_ids": ids}, err case "get_finding_case": c, err := s.m.pg.GetFindingCase(cid) if err != nil { @@ -378,14 +391,22 @@ func (s *Server) seedFindingCaseTools() { if err != nil || a == nil { return } - const flag = "finding_case_tools_v2_review" + const flag = "finding_case_tools_v3_review" if v, _, _ := s.m.pg.GetSetting(flag); v == "true" { return } names := []string{} previous, _, _ := s.m.pg.GetSetting("finding_case_tools_v1") + // v3 refreshes descriptions and schemas without repeating the v2 binding migration. + previousReview, _, _ := s.m.pg.GetSetting("finding_case_tools_v2_review") for _, t := range s.findingCaseTools() { schema, _ := json.Marshal(t.InputSchema()) + if err := s.m.pg.RefreshToolDefaults(t.Name(), t.Description(), schema); err != nil { + log.Printf("[finding-cases] refresh %s schema failed: %v", t.Name(), err) + } + if previousReview == "true" { + continue + } if err := s.m.pg.SeedTool(t.Name(), t.Description(), schema, json.RawMessage(`["reporter"]`)); err != nil { return } @@ -393,11 +414,19 @@ func (s *Server) seedFindingCaseTools() { names = append(names, t.Name()) } } - if err := s.m.pg.AddAgentToToolBinding("reporter", names); err != nil { - return + if previousReview != "true" { + if err := s.m.pg.AddAgentToToolBinding("reporter", names); err != nil { + return + } + if err := s.m.pg.SetSetting("finding_case_tools_v2_review", "true"); err != nil { + return + } } _ = s.m.pg.SetSetting(flag, "true") } + +const findingCaseReviewTaskMessage = "人工选择整理（只读原始记录）：本任务所选 finding_ids=%s。逐条读取已有报告，仅判断同一实际根因；根因明确相同则直接归并。只有证据指向相同端点、处理器/后端组件、参数或缺失检查等具体共同点，但缺少确认共同根因的关键证据时，才提交一个覆盖2..20条同任务疑似成员的分组，不要拆成两两建议；仅类型或资产相同不足以建议，没有具体共同点时保持 independent，不要因为不确定就提交建议。证据表明属于不同代码路径时可判 independent。已有或新建组应尽力更新统一报告；统一报告暂缺或待更新不阻塞 complete_finding_case_review，可由后续运行或人工补齐。禁止生成或覆盖原始报告、改评级状态、补绑证据，不归并未选记录，不探测目标，遵守人工否决。结束前调用 complete_finding_case_review 提交逐条判断并通过验收。" + func (s *Server) reviewFindingCases(w http.ResponseWriter, r *http.Request) { var a findingCaseRequest if !decodeCaseRequest(w, r, &a) { @@ -463,7 +492,7 @@ func (s *Server) reviewFindingCases(w http.ResponseWriter, r *http.Request) { taskID, fids := reservation.TaskID, reservation.FindingIDs taskContext := taskContexts[taskID] raw, _ := json.Marshal(fids) - msg := fmt.Sprintf("人工选择整理（只读原始记录）：本任务所选 finding_ids=%s。逐条读取已有报告，仅判断同一实际根因，明确则归并，不确定提交建议；已有组更新统一报告。禁止生成或覆盖原始报告、改评级状态、补绑证据，不归并未选记录，不探测目标。结束前调用 complete_finding_case_review 提交逐条判断并通过验收。", raw) + msg := fmt.Sprintf(findingCaseReviewTaskMessage, raw) item := triggeredRun{agentKey: "reporter", title: "漏洞整理 · task#" + i64s(taskID), message: msg, taskID: taskID, taskDesc: taskContext.taskDesc, taskGoal: taskContext.taskGoal, conversationID: reservation.ConversationID} cfg := s.readTriggerBehavior("reporter") s.queueMu.Lock() diff --git a/server/finding_cases_test.go b/server/finding_cases_test.go index 09037d8..23be7b7 100644 --- a/server/finding_cases_test.go +++ b/server/finding_cases_test.go @@ -42,6 +42,119 @@ func testCaseServer(t *testing.T) (*Server, int64, []int64) { t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM findings WHERE task_id=$1`, task.ID); _ = pg.DeleteTask(task.ID) }) return &Server{m: &Manager{pg: pg}}, task.ID, ids } + +func TestFindingCaseToolMigrationPreservesBindings(t *testing.T) { + s, _, _ := testCaseServer(t) + pg := s.m.pg + for _, flag := range []string{"finding_case_tools_v1", "finding_case_tools_v2_review", "finding_case_tools_v3_review"} { + value, exists, err := pg.GetSetting(flag) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if exists { + if err := pg.SetSetting(flag, value); err != nil { + t.Error(err) + } + } else if _, err := pg.Exec(`DELETE FROM settings WHERE key=$1`, flag); err != nil { + t.Error(err) + } + }) + } + for _, tool := range s.findingCaseTools() { + key := tool.Name() + old, err := pg.GetTool(key) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if old == nil { + if _, err := pg.Exec(`DELETE FROM tools WHERE key=$1`, key); err != nil { + t.Error(err) + } + return + } + bindings, _ := json.Marshal(old.Agents) + if err := pg.UpdateTool(key, old.Description, old.Schema, bindings, old.Enabled); err != nil { + t.Error(err) + } + }) + if err := pg.SeedTool(key, tool.Description(), nil, nil); err != nil { + t.Fatal(err) + } + } + for _, tc := range []struct { + name string + v1, v2, fresh bool + }{ + {name: "fresh", fresh: true}, + {name: "initial"}, + {name: "v1_to_v3", v1: true}, + {name: "v2_to_v3", v1: true, v2: true}, + {name: "v2_without_v1_flag", v2: true}, + } { + t.Run(tc.name, func(t *testing.T) { + for flag, value := range map[string]bool{ + "finding_case_tools_v1": tc.v1, + "finding_case_tools_v2_review": tc.v2, + "finding_case_tools_v3_review": false, + } { + if value { + if err := pg.SetSetting(flag, "true"); err != nil { + t.Fatal(err) + } + } else if _, err := pg.Exec(`DELETE FROM settings WHERE key=$1`, flag); err != nil { + t.Fatal(err) + } + } + for _, tool := range s.findingCaseTools() { + if tc.fresh { + if _, err := pg.Exec(`DELETE FROM tools WHERE key=$1`, tool.Name()); err != nil { + t.Fatal(err) + } + } else if err := pg.UpdateTool(tool.Name(), "old description", json.RawMessage(`{}`), json.RawMessage(`["custom-agent"]`), false); err != nil { + t.Fatal(err) + } + } + s.seedFindingCaseTools() + for _, tool := range s.findingCaseTools() { + row, err := pg.GetTool(tool.Name()) + if err != nil || row == nil { + t.Fatal("missing tool", tool.Name(), err) + } + wantReporter := !tc.v2 && (!tc.v1 || tool.Name() == "complete_finding_case_review") + if contains(row.Agents, "reporter") != wantReporter || contains(row.Agents, "custom-agent") == tc.fresh || row.Enabled != tc.fresh { + t.Fatalf("binding/enabled configuration changed incorrectly: %+v", row) + } + if row.Description != tool.Description() || len(row.Schema) <= 2 { + t.Fatalf("description/schema not refreshed: %+v", row) + } + } + for _, flag := range []string{"finding_case_tools_v2_review", "finding_case_tools_v3_review"} { + if value, _, err := pg.GetSetting(flag); err != nil || value != "true" { + t.Fatalf("migration flag %s not saved: %q, %v", flag, value, err) + } + } + for _, tool := range s.findingCaseTools() { + if err := pg.RemoveAgentFromTool("reporter", tool.Name()); err != nil { + t.Fatal(err) + } + } + // A refresh retry must not replay a completed legacy binding migration. + if err := pg.SetSetting("finding_case_tools_v3_review", "false"); err != nil { + t.Fatal(err) + } + s.seedFindingCaseTools() + for _, tool := range s.findingCaseTools() { + row, err := pg.GetTool(tool.Name()) + if err != nil || row == nil || contains(row.Agents, "reporter") { + t.Fatal("refresh restored an operator-removed binding", row, err) + } + } + }) + } +} + func TestFindingCaseAPIPaginationAndTools(t *testing.T) { s, task, ids := testCaseServer(t) ctx := context.Background() diff --git a/server/manager.go b/server/manager.go index 09fdcd3..9b61e13 100644 --- a/server/manager.go +++ b/server/manager.go @@ -1,6 +1,7 @@ package server import ( + "context" "database/sql" "encoding/json" "errors" @@ -59,6 +60,7 @@ type Task struct { TimeoutSeconds int `json:"timeout_seconds"` PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒) CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定,默认开) + WorkerCount int `json:"worker_count"` // 本任务 Worker 数;0=用全局默认 FirstRunAt int64 `json:"first_run_at,omitempty"` DeadlineAt int64 `json:"deadline_at,omitempty"` Store *pgdb.ExplorationStore `json:"-"` @@ -265,6 +267,7 @@ const ( // only routes the search backend) and the per-profile LLM proxy. settingGlobalProxy = "global_proxy" settingWorkers = "workers" + settingWorkerPools = "worker_pools" settingLLMRecord = "llm_record" // LLM 轮询(故障转移)。默认关闭——开启后走「全局激活配置」的 agent 在当前配置 // 不可用(余额不足/key 失效/限流/服务异常)时自动切到下一个配置。 @@ -291,6 +294,102 @@ const ( defaultConcurrencyLimit = 5 ) +// WorkerPool binds a number of worker slots to one LLM profile. Profile IDs are +// encoded as JSON strings to match the frontend LLMProfile API. +type WorkerPool struct { + ProfileID int64 `json:"profile_id,string"` + Workers int `json:"workers"` +} + +func (p *WorkerPool) UnmarshalJSON(data []byte) error { + var raw struct { + ProfileID json.RawMessage `json:"profile_id"` + Workers int `json:"workers"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + var id int64 + if err := json.Unmarshal(raw.ProfileID, &id); err != nil { + var text string + if err := json.Unmarshal(raw.ProfileID, &text); err != nil { + return fmt.Errorf("profile_id 必须是整数或数字字符串") + } + parsed, err := strconv.ParseInt(strings.TrimSpace(text), 10, 64) + if err != nil { + return fmt.Errorf("profile_id 必须是整数或数字字符串") + } + id = parsed + } + p.ProfileID, p.Workers = id, raw.Workers + return nil +} + +// WorkerSlot is the concrete routing plan for one task run. +type WorkerSlot struct { + Name string + ProfileID int64 // zero means legacy task-router resolution +} + +// ValidateWorkerPools validates pool shape without requiring a database. Pool +// entries with duplicate profile IDs are rejected so each profile has one +// unambiguous slot group in the saved order. +func ValidateWorkerPools(pools []WorkerPool) error { + if len(pools) == 0 { + return nil + } + seen := make(map[int64]bool, len(pools)) + total := 0 + for _, pool := range pools { + if pool.ProfileID <= 0 { + return fmt.Errorf("worker_pools 中 profile_id 必须大于 0") + } + if pool.Workers <= 0 { + return fmt.Errorf("worker_pools 中 workers 必须大于 0") + } + if seen[pool.ProfileID] { + return fmt.Errorf("worker_pools 中 profile_id %d 重复", pool.ProfileID) + } + seen[pool.ProfileID] = true + if total > 50-pool.Workers { + return fmt.Errorf("worker_pools 的 Worker 总数必须在 1 到 50 之间") + } + total += pool.Workers + } + if total < 1 || total > 50 { + return fmt.Errorf("worker_pools 的 Worker 总数必须在 1 到 50 之间") + } + return nil +} + +// ExpandWorkerSlots expands pools in their saved order; invalid or empty legacy +// settings use the historical worker count (default 3). +func ExpandWorkerSlots(pools []WorkerPool, legacy int) ([]WorkerSlot, error) { + if err := ValidateWorkerPools(pools); err != nil { + return nil, err + } + if len(pools) == 0 { + if legacy <= 0 { + legacy = defaultWorkers + } + slots := make([]WorkerSlot, legacy) + for i := range slots { + slots[i] = WorkerSlot{Name: fmt.Sprintf("work#%d", i+1)} + } + return slots, nil + } + slots := make([]WorkerSlot, 0) + for _, pool := range pools { + for i := 0; i < pool.Workers; i++ { + slots = append(slots, WorkerSlot{ProfileID: pool.ProfileID}) + } + } + for i := range slots { + slots[i].Name = fmt.Sprintf("work#%d", i+1) + } + return slots, nil +} + // ConcurrencyLimit returns whether the simultaneous-running-task cap is enabled and // its limit (default 5 when enabled but unset). limit is always >=1 when enabled. func (m *Manager) ConcurrencyLimit() (enabled bool, limit int) { @@ -332,6 +431,64 @@ func (m *Manager) Workers() int { return n } +func (m *Manager) WorkerPools() []WorkerPool { + v, ok, err := m.pg.GetSetting(settingWorkerPools) + if err != nil || !ok || strings.TrimSpace(v) == "" { + return []WorkerPool{} + } + var pools []WorkerPool + if json.Unmarshal([]byte(v), &pools) != nil || ValidateWorkerPools(pools) != nil { + return []WorkerPool{} + } + if pools == nil { + return []WorkerPool{} + } + return pools +} + +// SetWorkerPools validates profile references at save time and atomically +// synchronizes the legacy workers setting with the pool total. An empty slice +// restores legacy mode; an optional legacy count is used when clearing pools. +func (m *Manager) SetWorkerPools(pools []WorkerPool, legacyWorkers ...int) error { + if pools == nil { + pools = []WorkerPool{} + } + if err := ValidateWorkerPools(pools); err != nil { + return err + } + if len(pools) > 0 { + for _, pool := range pools { + profile, err := m.pg.ProfileByID(pool.ProfileID) + if err != nil { + return err + } + if profile == nil { + return fmt.Errorf("worker_pools 中 profile_id %d 不存在", pool.ProfileID) + } + } + } + b, err := json.Marshal(pools) + if err != nil { + return err + } + workerCount := m.Workers() + if len(pools) > 0 { + workerCount = 0 + for _, pool := range pools { + workerCount += pool.Workers + } + } else if len(legacyWorkers) > 0 { + if legacyWorkers[0] <= 0 { + return fmt.Errorf("workers 必须 >0") + } + workerCount = legacyWorkers[0] + } + return m.pg.SetSettingsAtomic(context.Background(), map[string]string{ + settingWorkerPools: string(b), + settingWorkers: strconv.Itoa(workerCount), + }) +} + // SetWorkers persists the concurrent work-agent count. Values <=0 are rejected. func (m *Manager) SetWorkers(n int) error { if n <= 0 { @@ -849,8 +1006,8 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...), CompanyIDs: append([]int64(nil), pt.CompanyIDs...), TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds, - CoverageEnabled: pt.CoverageEnabled, - FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), + CoverageEnabled: pt.CoverageEnabled, WorkerCount: pt.WorkerCount, + FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1), } } diff --git a/server/orchestration.go b/server/orchestration.go index 3dd95c9..44eb0d6 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -524,7 +524,7 @@ func (s *Server) seedOrchestrationTools() { // reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag. // Bump the flag whenever these tools' schemas/descriptions change in code. func (s *Server) refreshBuiltinToolSchemas() { - const flag = "tool_schema_refresh_v7_list_facts_paging" + const flag = "tool_schema_refresh_v8_finding_case_tools" if v, _, _ := s.m.pg.GetSetting(flag); v == "true" { return } diff --git a/server/remarks.go b/server/remarks.go index c33ecc6..0864eb9 100644 --- a/server/remarks.go +++ b/server/remarks.go @@ -10,7 +10,7 @@ import ( "github.com/Autumn-27/artex/db" ) -const maxRemarkRunes = 2000 +const maxRemarkRunes = db.MaxTaskRemarkRunes func normalizeRemark(value string) (string, error) { value = strings.TrimSpace(value) diff --git a/server/repeater.go b/server/repeater.go new file mode 100644 index 0000000..3f157d9 --- /dev/null +++ b/server/repeater.go @@ -0,0 +1,416 @@ +package server + +import ( + "bytes" + "context" + "crypto/tls" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "syscall" + "time" +) + +const ( + repeaterMaxBody = 2 << 20 // 2 MiB, for both the JSON request and response preview + repeaterDefaultTimeout = 30_000 + repeaterMinTimeout = 1_000 + repeaterMaxTimeout = 120_000 + repeaterAllowLocalEnv = "ARTEX_REPEATER_ALLOW_LOCAL" +) + +var repeaterHopByHop = map[string]struct{}{ + "Connection": {}, "Keep-Alive": {}, "Proxy-Connection": {}, "Transfer-Encoding": {}, + "Te": {}, "Trailer": {}, "Upgrade": {}, "Content-Length": {}, "Proxy-Authorization": {}, +} + +type repeaterSendRequest struct { + Method *string `json:"method"` + URL string `json:"url"` + Headers map[string]string `json:"headers"` + Body string `json:"body"` + TimeoutMs int `json:"timeout_ms"` + InsecureSkipVerify bool `json:"insecure_skip_verify"` + UseRecordingProxy *bool `json:"use_recording_proxy"` +} + +type repeaterSendResponse struct { + Status int `json:"status"` + StatusText string `json:"status_text"` + Proto string `json:"proto"` + Headers map[string][]string `json:"headers"` + Body string `json:"body"` + BodyTruncated bool `json:"body_truncated"` + BodyBytes int `json:"body_bytes"` + ElapsedMs int64 `json:"elapsed_ms"` + URL string `json:"url"` + Error string `json:"error,omitempty"` +} + +func (s *Server) registerRepeater(mux *http.ServeMux) { + mux.HandleFunc("POST /api/repeater/send", s.sendRepeater) +} + +func (s *Server) sendRepeater(w http.ResponseWriter, r *http.Request) { + var in repeaterSendRequest + dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, repeaterMaxBody)) + if err := dec.Decode(&in); err != nil { + writeErr(w, http.StatusBadRequest, "请求体无效") + return + } + // Reject a second JSON value (or oversized trailing data) instead of silently + // accepting bytes beyond the documented 2 MiB request limit. + var extra any + if err := dec.Decode(&extra); err != io.EOF { + writeErr(w, http.StatusBadRequest, "请求体无效") + return + } + + method := http.MethodGet + if in.Method != nil { + method = strings.ToUpper(strings.TrimSpace(*in.Method)) + } + if method == "" { + writeErr(w, http.StatusBadRequest, "HTTP 方法不能为空") + return + } + switch method { + case http.MethodGet, http.MethodPost, http.MethodPut, http.MethodPatch, + http.MethodDelete, http.MethodHead, http.MethodOptions: + default: + writeErr(w, http.StatusBadRequest, "不支持的 HTTP 方法") + return + } + + rawURL := strings.TrimSpace(in.URL) + if rawURL == "" { + writeErr(w, http.StatusBadRequest, "url 必填") + return + } + u, err := url.Parse(rawURL) + if err != nil || u.Host == "" || u.Hostname() == "" || + !strings.EqualFold(u.Scheme, "http") && !strings.EqualFold(u.Scheme, "https") { + writeErr(w, http.StatusBadRequest, "仅支持绝对 http/https URL") + return + } + // This check is needed even when a proxy is used: the proxy must remain usable + // for normal targets, while obvious literal metadata/loopback URLs stay blocked. + if err := repeaterCheckURL(u); err != nil { + writeErr(w, http.StatusBadRequest, err.Error()) + return + } + + timeoutMs := in.TimeoutMs + if timeoutMs <= 0 { + timeoutMs = repeaterDefaultTimeout + } + if timeoutMs < repeaterMinTimeout { + timeoutMs = repeaterMinTimeout + } + if timeoutMs > repeaterMaxTimeout { + timeoutMs = repeaterMaxTimeout + } + useRecordingProxy := true + if in.UseRecordingProxy != nil { + useRecordingProxy = *in.UseRecordingProxy + } + + var body io.Reader + if in.Body != "" && method != http.MethodHead { + body = strings.NewReader(in.Body) + } + ctx, cancel := context.WithTimeout(r.Context(), time.Duration(timeoutMs)*time.Millisecond) + defer cancel() + req, err := http.NewRequestWithContext(ctx, method, u.String(), body) + if err != nil { + writeErr(w, http.StatusBadRequest, "构造请求失败") + return + } + for key, value := range in.Headers { + name := strings.TrimSpace(key) + if name == "" { + continue + } + if _, blocked := repeaterHopByHop[http.CanonicalHeaderKey(name)]; blocked { + continue + } + if strings.EqualFold(name, "Host") { + // Host is a special Request field, but retaining it in Header too makes + // the editor semantics explicit; net/http uses req.Host on the wire. + req.Host = strings.TrimSpace(value) + req.Header.Set("Host", strings.TrimSpace(value)) + continue + } + req.Header.Set(name, value) + } + + client := &http.Client{ + Timeout: time.Duration(timeoutMs) * time.Millisecond, + Transport: s.repeaterTransport(useRecordingProxy, in.InsecureSkipVerify), + CheckRedirect: func(next *http.Request, via []*http.Request) error { + if len(via) >= 10 { + return errors.New("重定向次数过多") + } + if err := repeaterCheckURL(next.URL); err != nil { + return err + } + if len(via) > 0 && !strings.EqualFold(next.URL.Host, via[0].URL.Host) { + return fmt.Errorf("拒绝跨主机重定向（%s → %s）", via[0].URL.Host, next.URL.Host) + } + return nil + }, + } + // Each send owns its transport. Registered before Body.Close so idle + // connections are released after the response body, including on errors. + defer client.CloseIdleConnections() + + started := time.Now() + resp, err := client.Do(req) + if err != nil { + writeJSON(w, http.StatusOK, repeaterSendResponse{ + Headers: map[string][]string{}, URL: u.String(), + ElapsedMs: time.Since(started).Milliseconds(), Error: repeaterRedactError(err), + }) + return + } + defer resp.Body.Close() + + out := repeaterSendResponse{ + Status: resp.StatusCode, StatusText: resp.Status, Proto: resp.Proto, + Headers: resp.Header.Clone(), URL: resp.Request.URL.String(), + } + if method != http.MethodHead { + // Read one byte beyond the cap so truncation is exact. The MVP keeps the + // preview as UTF-8 text; invalid bytes are replaced instead of base64. + data, readErr := io.ReadAll(io.LimitReader(resp.Body, repeaterMaxBody+1)) + if len(data) > repeaterMaxBody { + out.BodyTruncated = true + data = data[:repeaterMaxBody] + } + out.Body = repeaterBodyText(data) + out.BodyBytes = len(data) + if readErr != nil && !errors.Is(readErr, io.EOF) { + out.Error = "读取响应失败: " + readErr.Error() + } + } + out.ElapsedMs = time.Since(started).Milliseconds() + writeJSON(w, http.StatusOK, out) +} + +// repeaterBodyText keeps the response preview as text for the MVP. Invalid +// UTF-8 and NUL bytes are replaced so a binary response cannot inject control +// characters into the JSON/UI raw view; the original bytes are not base64 +// encoded and the preview remains bounded by repeaterMaxBody. +func repeaterBodyText(data []byte) string { + text := string(bytes.ToValidUTF8(data, []byte("�"))) + return strings.ReplaceAll(text, "\x00", "�") +} + +func (s *Server) repeaterTransport(useRecordingProxy, insecure bool) http.RoundTripper { + var transport *http.Transport + proxied := false + if useRecordingProxy && s != nil && s.m != nil { + transport = s.httpProxyTransport(httpExec{UseRecordingProxy: true}) + proxied = transport != nil + } + if transport == nil && s != nil && s.m != nil && strings.TrimSpace(s.m.GlobalProxy()) != "" { + transport = s.httpProxyTransport(httpExec{Proxy: s.m.GlobalProxy()}) + proxied = transport != nil + } + if transport == nil { + if base, ok := http.DefaultTransport.(*http.Transport); ok { + transport = base.Clone() + } else { + transport = &http.Transport{} + } + // "Direct" means direct even if the process inherited HTTP_PROXY. The + // configured recording/global proxy paths are handled above explicitly. + transport.Proxy = nil + } + transport = transport.Clone() + // Keep response bytes faithful and make body_bytes match what was received; + // transparent gzip decompression would otherwise obscure both values. + transport.DisableCompression = true + // Use the request's timeout_ms budget instead of inherited phase deadlines. + transport.TLSHandshakeTimeout = 0 + transport.ExpectContinueTimeout = 0 + if !proxied { + // Clear a custom TLS dialer before installing the guarded dialer; when + // DialTLSContext is set, net/http bypasses DialContext entirely. + transport.DialTLSContext = nil + transport.DialTLS = nil + transport.Dial = nil + // Let the request context enforce timeout_ms for the entire dial. + transport.DialContext = (&net.Dialer{ + KeepAlive: 30 * time.Second, + Control: repeaterBlockSensitiveDial, + }).DialContext + } + if insecure { + cfg := &tls.Config{InsecureSkipVerify: true} //nolint:gosec // explicit user option for pentest targets + if transport.TLSClientConfig != nil { + cfg = transport.TLSClientConfig.Clone() + cfg.InsecureSkipVerify = true + } + transport.TLSClientConfig = cfg + } + return transport +} + +func repeaterAllowLocal() bool { + v := strings.TrimSpace(os.Getenv(repeaterAllowLocalEnv)) + return v == "1" || strings.EqualFold(v, "true") +} + +func repeaterCheckURL(u *url.URL) error { + if repeaterAllowLocal() { + return nil + } + host := strings.TrimSuffix(strings.TrimSpace(u.Hostname()), ".") + // URL.Hostname may retain an IPv6 zone (for example fe80::1%eth0); the + // zone is not part of the address classification. + if zone := strings.IndexByte(host, '%'); zone >= 0 { + host = host[:zone] + } + ip := net.ParseIP(host) + if ip == nil && repeaterLooksNumericHost(host) { + var ok bool + ip, ok = repeaterParseLegacyIPv4(host) + if !ok { + return fmt.Errorf("拒绝无法解析的数字主机地址 %q", host) + } + } + if strings.EqualFold(host, "localhost") || (ip != nil && repeaterIsBlockedDialIP(ip)) { + return fmt.Errorf("拒绝访问本机/链路本地地址（如确需访问本机服务，设置 %s=1）", repeaterAllowLocalEnv) + } + return nil +} + +// repeaterLooksNumericHost identifies the legacy numeric spellings accepted by +// common proxy/resolver implementations (decimal, octal, hexadecimal and +// dotted inet_aton forms). A numeric-looking value that fails to parse is +// rejected instead of being handed to a proxy for its own interpretation. +func repeaterLooksNumericHost(host string) bool { + if host == "" || host[0] < '0' || host[0] > '9' { + return false + } + for i := 0; i < len(host); i++ { + c := host[i] + switch { + case c == '.': + case c >= '0' && c <= '9': + case c >= 'a' && c <= 'f', c >= 'A' && c <= 'F': + case c == 'x' || c == 'X': + default: + return false + } + } + return true +} + +// repeaterParseLegacyIPv4 follows inet_aton's 1-, 2-, 3- and 4-part IPv4 +// notation. Parts with a leading zero are octal and parts with 0x are hex. +func repeaterParseLegacyIPv4(host string) (net.IP, bool) { + parts := strings.Split(host, ".") + if len(parts) < 1 || len(parts) > 4 { + return nil, false + } + values := make([]uint64, len(parts)) + for i, part := range parts { + if part == "" { + return nil, false + } + base := 10 + digits := part + if len(digits) > 2 && digits[0] == '0' && (digits[1] == 'x' || digits[1] == 'X') { + base, digits = 16, digits[2:] + } else if len(digits) > 1 && digits[0] == '0' { + base = 8 + } + if digits == "" { + return nil, false + } + value, err := strconv.ParseUint(digits, base, 32) + if err != nil { + return nil, false + } + values[i] = value + } + var value uint64 + switch len(values) { + case 1: + value = values[0] + case 2: + if values[0] > 0xff || values[1] > 0xffffff { + return nil, false + } + value = values[0]<<24 | values[1] + case 3: + if values[0] > 0xff || values[1] > 0xff || values[2] > 0xffff { + return nil, false + } + value = values[0]<<24 | values[1]<<16 | values[2] + case 4: + for _, part := range values { + if part > 0xff { + return nil, false + } + } + value = values[0]<<24 | values[1]<<16 | values[2]<<8 | values[3] + } + if value > 0xffffffff { + return nil, false + } + return net.IPv4(byte(value>>24), byte(value>>16), byte(value>>8), byte(value)), true +} + +func repeaterIsBlockedDialIP(ip net.IP) bool { + if ip == nil { + return true + } + if v4 := ip.To4(); v4 != nil { + ip = v4 + } + return ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || + ip.IsInterfaceLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast() +} + +func repeaterBlockSensitiveDial(_, address string, _ syscall.RawConn) error { + if repeaterAllowLocal() { + return nil + } + host, _, err := net.SplitHostPort(address) + if err != nil { + return err + } + ip := net.ParseIP(host) + if ip == nil { + return fmt.Errorf("无法解析目标地址 %q", host) + } + if repeaterIsBlockedDialIP(ip) { + return fmt.Errorf("拒绝访问本机/链路本地地址 %s（如确需访问本机服务，设置 %s=1）", ip, repeaterAllowLocalEnv) + } + return nil +} + +func repeaterRedactError(err error) string { + var urlErr *url.Error + if errors.As(err, &urlErr) { + host := "" + if parsed, parseErr := url.Parse(urlErr.URL); parseErr == nil { + host = parsed.Host + } + if urlErr.Err != nil { + return fmt.Sprintf("%s %s: %s", urlErr.Op, host, urlErr.Err) + } + return fmt.Sprintf("%s %s: 未知错误", urlErr.Op, host) + } + return err.Error() +} diff --git a/server/repeater_test.go b/server/repeater_test.go new file mode 100644 index 0000000..608f04d --- /dev/null +++ b/server/repeater_test.go @@ -0,0 +1,404 @@ +package server + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/traffic" +) + +func repeaterTestSend(t *testing.T, s *Server, payload any) (int, repeaterSendResponse, string) { + t.Helper() + raw, err := json.Marshal(payload) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest(http.MethodPost, "/api/repeater/send", bytes.NewReader(raw)) + w := httptest.NewRecorder() + s.sendRepeater(w, r) + var out repeaterSendResponse + _ = json.Unmarshal(w.Body.Bytes(), &out) + return w.Code, out, w.Body.String() +} + +func TestRepeaterRejectsInvalidInput(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "") + s := &Server{} + for _, tc := range []struct { + name string + body string + }{ + {"missing url", `{"method":"GET"}`}, + {"empty method", `{"method":" ","url":"https://example.test/"}`}, + {"bad scheme", `{"method":"GET","url":"file:///etc/passwd"}`}, + {"gopher", `{"method":"GET","url":"gopher://example.test/1"}`}, + {"unknown method", `{"method":"TRACE","url":"http://example.test/"}`}, + } { + t.Run(tc.name, func(t *testing.T) { + r := httptest.NewRequest(http.MethodPost, "/api/repeater/send", strings.NewReader(tc.body)) + w := httptest.NewRecorder() + s.sendRepeater(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("status=%d want 400 body=%s", w.Code, w.Body.String()) + } + }) + } +} + +func TestRepeaterBlockedIPs(t *testing.T) { + for _, raw := range []string{"127.0.0.1", "::1", "169.254.169.254", "0.0.0.0", "224.0.0.1"} { + if !repeaterIsBlockedDialIP(net.ParseIP(raw)) { + t.Errorf("expected %s blocked", raw) + } + } + for _, raw := range []string{"8.8.8.8", "10.0.0.1", "192.168.1.20", "172.16.0.5"} { + if repeaterIsBlockedDialIP(net.ParseIP(raw)) { + t.Errorf("expected %s allowed", raw) + } + } + + t.Setenv(repeaterAllowLocalEnv, "") + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("loopback handler should not be reached") + })) + defer srv.Close() + code, _, body := repeaterTestSend(t, &Server{}, map[string]any{"url": srv.URL, "use_recording_proxy": false}) + if code != http.StatusBadRequest || !strings.Contains(body, "拒绝") { + t.Fatalf("loopback status=%d body=%s", code, body) + } +} + +func TestRepeaterRejectsLegacyNumericHostsIncludingProxyPath(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "") + for _, raw := range []string{ + "http://2130706433/", // decimal 127.0.0.1 + "http://0177.0.0.1/", // octal 127.0.0.1 + "http://127.1/", // two-part inet_aton form + "http://0x7f000001/", // hexadecimal 127.0.0.1 + "http://9999999999/", // numeric-looking but invalid + } { + u, err := url.Parse(raw) + if err != nil { + t.Fatal(err) + } + if err := repeaterCheckURL(u); err == nil { + t.Errorf("repeaterCheckURL(%q) allowed a legacy numeric host", raw) + } + } + + // The check must run before a configured proxy gets a chance to interpret + // the legacy spelling and connect to loopback remotely. + proxy := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Error("proxy received a blocked numeric target") + })) + defer proxy.Close() + s := &Server{m: &Manager{globalProxy: proxy.URL}} + code, _, body := repeaterTestSend(t, s, map[string]any{"url": "http://2130706433/"}) + if code != http.StatusBadRequest || !strings.Contains(body, "拒绝") { + t.Fatalf("proxied legacy host status=%d body=%s", code, body) + } +} + +func TestRepeaterBlockSensitiveDial(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "") + for _, address := range []string{"127.0.0.1:80", "169.254.169.254:80"} { + if err := repeaterBlockSensitiveDial("tcp", address, nil); err == nil { + t.Errorf("repeaterBlockSensitiveDial(%q) allowed blocked address", address) + } + } +} + +func TestRepeaterTransportBlocksSensitiveDial(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "") + tr := (&Server{}).repeaterTransport(false, false).(*http.Transport) + defer tr.CloseIdleConnections() + if tr.DialContext == nil { + t.Fatal("direct transport has no guarded DialContext") + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + conn, err := tr.DialContext(ctx, "tcp", "127.0.0.1:1") + if conn != nil { + conn.Close() + } + if err == nil || !strings.Contains(err.Error(), "拒绝访问本机/链路本地地址 127.0.0.1") { + t.Fatalf("dial should fail at the sensitive-address guard, got %v", err) + } +} + +func TestRepeaterDirectHTTP2(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.ProtoMajor != 2 { + t.Errorf("request protocol=%s, want HTTP/2", r.Proto) + } + _, _ = io.WriteString(w, "h2 response") + })) + srv.EnableHTTP2 = true + srv.TLS = &tls.Config{NextProtos: []string{"h2"}} + srv.StartTLS() + defer srv.Close() + + code, out, raw := repeaterTestSend(t, &Server{}, map[string]any{ + "url": srv.URL, "use_recording_proxy": false, "insecure_skip_verify": true, + }) + if code != http.StatusOK || out.Error != "" || out.Status != http.StatusOK || out.Proto != "HTTP/2.0" || out.Body != "h2 response" { + t.Fatalf("HTTP/2 response status=%d raw=%s", code, raw) + } +} + +func TestRepeaterTLSHandshakeUsesRequestTimeout(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + // Scale down the inherited timeout so a slow handshake can be tested + // without waiting for the default transport's ten-second deadline. + base := http.DefaultTransport + short := base.(*http.Transport).Clone() + short.TLSHandshakeTimeout = 10 * time.Millisecond + http.DefaultTransport = short + defer func() { http.DefaultTransport = base }() + defer short.CloseIdleConnections() + + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "handshake completed") + })) + srv.TLS = &tls.Config{GetConfigForClient: func(*tls.ClientHelloInfo) (*tls.Config, error) { + time.Sleep(100 * time.Millisecond) + return nil, nil + }} + srv.StartTLS() + defer srv.Close() + + code, out, raw := repeaterTestSend(t, &Server{}, map[string]any{ + "url": srv.URL, "use_recording_proxy": false, "insecure_skip_verify": true, "timeout_ms": 1000, + }) + if code != http.StatusOK || out.Error != "" || out.Status != http.StatusOK || out.Body != "handshake completed" { + t.Fatalf("handshake within request budget failed: status=%d raw=%s", code, raw) + } +} + +func TestRepeaterClosesIdleConnections(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + for _, mode := range []string{"direct", "http2", "global proxy", "recording proxy fallback"} { + t.Run(mode, func(t *testing.T) { + closed := make(chan struct{}, 1) + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "response") + })) + srv.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateClosed { + closed <- struct{}{} + } + } + if mode == "http2" { + srv.EnableHTTP2 = true + srv.StartTLS() + } else { + srv.Start() + } + defer srv.Close() + + s := &Server{} + payload := map[string]any{"url": srv.URL, "use_recording_proxy": false, "insecure_skip_verify": true} + if mode == "global proxy" || mode == "recording proxy fallback" { + s.m = &Manager{globalProxy: srv.URL} + payload["url"] = "http://target.example/" + if mode == "recording proxy fallback" { + delete(payload, "use_recording_proxy") + } + } + code, out, raw := repeaterTestSend(t, s, payload) + if code != http.StatusOK || out.Error != "" || out.Status != http.StatusOK || out.Body != "response" { + t.Fatalf("status=%d raw=%s", code, raw) + } + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("connection still open after sendRepeater returned") + } + }) + } +} + +func TestRepeaterRoundTripHeadersAndBody(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + var got http.Request + var gotBody string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = *r + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + w.Header().Set("X-Echo", "ok") + w.WriteHeader(http.StatusCreated) + _, _ = io.WriteString(w, "hello") + })) + defer srv.Close() + + code, out, body := repeaterTestSend(t, &Server{}, map[string]any{ + "method": "POST", "url": srv.URL, "headers": map[string]string{ + "X-Test": "yes", "Host": "custom.example", "Connection": "keep-alive", + "Keep-Alive": "timeout=5", "Transfer-Encoding": "chunked", "Content-Length": "99", + }, "body": "payload", "use_recording_proxy": false, + }) + if code != http.StatusOK || out.Error != "" { + t.Fatalf("status=%d out=%+v raw=%s", code, out, body) + } + if out.Status != http.StatusCreated || out.Body != "hello" || out.BodyBytes != 5 { + t.Fatalf("response=%+v", out) + } + if out.Headers["X-Echo"][0] != "ok" || got.Method != http.MethodPost || gotBody != "payload" || got.Host != "custom.example" { + t.Fatalf("request method=%s host=%s body=%q headers=%v", got.Method, got.Host, gotBody, got.Header) + } + for _, name := range []string{"Connection", "Keep-Alive", "Transfer-Encoding"} { + if got.Header.Get(name) != "" { + t.Errorf("hop-by-hop header %s was forwarded: %q", name, got.Header.Get(name)) + } + } + if got.Header.Get("Content-Length") != "7" { + t.Errorf("content length should match body, got %q", got.Header.Get("Content-Length")) + } +} + +func TestRepeaterTruncatesAndHonoursHEAD(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + large := strings.Repeat("x", repeaterMaxBody+1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + _, _ = io.WriteString(w, "should-not-be-read") + return + } + _, _ = io.WriteString(w, large) + })) + defer srv.Close() + _, out, _ := repeaterTestSend(t, &Server{}, map[string]any{"url": srv.URL, "use_recording_proxy": false}) + if !out.BodyTruncated || out.BodyBytes != repeaterMaxBody || len(out.Body) != repeaterMaxBody { + t.Fatalf("truncation=%+v", out) + } + _, head, _ := repeaterTestSend(t, &Server{}, map[string]any{"method": "HEAD", "url": srv.URL, "use_recording_proxy": false}) + if head.Body != "" || head.BodyBytes != 0 { + t.Fatalf("HEAD body=%q bytes=%d", head.Body, head.BodyBytes) + } +} + +func TestRepeaterBodyTextReplacesBinaryControls(t *testing.T) { + if got := repeaterBodyText([]byte{'a', 0, 0xff, 'b'}); got != "a��b" { + t.Fatalf("body text=%q", got) + } +} + +func TestRepeaterRejectsCrossHostRedirect(t *testing.T) { + t.Setenv(repeaterAllowLocalEnv, "1") + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "target") })) + defer target.Close() + redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, target.URL, http.StatusFound) + })) + defer redirect.Close() + _, out, _ := repeaterTestSend(t, &Server{}, map[string]any{"url": redirect.URL, "use_recording_proxy": false}) + if out.Status != 0 || !strings.Contains(out.Error, "跨主机") { + t.Fatalf("redirect response=%+v", out) + } +} + +func TestRepeaterProxyTransportConstruction(t *testing.T) { + s := &Server{m: &Manager{}} + for _, scheme := range []string{"http", "https", "socks5", "socks5h"} { + tr := s.httpProxyTransport(httpExec{Proxy: scheme + "://user:pass@127.0.0.1:1080"}) + if tr == nil || tr.Proxy == nil { + t.Fatalf("%s proxy transport not configured: %#v", scheme, tr) + } + proxyURL, err := tr.Proxy(httptest.NewRequest(http.MethodGet, "https://target.example/", nil)) + if err != nil || proxyURL == nil || proxyURL.Scheme != scheme || proxyURL.Host != "127.0.0.1:1080" { + t.Fatalf("%s proxy URL = %#v, err=%v", scheme, proxyURL, err) + } + if user := proxyURL.User; user == nil || user.Username() != "user" { + t.Fatalf("%s proxy credentials missing: %v", scheme, proxyURL.User) + } + password, ok := proxyURL.User.Password() + if !ok || password != "pass" { + t.Fatalf("%s proxy password missing: %q (ok=%v)", scheme, password, ok) + } + } + + // With capture disabled, the recording-proxy option must preserve the + // configured global proxy URL instead of treating it as a host:port. + s.m.globalProxy = "socks5://user:pass@127.0.0.1:1080" + tr := s.httpProxyTransport(httpExec{UseRecordingProxy: true}) + if tr == nil || tr.Proxy == nil { + t.Fatal("global SOCKS5 proxy was not configured") + } + proxyURL, err := tr.Proxy(httptest.NewRequest(http.MethodGet, "https://target.example/", nil)) + if err != nil || proxyURL == nil || proxyURL.Scheme != "socks5" { + t.Fatalf("global proxy URL = %#v, err=%v", proxyURL, err) + } + + // Capture mode supplies a complete HTTP URL from Traffic.ProxyAddr. The + // helper must preserve that scheme and host rather than prefixing http:// a + // second time. + captured, err := traffic.Open(t.TempDir(), "127.0.0.1:18788") + if err != nil { + t.Fatalf("traffic.Open: %v", err) + } + defer captured.Close() + s.m.traffic = captured + s.m.trafficOn = true + tr = s.httpProxyTransport(httpExec{UseRecordingProxy: true}) + if tr == nil || tr.Proxy == nil { + t.Fatal("recording proxy was not configured") + } + proxyURL, err = tr.Proxy(httptest.NewRequest(http.MethodGet, "https://target.example/", nil)) + if err != nil || proxyURL == nil || proxyURL.Scheme != "http" || proxyURL.Host != "127.0.0.1:18788" { + t.Fatalf("recording proxy URL = %#v, err=%v", proxyURL, err) + } + caFile := s.m.ProxyCACert() + if caFile == "" { + t.Fatal("recording proxy CA path is empty") + } + caPEM, err := os.ReadFile(caFile) + if err != nil { + t.Fatalf("read recording proxy CA: %v", err) + } + wantRoots := x509.NewCertPool() + if !wantRoots.AppendCertsFromPEM(caPEM) { + t.Fatal("recording proxy CA is not valid PEM") + } + if tr.TLSClientConfig == nil || tr.TLSClientConfig.RootCAs == nil || !tr.TLSClientConfig.RootCAs.Equal(wantRoots) { + t.Fatal("recording proxy CA was not loaded into transport RootCAs") + } + if s.repeaterTransport(true, false).(*http.Transport).ForceAttemptHTTP2 { + t.Fatal("recording proxy transport must not force HTTP/2") + } +} + +func TestRepeaterRedactProxyCredentials(t *testing.T) { + for _, scheme := range []string{"http", "https", "socks5", "socks5h"} { + t.Run(scheme, func(t *testing.T) { + proxyErr := &url.Error{ + Op: "proxyconnect", URL: scheme + "://proxy-user:proxy-secret@127.0.0.1:1080", + Err: errors.New("connection refused"), + } + for _, err := range []error{proxyErr, fmt.Errorf("send request: %w", proxyErr)} { + got := repeaterRedactError(err) + if strings.Contains(got, "proxy-secret") || strings.Contains(got, "proxy-user") { + t.Fatalf("proxy credentials leaked: %q", got) + } + if got != "proxyconnect 127.0.0.1:1080: connection refused" { + t.Fatalf("unexpected redacted error: %q", got) + } + } + }) + } +} diff --git a/server/server.go b/server/server.go index 969216c..672069e 100644 --- a/server/server.go +++ b/server/server.go @@ -177,6 +177,9 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi b := s.agentsForTask(t) return b.pl, b.wk }) + s.engine.SetWorkerResolver(func(t *Task, slot WorkerSlot) (*agent.Worker, error) { + return s.workerForTaskSlot(t, slot) + }) // Global readiness (Ready()/llm_configured): a global active LLM provider is // installed. Task-level runnability is separate (ReadyFor → the resolver above). s.engine.SetReadiness(func() bool { @@ -575,7 +578,8 @@ func (s *Server) chatUnavailableReason() string { // providerForProfile returns a cached provider+cfg for a profile id, so every agent // bound/pinned to the same profile shares one provider instance (one rate limiter). -// ok=false when the profile is missing/invalid → caller falls back to the global pair. +// ok=false when the profile is missing/invalid; pinned callers surface that error +// while legacy callers retain their existing fallback behavior. func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) { s.provCacheMu.Lock() if e := s.provByProfile[id]; e != nil { @@ -824,6 +828,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats) mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport) s.registerFindingTraffic(mux) + s.registerRepeater(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) @@ -1540,6 +1545,8 @@ type createTaskReq struct { PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600 SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。 CoverageEnabled *bool `json:"coverage_enabled,omitempty"` // 资产覆盖度功能;省略/null=默认开(true)。false=关闭覆盖度计算/展示/自动累积范围+隐藏 add_task_scope/list_untested_assets。company 关联不受影响。 + Remark string `json:"remark,omitempty"` // 任务备注(最多2000字);空=无 + WorkerCount int `json:"worker_count,omitempty"` // 本任务 Worker 数;0/省略=用全局默认;上限50 // InterceptRules 任务级资产拦截/允许规则(创建时录入,存 task_intercept_rules,不进全局表)。 InterceptRules []taskInterceptRuleReq `json:"intercept_rules,omitempty"` } @@ -1603,12 +1610,27 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { } schedule = &input } + remark, err := normalizeRemark(req.Remark) + if err != nil { + writeErr(w, 400, err.Error()) + return + } + if req.WorkerCount < 0 { + writeErr(w, 400, "Worker 运行数量不能为负数") + return + } + if req.WorkerCount > db.MaxTaskWorkerCount { + writeErr(w, 400, fmt.Sprintf("Worker 运行数量最多 %d", db.MaxTaskWorkerCount)) + return + } t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{ Schedule: schedule, Name: strings.TrimSpace(req.Name), CategoryID: req.CategoryID, SourceTaskIDs: sourceIDs, CompanyIDs: req.CompanyIDs, LLMProfileIDs: req.LLMProfileIDs, TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds, CoverageEnabled: req.CoverageEnabled, + Remark: remark, + WorkerCount: req.WorkerCount, InterceptRules: interceptRules, }) if err != nil { @@ -3432,6 +3454,7 @@ func (s *Server) settingsPayload() map[string]any { "global_proxy": s.m.GlobalProxy(), // 全局出口代理(http/https/socks5)，所有目标流量走它，空=直连 "python_interpreter": strings.TrimSpace(pyStored), // 用户/自动设的值(空=用运行时检测) "workers": s.m.Workers(), // 并发工作 agent 数(默认3)；对之后启动的任务生效 + "worker_pools": s.m.WorkerPools(), // 非空时按池顺序绑定 worker 会话 "task_concurrency_enabled": concOn, // 任务并发上限开关(默认关) "task_concurrency_limit": concLimit, // 同时运行任务上限(开启后默认5) // LLM 轮询(故障转移)。默认关；开启后走全局激活配置的 agent 在当前配置不可用时 @@ -3499,14 +3522,15 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { 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"` - WebSearchBackend *string `json:"web_search_backend"` - BraveKey *string `json:"brave_search_api_key"` - TavilyKey *string `json:"tavily_search_api_key"` - WebSearchProxy *string `json:"web_search_proxy"` // 独立出口代理(http/https/socks5)；null=不改，""=清空 - GlobalProxy *string `json:"global_proxy"` // 全局出口代理(http/https/socks5)；null=不改，""=清空(直连) - PythonInterp *string `json:"python_interpreter"` // 自定义脚本工具的 python 解释器路径 - Workers *int `json:"workers"` // 并发工作 agent 数(>0)；对之后启动的任务生效 + WebSearchEnabled *bool `json:"web_search_enabled"` + WebSearchBackend *string `json:"web_search_backend"` + BraveKey *string `json:"brave_search_api_key"` + TavilyKey *string `json:"tavily_search_api_key"` + WebSearchProxy *string `json:"web_search_proxy"` // 独立出口代理(http/https/socks5)；null=不改，""=清空 + GlobalProxy *string `json:"global_proxy"` // 全局出口代理(http/https/socks5)；null=不改，""=清空(直连) + PythonInterp *string `json:"python_interpreter"` // 自定义脚本工具的 python 解释器路径 + Workers *int `json:"workers"` // 并发工作 agent 数(>0)；对之后启动的任务生效 + WorkerPools *[]WorkerPool `json:"worker_pools"` // 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。 ConcurrencyEnabled *bool `json:"task_concurrency_enabled"` ConcurrencyLimit *int `json:"task_concurrency_limit"` @@ -3579,11 +3603,32 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { return } } - if req.Workers != nil { - if err := s.m.SetWorkers(*req.Workers); err != nil { + // A non-empty pool list takes precedence when both fields are present and + // synchronizes workers to its total. An empty list clears pools; when workers + // is also present it sets the legacy count for the newly restored legacy mode. + if req.WorkerPools != nil { + if len(*req.WorkerPools) == 0 && req.Workers != nil && *req.Workers <= 0 { + writeErr(w, 400, "workers 必须 >0") + return + } + legacyWorkers := []int(nil) + if req.Workers != nil { + legacyWorkers = []int{*req.Workers} + } + if err := s.m.SetWorkerPools(*req.WorkerPools, legacyWorkers...); err != nil { writeErr(w, 400, err.Error()) return } + } else if req.Workers != nil { + // An omitted worker_pools field leaves the saved mode unchanged. If a + // non-empty pool configuration is already saved, its total remains the + // authoritative concurrency and workers is deliberately ignored. + if len(s.m.WorkerPools()) == 0 { + if err := s.m.SetWorkers(*req.Workers); err != nil { + writeErr(w, 400, err.Error()) + return + } + } } if req.ConcurrencyEnabled != nil || req.ConcurrencyLimit != nil { // 部分 PUT:未给的字段用当前值兜底,避免只改一个把另一个重置。 diff --git a/server/task_llm.go b/server/task_llm.go index 2cae9ce..8063a3e 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -27,6 +27,7 @@ type taskAgentBundle struct { pl *agent.Planner wk *agent.Worker main *agent.MainAgent + pinnedWorkers map[int64]*agent.Worker } type llmAuditProfile struct { @@ -48,9 +49,10 @@ type llmActivityMetadata struct { } type taskLLMRuntime struct { - s *Server - taskID string - agentKey string + s *Server + taskID string + agentKey string + pinProfileID *int64 } type taskLLMError struct { @@ -101,13 +103,24 @@ type taskLLMStreamHooks struct { transition func(taskLLMSelection, db.TaskLLMTransition, error) } -// current resolves the provider this role runs on, by precedence: -// Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置。 +// current resolves the provider this role runs on. Pinned workers use only their +// assigned profile; legacy task agents use Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置. // 绑定优先于任务链：某个角色被显式指定了模型，就一直跑在那个模型上；绑定不存在或 // 构建失败时才降级到任务链，任务链为空时再降级到全局配置。 // 返回的 profile id 只有走任务链时才非零 —— streamTaskLLM 以此判断额度错误是否 // 应该推进任务的故障转移状态（绑定/全局路径不改任务链状态，沿用既有语义）。 func (r *taskLLMRuntime) current() (taskLLMSelection, error) { + if r.pinProfileID != nil { + // Pinned worker runtimes are exclusive: resolve their profile before + // touching bindings or the task failover state. A missing or invalid + // profile is an error, never an invitation to use another model. + prov, cfg, ok := r.s.providerForProfile(*r.pinProfileID) + if !ok { + return taskLLMSelection{}, fmt.Errorf("LLM profile #%d is missing or invalid", *r.pinProfileID) + } + t, _ := r.s.m.Task(r.taskID) + return taskLLMSelection{task: t, provider: prov, retry: cfg.Retry}, nil + } taskNum, err := parseTaskID(r.taskID) if err != nil { return taskLLMSelection{}, err @@ -152,6 +165,10 @@ func (r *taskLLMRuntime) current() (taskLLMSelection, error) { // the caller. If failover switches profiles, the change takes effect on the next // agent run (a fresh Session is built per run in captureRun). func (r *taskLLMRuntime) activeCfg() (agent.Config, bool) { + if r.pinProfileID != nil { + _, cfg, ok := r.s.providerForProfile(*r.pinProfileID) + return cfg, ok + } taskNum, err := parseTaskID(r.taskID) if err != nil { return agent.Config{}, false @@ -454,6 +471,12 @@ func isRetryableStreamError(err error) bool { // CompactionWindow mirrors current()'s precedence so the context window always // matches the provider the role will actually stream on. func (r *taskLLMRuntime) CompactionWindow() int { + if r.pinProfileID != nil { + if _, cfg, ok := r.s.providerForProfile(*r.pinProfileID); ok { + return cfg.CompactionWindow() + } + return (agent.Config{}).CompactionWindow() + } if _, cfg, ok := r.s.agentBindingProvider(r.agentKey); ok { return cfg.CompactionWindow() } @@ -597,12 +620,54 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { main.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏 bundle := &taskAgentBundle{ runtime: goalRuntime, plannerRuntime: plannerRuntime, workerRuntime: workerRuntime, - mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, + mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, pinnedWorkers: map[int64]*agent.Worker{}, } s.taskAgents[t.ID] = bundle return bundle } +func (s *Server) workerForTaskSlot(t *Task, slot WorkerSlot) (*agent.Worker, error) { + if slot.ProfileID <= 0 { + return nil, errors.New("worker slot is not pinned") + } + _, cfg, ok := s.providerForProfile(slot.ProfileID) + if !ok { + return nil, fmt.Errorf("LLM profile #%d is missing or invalid", slot.ProfileID) + } + s.taskAgentMu.Lock() + bundle := s.taskAgents[t.ID] + s.taskAgentMu.Unlock() + if bundle == nil { + bundle = s.agentsForTask(t) + } + s.taskAgentMu.Lock() + defer s.taskAgentMu.Unlock() + if current := s.taskAgents[t.ID]; current != nil { + bundle = current + } + if bundle.pinnedWorkers == nil { + bundle.pinnedWorkers = make(map[int64]*agent.Worker) + } + if worker := bundle.pinnedWorkers[slot.ProfileID]; worker != nil { + return worker, nil + } + id := slot.ProfileID + runtime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "worker", pinProfileID: &id} + tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) + worker := agent.NewWorker(runtime, cfg.Model, s.m.dir, tx, runtime.CompactionWindow(), s.agentMaxTurns("worker")) + worker.SetFindingRecorder(s.evidenceStore()) + worker.SetCompactionWindowResolver(runtime.CompactionWindow) + worker.SetNonStreaming(runtime.nonStreaming) + worker.SetMaxTokens(runtime.maxTokens) + worker.SetNoaEnabled(s.m.NoaCompactionEnabled) + worker.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) + worker.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + worker.SetWebSearch(s.webSearchFor("worker")) + worker.SetConstraintInject(s.constraintInjectWorker) + bundle.pinnedWorkers[id] = worker + return worker, nil +} + func (s *Server) syncTaskLLMState(pt *db.Task) { if pt == nil { return diff --git a/server/task_templates.go b/server/task_templates.go index 7125177..f94de47 100644 --- a/server/task_templates.go +++ b/server/task_templates.go @@ -22,8 +22,10 @@ type taskTemplateRequest struct { Name *string `json:"name"` Description *string `json:"description"` Goal *string `json:"goal"` + Remark *string `json:"remark"` CategoryID *int64 `json:"category_id"` InterceptRules []taskInterceptRuleReq `json:"intercept_rules"` + WorkerCount *int `json:"worker_count"` } // decodeTaskTemplateRequest decodes the body into req and returns the set of @@ -65,6 +67,7 @@ func validateTaskTemplateRequest(req taskTemplateRequest) error { {name: "name", value: req.Name, limit: db.MaxTaskTemplateNameRunes}, {name: "description", value: req.Description, limit: db.MaxTaskTemplateTextRunes}, {name: "goal", value: req.Goal, limit: db.MaxTaskTemplateTextRunes}, + {name: "remark", value: req.Remark, limit: db.MaxTaskRemarkRunes}, // 与任务备注 normalizeRemark 上限一致 } for _, check := range checks { if check.value == nil { @@ -75,9 +78,20 @@ func validateTaskTemplateRequest(req taskTemplateRequest) error { value = strings.Join(strings.Fields(value), " ") } if utf8.RuneCountInString(value) > check.limit { + if check.name == "remark" { + return fmt.Errorf("备注最多 %d 个字符", check.limit) + } return fmt.Errorf("%s 最多 %d 个字符", check.name, check.limit) } } + if req.WorkerCount != nil { + if *req.WorkerCount < 0 { + return fmt.Errorf("Worker 运行数量不能为负数") + } + if *req.WorkerCount > db.MaxTaskWorkerCount { + return fmt.Errorf("Worker 运行数量最多 %d", db.MaxTaskWorkerCount) + } + } return nil } @@ -125,12 +139,18 @@ func (s *Server) pgCreateTaskTemplate(w http.ResponseWriter, r *http.Request) { writeErr(w, http.StatusBadRequest, "拦截/允许规则无效："+err.Error()) return } + workerCount := 0 + if req.WorkerCount != nil { + workerCount = *req.WorkerCount + } template, err := pg.CreateTaskTemplate(db.TaskTemplateInput{ Name: stringValue(req.Name), Description: stringValue(req.Description), Goal: stringValue(req.Goal), + Remark: stringValue(req.Remark), CategoryID: req.CategoryID, InterceptRules: rules, + WorkerCount: workerCount, }) if err != nil { writeTaskTemplateErr(w, err) @@ -160,11 +180,21 @@ func (s *Server) pgUpdateTaskTemplate(w http.ResponseWriter, r *http.Request) { } _, catPresent := present["category_id"] _, rulesPresent := present["intercept_rules"] - if req.Name == nil && req.Description == nil && req.Goal == nil && !catPresent && !rulesPresent { - writeErr(w, http.StatusBadRequest, "至少需要提供 name、description、goal、category_id 或 intercept_rules") + _, remarkPresent := present["remark"] + _, workerPresent := present["worker_count"] + if req.Name == nil && req.Description == nil && req.Goal == nil && !catPresent && !rulesPresent && !remarkPresent && !workerPresent { + writeErr(w, http.StatusBadRequest, "至少需要提供 name、description、goal、remark、category_id、intercept_rules 或 worker_count") return } patch := db.TaskTemplatePatch{Name: req.Name, Description: req.Description, Goal: req.Goal} + if remarkPresent { + if req.Remark == nil { + empty := "" + patch.Remark = &empty + } else { + patch.Remark = req.Remark + } + } if catPresent { patch.SetCategoryID = true patch.CategoryID = req.CategoryID @@ -178,6 +208,13 @@ func (s *Server) pgUpdateTaskTemplate(w http.ResponseWriter, r *http.Request) { patch.SetInterceptRules = true patch.InterceptRules = rules } + if workerPresent { + n := 0 + if req.WorkerCount != nil { + n = *req.WorkerCount + } + patch.WorkerCount = &n + } template, err := pg.PatchTaskTemplate(id, patch) if err != nil { writeTaskTemplateErr(w, err) diff --git a/server/tools_wire_test.go b/server/tools_wire_test.go index dd4850e..705138a 100644 --- a/server/tools_wire_test.go +++ b/server/tools_wire_test.go @@ -5,9 +5,9 @@ import ( "encoding/json" "testing" - actool "github.com/Autumn-27/norma/tool" "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" + actool "github.com/Autumn-27/norma/tool" ) func names(tools []actool.CoreTool) map[string]actool.CoreTool { diff --git a/server/worker_pools_test.go b/server/worker_pools_test.go new file mode 100644 index 0000000..4facf4a --- /dev/null +++ b/server/worker_pools_test.go @@ -0,0 +1,393 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "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/norma/llm" +) + +func TestExpandWorkerSlotsLegacy(t *testing.T) { + for _, tc := range []struct { + name string + count int + want int + }{ + {"configured", 2, 2}, + {"zero defaults", 0, 3}, + {"negative defaults", -1, 3}, + } { + t.Run(tc.name, func(t *testing.T) { + slots, err := ExpandWorkerSlots(nil, tc.count) + if err != nil || len(slots) != tc.want { + t.Fatalf("slots=%v err=%v, want %d", slots, err, tc.want) + } + for i, slot := range slots { + if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != 0 { + t.Fatalf("slot[%d]=%+v", i, slot) + } + } + }) + } +} + +func TestExpandWorkerSlotsPoolOrder(t *testing.T) { + slots, err := ExpandWorkerSlots([]WorkerPool{{ProfileID: 7, Workers: 2}, {ProfileID: 12, Workers: 3}}, 99) + if err != nil || len(slots) != 5 { + t.Fatalf("slots=%v err=%v", slots, err) + } + for i, slot := range slots { + want := int64(7) + if i >= 2 { + want = 12 + } + if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != want { + t.Fatalf("slot[%d]=%+v want profile %d", i, slot, want) + } + } +} + +func TestValidateWorkerPools(t *testing.T) { + for _, tc := range []struct { + name string + pools []WorkerPool + valid bool + }{ + {"empty", nil, true}, + {"bad profile", []WorkerPool{{ProfileID: 0, Workers: 1}}, false}, + {"bad workers", []WorkerPool{{ProfileID: 1, Workers: 0}}, false}, + {"duplicate", []WorkerPool{{ProfileID: 1, Workers: 1}, {ProfileID: 1, Workers: 1}}, false}, + {"too many", []WorkerPool{{ProfileID: 1, Workers: 51}}, false}, + {"valid", []WorkerPool{{ProfileID: 1, Workers: 1}}, true}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := ValidateWorkerPools(tc.pools) == nil; got != tc.valid { + t.Fatalf("valid=%v want %v", got, tc.valid) + } + }) + } +} + +func TestWorkerPoolSettingsRoundTrip(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + pID, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pool-test", Format: "openai", Model: "pool-test", APIKey: "pool-test-key"}) + if err != nil { + t.Fatal(err) + } + defer m.pg.DeleteProfile(pID) + oldPools, _, _ := m.pg.GetSetting(settingWorkerPools) + oldWorkers, workersSet, _ := m.pg.GetSetting(settingWorkers) + defer func() { + if workersSet { + _ = m.pg.SetSetting(settingWorkers, oldWorkers) + } else { + _, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkers) + } + if oldPools != "" { + _ = m.pg.SetSetting(settingWorkerPools, oldPools) + } else { + _, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools) + } + }() + if err := m.SetWorkerPools([]WorkerPool{{ProfileID: 9223372036854775807, Workers: 1}}); err == nil { + t.Fatal("nonexistent profile accepted") + } + if err := m.SetWorkerPools([]WorkerPool{{ProfileID: pID, Workers: 2}}); err != nil { + t.Fatal(err) + } + got := m.WorkerPools() + if len(got) != 1 || got[0].ProfileID != pID || got[0].Workers != 2 || m.Workers() != 2 { + t.Fatalf("round trip pools=%+v workers=%d", got, m.Workers()) + } + if err := m.SetWorkerPools(nil); err != nil { + t.Fatal(err) + } + if got := m.WorkerPools(); got == nil || len(got) != 0 { + t.Fatalf("cleared pools=%+v", got) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + if _, err := m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools); err != nil { + t.Fatal(err) + } + res := httptest.NewRecorder() + s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil)) + var unset map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &unset); err != nil { + t.Fatal(err) + } + if pools, ok := unset["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 { + t.Fatalf("GET unset worker_pools=%#v, want empty array", unset["worker_pools"]) + } + put := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":2,"worker_pools":[{"profile_id":"`+strconv.FormatInt(pID, 10)+`","workers":2}]}`)) + putResult := httptest.NewRecorder() + s.putSettings(putResult, put) + if putResult.Code != http.StatusOK { + t.Fatalf("PUT worker pools status=%d body=%s", putResult.Code, putResult.Body.String()) + } + if m.Workers() != 2 || len(m.WorkerPools()) != 1 { + t.Fatalf("PUT worker pools did not sync settings: workers=%d pools=%+v", m.Workers(), m.WorkerPools()) + } + workersOnly := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":9}`)) + workersOnlyResult := httptest.NewRecorder() + s.putSettings(workersOnlyResult, workersOnly) + if workersOnlyResult.Code != http.StatusOK || m.Workers() != 2 { + t.Fatalf("workers-only PUT changed pool total: status=%d workers=%d", workersOnlyResult.Code, m.Workers()) + } + req := httptest.NewRequest(http.MethodGet, "/api/settings", nil) + res = httptest.NewRecorder() + s.getSettings(res, req) + if res.Code != http.StatusOK { + t.Fatalf("GET settings status=%d", res.Code) + } + var payload map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &payload); err != nil { + t.Fatal(err) + } + pools, ok := payload["worker_pools"].([]any) + if !ok || len(pools) != 1 || pools[0].(map[string]any)["profile_id"] != strconv.FormatInt(pID, 10) { + t.Fatalf("GET worker_pools=%#v, want string profile id", payload["worker_pools"]) + } + clear := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":4,"worker_pools":[]}`)) + clearResult := httptest.NewRecorder() + s.putSettings(clearResult, clear) + if clearResult.Code != http.StatusOK || m.Workers() != 4 || len(m.WorkerPools()) != 0 { + t.Fatalf("PUT legacy clear status=%d body=%s workers=%d pools=%+v", clearResult.Code, clearResult.Body.String(), m.Workers(), m.WorkerPools()) + } + res = httptest.NewRecorder() + s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil)) + var cleared map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &cleared); err != nil { + t.Fatal(err) + } + if pools, ok := cleared["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 { + t.Fatalf("GET cleared worker_pools=%#v, want empty array", cleared["worker_pools"]) + } +} + +func TestPinnedRuntimeMissingProfileDoesNotFallback(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned runtime test", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + fallback, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pinned-fallback", Format: "openai", Model: "fallback", APIKey: "fallback-key"}) + if err != nil { + t.Fatal(err) + } + defer m.pg.DeleteProfile(fallback) + if _, err := m.ReplaceTaskLLMProfiles(task.ID, []int64{fallback}, fallback); err != nil { + t.Fatal(err) + } + missing := int64(9223372036854775807) + s := &Server{m: m, llmOn: true, llmProv: &scriptedLLMProvider{}} + runtime := &taskLLMRuntime{s: s, taskID: task.ID, agentKey: "worker", pinProfileID: &missing} + if _, err := runtime.current(); err == nil { + t.Fatal("missing pinned profile silently fell back") + } + if _, ok := runtime.activeCfg(); ok || runtime.nonStreaming() || runtime.maxTokens() != 0 { + t.Fatal("missing pinned profile exposed fallback runtime settings") + } +} + +func TestPinnedRuntimeQuotaDoesNotAdvanceTaskFailover(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned quota", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + profileID, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pinned-quota", Format: "openai", Model: "pinned", APIKey: "key"}) + if err != nil { + t.Fatal(err) + } + defer m.pg.DeleteProfile(profileID) + if _, err := m.ReplaceTaskLLMProfiles(task.ID, []int64{profileID}, profileID); err != nil { + t.Fatal(err) + } + before, err := m.pg.GetTask(taskIDInt(t, task.ID)) + if err != nil || before == nil || before.ActiveLLMProfileID == nil { + t.Fatalf("load task failover state: task=%+v err=%v", before, err) + } + provider := &scriptedLLMProvider{err: errors.New("insufficient_quota")} + s := &Server{m: m, provByProfile: map[int64]*provEntry{profileID: {prov: provider, cfg: agent.Config{}}}} + id := profileID + runtime := &taskLLMRuntime{s: s, taskID: task.ID, agentKey: "worker", pinProfileID: &id} + if _, _, _, err := runtime.Complete(context.Background(), llm.CompletionRequest{}); err == nil { + t.Fatal("pinned quota error unexpectedly succeeded") + } + after, err := m.pg.GetTask(taskIDInt(t, task.ID)) + if err != nil || after == nil || after.LLMChainRevision != before.LLMChainRevision || after.ActiveLLMProfileID == nil || *after.ActiveLLMProfileID != *before.ActiveLLMProfileID { + t.Fatalf("pinned quota changed task failover state: before=%+v after=%+v err=%v", before, after, err) + } +} + +func taskIDInt(t *testing.T, id string) int64 { + t.Helper() + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + t.Fatal(err) + } + return n +} + +func TestPinnedWorkerLoopDoesNotConsultTaskRouterAndReportsFailureOnce(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned worker loop", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + + e := NewEngine(m) + var routerCalls atomic.Int32 + e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + routerCalls.Add(1) + return nil, nil + }) + e.SetWorkerResolver(func(*Task, WorkerSlot) (*agent.Worker, error) { + return nil, errors.New("profile unavailable") + }) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + e.workerLoop(ctx, task, WorkerSlot{Name: "work#4", ProfileID: 12}) + close(done) + }() + time.Sleep(100 * time.Millisecond) + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("worker loop did not stop") + } + if routerCalls.Load() != 0 { + t.Fatalf("pinned worker consulted task router %d times", routerCalls.Load()) + } + activities, _, err := task.Store.ActivityList(nil, 0, 20) + if err != nil { + t.Fatal(err) + } + var failures int + for _, activity := range activities { + if activity.Worker == "system" && activity.IsError { + failures++ + } + } + if failures != 1 { + t.Fatalf("pinned failure activities=%d, want one throttled entry: %+v", failures, activities) + } +} + +func TestPinnedWorkerAvailabilityRecoveryUsesSystemSession(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned availability", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + e := NewEngine(m) + slot := WorkerSlot{Name: "work#4", ProfileID: 12} + e.reportSlotAvailability(task, slot, errors.New("profile unavailable")) + e.reportSlotAvailability(task, slot, nil) + activities, _, err := task.Store.ActivityList(nil, 0, 20) + if err != nil { + t.Fatal(err) + } + var failure, recovery int + for _, activity := range activities { + if activity.Worker != "system" { + continue + } + if activity.IsError && strings.Contains(activity.Summary, "work#4") && strings.Contains(activity.Summary, "#12") { + failure++ + } + if !activity.IsError && strings.Contains(activity.Summary, "work#4") && strings.Contains(activity.Summary, "#12") { + recovery++ + } + } + if failure != 1 || recovery != 1 { + t.Fatalf("system availability activities failure=%d recovery=%d: %+v", failure, recovery, activities) + } +} + +func TestDetachedIntentPinnedOwnerUsesPinnedResolver(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("detached pinned owner", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + intentID, err := task.Store.AddIntent(map[string]any{"text": "continue"}, 0, nil, "test") + if err != nil { + t.Fatal(err) + } + if ok, err := task.Store.ClaimIntent(intentID, "work#4"); err != nil || !ok { + t.Fatalf("claim intent: ok=%v err=%v", ok, err) + } + if changed, err := task.Store.CompareAndSetIntentState(intentID, "running", "paused"); err != nil || !changed { + t.Fatalf("pause intent: changed=%v err=%v", changed, err) + } + e := NewEngine(m) + e.slotPlans.Store(task.ID, map[string]int64{"work#4": 12}) + var pinnedCalls, snapshotCalls atomic.Int32 + e.SetWorkerResolver(func(_ *Task, slot WorkerSlot) (*agent.Worker, error) { + pinnedCalls.Add(1) + if slot.Name != "work#4" || slot.ProfileID != 12 { + t.Fatalf("unexpected pinned slot: %+v", slot) + } + return nil, errors.New("profile unavailable") + }) + e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + snapshotCalls.Add(1) + return nil, nil + }) + err = e.runDetachedIntent(context.Background(), task, intentID, "request", "human", "agent") + if err == nil || !strings.Contains(err.Error(), "LLM 配置 #12 不可用") { + t.Fatalf("runDetachedIntent err=%v", err) + } + if pinnedCalls.Load() != 1 || snapshotCalls.Load() != 0 { + t.Fatalf("pinned calls=%d snapshot calls=%d", pinnedCalls.Load(), snapshotCalls.Load()) + } + node, err := task.Store.GetNode(intentID) + if err != nil || node == nil || node.State != "paused" { + t.Fatalf("intent changed after unavailable resolver: node=%+v err=%v", node, err) + } +} diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index f771776..d4a640a 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -58,6 +58,7 @@ function FindingDetailInner() { const searchParams = useSearchParams(); const id = searchParams.get("id") ?? ""; const contextTaskId = searchParams.get("context_task") ?? ""; + const backHref = contextTaskId ? `/function/tasks/detail?id=${contextTaskId}&tab=findings` : "/function/findings"; const [finding, setFinding] = React.useState(null); const [loaded, setLoaded] = React.useState(false); const [tab, setTab] = React.useState("overview"); @@ -117,7 +118,7 @@ function FindingDetailInner() {

{loaded ? `未找到发现 ${id}` : "加载中…"}

{loaded && ( @@ -135,7 +136,7 @@ function FindingDetailInner() {
diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx index 679e23f..577944f 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx @@ -115,11 +115,7 @@ function Row({ {f.finding_id && ( @@ -205,6 +201,7 @@ export function FindingsTab({ taskId }: { taskId: string }) { {[ ...new Set( diff --git a/web/src/app/(main)/function/tasks/detail/page.tsx b/web/src/app/(main)/function/tasks/detail/page.tsx index 3c23d3e..db99f6c 100644 --- a/web/src/app/(main)/function/tasks/detail/page.tsx +++ b/web/src/app/(main)/function/tasks/detail/page.tsx @@ -354,7 +354,17 @@ function TaskDetailInner() { ); return ( - + { + setTab(value); + const params = new URLSearchParams(searchParams.toString()); + params.set("tab", value); + params.delete("session"); + router.replace(`/function/tasks/detail?${params.toString()}`, { scroll: false }); + }} + className="flex flex-1 flex-col gap-0" + > {/* Top fixed area */}
@@ -404,6 +414,11 @@ function TaskDetailInner() {

{task.goal}

+ {task.remark ? ( +

+ 备注：{task.remark} +

+ ) : null} {/* Tabs */}
diff --git a/web/src/app/(main)/function/tasks/page.tsx b/web/src/app/(main)/function/tasks/page.tsx index bc9ae91..3737d6d 100644 --- a/web/src/app/(main)/function/tasks/page.tsx +++ b/web/src/app/(main)/function/tasks/page.tsx @@ -137,6 +137,7 @@ import type { TaskCategory, TaskStatus, } from "@/lib/types"; +import { MAX_TASK_REMARK_RUNES, MAX_TASK_WORKER_COUNT, parseWorkerCountInput } from "@/lib/types"; import { cn } from "@/lib/utils"; // fmtBytes renders a human file size for the upload manifest. @@ -3104,6 +3105,7 @@ function CreateTaskSheet({ }) { const [open, setOpen] = React.useState(false); const [name, setName] = React.useState(""); + const [remark, setRemark] = React.useState(""); const [categoryID, setCategoryID] = React.useState(undefined); const [description, setDescription] = React.useState(""); const [goal, setGoal] = React.useState(""); @@ -3118,6 +3120,7 @@ function CreateTaskSheet({ const [heartbeatMin, setHeartbeatMin] = React.useState("10"); // planner 心跳(分钟);默认10,下限10(与后端一致) const [seedFirstIntent, setSeedFirstIntent] = React.useState(false); // 创建时下发种子意图,worker 免等首轮 planner 直接开跑;默认关闭,走标准先规划再执行 const [coverageEnabled, setCoverageEnabled] = React.useState(true); // 资产覆盖度功能;默认开。关闭=不计算/展示覆盖度、不累积范围、隐藏范围类工具(company 关联不受影响) + const [workerCount, setWorkerCount] = React.useState(""); // 本任务 Worker 数;空/0=用全局默认 const [interceptRules, setInterceptRules] = React.useState([]); // 任务级资产拦截规则(仅本任务生效,不进全局表) const [scheduleEnabled, setScheduleEnabled] = React.useState(false); const [schedule, setSchedule] = React.useState({ @@ -3181,12 +3184,18 @@ function CreateTaskSheet({ toast.error(`最多关联 ${MAX_SOURCE_TASKS} 个来源任务`); return; } + const workerCountNum = parseWorkerCountInput(workerCount); + if (workerCountNum == null) { + toast.error(`Worker 运行数量必须是 0-${MAX_TASK_WORKER_COUNT} 的整数`); + return; + } setCreating(true); try { const timeoutSec = Math.max(0, Math.floor(Number(timeoutMin) || 0)) * 60; const heartbeatSec = Math.max(10, Math.floor(Number(heartbeatMin) || 10)) * 60; // 下限 10min，与后端归一一致 await api.createTask({ name: name.trim(), + remark: remark.trim(), categoryId: categoryID, description: description.trim(), goal: goal.trim(), @@ -3197,6 +3206,7 @@ function CreateTaskSheet({ seedFirstIntent, planHeartbeatSeconds: heartbeatSec, coverageEnabled, + workerCount: workerCountNum, schedule: scheduleEnabled ? { ...schedule, task_ids: [] } : undefined, interceptRules: interceptRules .map((r) => ({ ...r, pattern: r.pattern.trim() })) @@ -3204,6 +3214,7 @@ function CreateTaskSheet({ }); toast.success("任务已创建"); setName(""); + setRemark(""); setCategoryID(undefined); setDescription(""); setGoal(""); @@ -3215,6 +3226,7 @@ function CreateTaskSheet({ setHeartbeatMin("10"); setSeedFirstIntent(false); setCoverageEnabled(true); + setWorkerCount(""); setInterceptRules([]); setScheduleEnabled(false); setSchedule({ @@ -3263,15 +3275,19 @@ function CreateTaskSheet({ { setDescription(template.description); setGoal(template.goal); + setRemark(template.remark ?? ""); setCategoryID(template.category_id ?? undefined); setInterceptRules(template.intercept_rules ?? []); + setWorkerCount(template.worker_count ? String(template.worker_count) : ""); setUploadCount(0); }} portalContainer={sheetContentRef} @@ -3285,6 +3301,17 @@ function CreateTaskSheet({ onChange={(e) => setName(e.target.value)} />
+
+ +