From 5244500feec5b31500e389653eeee13b59320a92 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:17:04 -0400 Subject: [PATCH 01/70] =?UTF-8?q?fix(scheduler):=20=E4=BA=8B=E4=BB=B6?= =?UTF-8?q?=E6=B0=B4=E4=BD=8D=E6=AF=8F=20tick=20=E9=83=BD=E6=8E=A8?= =?UTF-8?q?=E8=BF=9B,=E9=81=BF=E5=85=8D=E6=96=B0=E5=BB=BA=E8=A7=A6?= =?UTF-8?q?=E5=8F=91=E5=99=A8=E5=9B=9E=E6=94=BE=E5=8E=86=E5=8F=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 原来 fire{Findings,Goals,TaskTimeouts,TaskCreates,ToolCalls} 在没有 对应触发器时直接 return,不推进水位。于是水位会一直停在旧值;等用户 新建/启用触发器(如预置的 reporter on_tool_call),下一个 tick 就把 水位之后的全部历史事件当成新事件触发一遍(逐条唤起 reporter 写报告)。 改为:即便当前没有对应触发器,也照常把水位推进到最新(静默消费), 触发动作只在存在触发器时执行。这样触发器只对启用之后出现的事件生效, 不再回放历史。goal 用 fired-set 同理:无触发器时也标记为已消费。 --- server/scheduler.go | 38 +++++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/server/scheduler.go b/server/scheduler.go index 1aeeb17..01fb6d1 100644 --- a/server/scheduler.go +++ b/server/scheduler.go @@ -155,19 +155,22 @@ func (sc *Scheduler) fireFindings(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64) events, err := sc.pg.NewFindingsSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark whether or not any on_finding trigger is active: a + // finding fires only for triggers live at the moment it appears. Otherwise a + // trigger enabled later would replay the entire historical backlog at once. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务发现 finding 触发】\n任务: #%d %s(目标:%s)\n发现: [%s/%s] %s", e.TaskID, e.TaskDesc, e.TaskGoal, e.VulnClass, e.Severity, e.Summary) for _, tr := range want { @@ -185,13 +188,12 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } events, err := sc.pg.MetGoals() if err != nil || len(events) == 0 { return } + // Mark goals as consumed whether or not a trigger is active, so enabling an + // on_goal_met trigger later doesn't replay every already-met goal. fired := sc.firedGoalSet() changed := false for _, e := range events { @@ -200,6 +202,9 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) { } fired[e.NodeID] = true changed = true + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s(目标:%s)\n达成目标: %s", e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary) for _, tr := range want { @@ -220,19 +225,20 @@ func (sc *Scheduler) fireTaskTimeouts(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64) events, err := sc.pg.TimedOutTasksSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s(目标:%s)", e.TaskID, e.TaskDesc, e.TaskGoal) for _, tr := range want { @@ -251,19 +257,20 @@ func (sc *Scheduler) fireTaskCreates(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTaskCreate), 10, 64) events, err := sc.pg.NewTasksSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务创建触发】\n任务ID: #%d\n任务描述: %s\n任务目标: %s", e.TaskID, e.TaskDesc, e.TaskGoal) for _, tr := range want { @@ -283,19 +290,20 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastToolCall), 10, 64) events, err := sc.pg.NewToolCallsSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } errTag := "" if e.ToolIsErr { errTag = "[error] " From 8e9b126972a7d27fa64b49d14a7f4701ce1f53fd Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:17:13 -0400 Subject: [PATCH 02/70] =?UTF-8?q?feat(assets):=20=E4=BB=BB=E5=8A=A1/?= =?UTF-8?q?=E4=BC=81=E4=B8=9A=E8=B5=84=E4=BA=A7=E5=88=97=E8=A1=A8=E6=94=B9?= =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E7=AB=AF=E5=88=86=E9=A1=B5,=E5=8E=BB?= =?UTF-8?q?=E6=8E=89=2050=20=E6=9D=A1=E4=B8=8A=E9=99=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - db: QueryByTask/QueryByCompany 加 offset 参数,limit<=0 不再兜底成 50/200 而是不限量;新增 CountByTask/CountByCompany/CountsByTypeForTask; 抽出共享 pageClause 与 scanTypeCounts。 - server: /assets 的 task_id/company_id 分支改用真实 total(不再是 len(assets));/assets/counts 支持 ?task_id= 返回按类型计数。 - web: taskAssets(taskId,type,limit,offset) 返回 {assets,total}; assetCounts(taskId?);assets-tab 改为只取当前分类当前页,分类徽章走 counts 接口,分页器 100/200 真正可用,轮询变轻。 - 补 TestQueryByTaskPaging/TestQueryByCompanyPaging 回归(60 条 25/页翻页 不重不漏)。 --- db/assets.go | 76 ++++++++-- db/assets_test.go | 143 +++++++++++++++++- server/assets.go | 27 +++- .../tasks/detail/_tabs/assets-tab.tsx | 72 +++++---- web/src/lib/api.ts | 12 +- 5 files changed, 268 insertions(+), 62 deletions(-) diff --git a/db/assets.go b/db/assets.go index fee5d94..3cf1c1b 100644 --- a/db/assets.go +++ b/db/assets.go @@ -913,10 +913,8 @@ func (s *AssetStore) CountByType(typ string) (int, error) { } // QueryByCompany returns assets for a company, optionally filtered by type. -func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit int) ([]*Asset, error) { - if limit <= 0 { - limit = 200 - } +// limit <= 0 means no limit — pair it with CountByCompany for server-side paging. +func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), COALESCE(c_segment::text,''), port, @@ -933,8 +931,7 @@ FROM assets WHERE company_id = $1` args = append(args, typ) q += fmt.Sprintf(` AND type = $%d`, len(args)) } - args = append(args, limit) - q += fmt.Sprintf(` ORDER BY last_seen DESC LIMIT $%d`, len(args)) + q += pageClause(&args, limit, offset) rows, err := s.db.Query(q, args...) if err != nil { return nil, err @@ -943,11 +940,38 @@ FROM assets WHERE company_id = $1` return scanAssets(rows) } -// QueryByTask returns assets rows that have a given task_id in task_ids. -func (s *AssetStore) QueryByTask(taskID int64, typ string, limit int) ([]*Asset, error) { - if limit <= 0 { - limit = 200 +// CountByCompany returns the full number of assets attributed to a company +// (optionally filtered by type), ignoring limit/offset. +func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) { + q := `SELECT count(*) FROM assets WHERE company_id = $1` + args := []any{companyID} + if typ != "" { + args = append(args, typ) + q += fmt.Sprintf(` AND type = $%d`, len(args)) } + var n int + err := s.db.QueryRow(q, args...).Scan(&n) + return n, err +} + +// pageClause appends limit/offset placeholders to args and returns the trailing +// ORDER BY … LIMIT … OFFSET fragment. limit <= 0 means no limit. +func pageClause(args *[]any, limit, offset int) string { + q := ` ORDER BY last_seen DESC` + if limit > 0 { + *args = append(*args, limit) + q += fmt.Sprintf(` LIMIT $%d`, len(*args)) + } + if offset > 0 { + *args = append(*args, offset) + q += fmt.Sprintf(` OFFSET $%d`, len(*args)) + } + return q +} + +// QueryByTask returns assets rows that have a given task_id in task_ids. +// limit <= 0 means no limit — pair it with CountByTask for server-side paging. +func (s *AssetStore) QueryByTask(taskID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), COALESCE(c_segment::text,''), port, @@ -964,8 +988,7 @@ FROM assets WHERE $1 = ANY(task_ids)` args = append(args, typ) q += fmt.Sprintf(` AND type = $%d`, len(args)) } - args = append(args, limit) - q += fmt.Sprintf(` ORDER BY last_seen DESC LIMIT $%d`, len(args)) + q += pageClause(&args, limit, offset) rows, err := s.db.Query(q, args...) if err != nil { return nil, err @@ -974,6 +997,30 @@ FROM assets WHERE $1 = ANY(task_ids)` return scanAssets(rows) } +// CountByTask returns the full number of assets attached to a task (optionally +// filtered by type), ignoring limit/offset — lets the UI paginate on the server. +func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { + q := `SELECT count(*) FROM assets WHERE $1 = ANY(task_ids)` + args := []any{taskID} + if typ != "" { + args = append(args, typ) + q += fmt.Sprintf(` AND type = $%d`, len(args)) + } + var n int + err := s.db.QueryRow(q, args...).Scan(&n) + return n, err +} + +// CountsByTypeForTask returns asset counts per type for a single task. +func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) { + rows, err := s.db.Query(`SELECT type, COUNT(*) FROM assets WHERE $1 = ANY(task_ids) GROUP BY type`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTypeCounts(rows) +} + const assetSelectCols = `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), COALESCE(c_segment::text,''), port, @@ -1072,6 +1119,11 @@ func (s *AssetStore) CountsByType() (map[string]int, error) { return nil, err } defer rows.Close() + return scanTypeCounts(rows) +} + +// scanTypeCounts scans a (type, count) result set into a map. +func scanTypeCounts(rows *sql.Rows) (map[string]int, error) { out := map[string]int{} for rows.Next() { var typ string diff --git a/db/assets_test.go b/db/assets_test.go index 4c5b5b2..2245192 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -1,6 +1,7 @@ package db import ( + "fmt" "testing" ) @@ -545,7 +546,7 @@ func TestQueryByTask(t *testing.T) { } defer deleteAsset(d, id) - assets, err := av2.QueryByTask(taskID, "", 10) + assets, err := av2.QueryByTask(taskID, "", 10, 0) if err != nil { t.Fatal(err) } @@ -558,6 +559,73 @@ func TestQueryByTask(t *testing.T) { if !found { t.Error("QueryByTask: inserted asset not found") } + + n, err := av2.CountByTask(taskID, "") + if err != nil { + t.Fatal(err) + } + if n < 1 { + t.Errorf("CountByTask: got %d, want >= 1", n) + } + if counts, err := av2.CountsByTypeForTask(taskID); err != nil { + t.Fatal(err) + } else if counts["root_domain"] < 1 { + t.Errorf("CountsByTypeForTask: root_domain = %d, want >= 1", counts["root_domain"]) + } + + // offset past the end must return nothing, not fall back to page 1 + rest, err := av2.QueryByTask(taskID, "", 10, n) + if err != nil { + t.Fatal(err) + } + if len(rest) != 0 { + t.Errorf("QueryByTask offset=%d: got %d rows, want 0", n, len(rest)) + } +} + +// 任务资产列表按页取,不再被固定条数截断:60 条资产用 25/页要能完整翻出来。 +func TestQueryByTaskPaging(t *testing.T) { + d, av2, _ := testSetup(t) + defer d.Close() + + const taskID = int64(99998) + const n = 60 + for i := 0; i < n; i++ { + id, err := av2.UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("paging-%02d.querybytask.test", i), + TaskID: taskID, + }) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, id) + } + + total, err := av2.CountByTask(taskID, "root_domain") + if err != nil { + t.Fatal(err) + } + if total != n { + t.Fatalf("CountByTask: got %d, want %d", total, n) + } + + const size = 25 + seen := map[int64]bool{} + for offset := 0; offset < total; offset += size { + page, err := av2.QueryByTask(taskID, "root_domain", size, offset) + if err != nil { + t.Fatal(err) + } + for _, a := range page { + if seen[a.ID] { + t.Errorf("offset=%d: asset %d returned twice", offset, a.ID) + } + seen[a.ID] = true + } + } + if len(seen) != n { + t.Errorf("paged through %d assets, want %d", len(seen), n) + } } func TestQueryByCompany(t *testing.T) { @@ -582,7 +650,7 @@ func TestQueryByCompany(t *testing.T) { t.Fatal(err) } - assets, err := av2.QueryByCompany(companyID, "root_domain", 10) + assets, err := av2.QueryByCompany(companyID, "root_domain", 10, 0) if err != nil { t.Fatal(err) } @@ -595,6 +663,77 @@ func TestQueryByCompany(t *testing.T) { if !found { t.Error("QueryByCompany: attributed asset not found") } + + n, err := av2.CountByCompany(companyID, "root_domain") + if err != nil { + t.Fatal(err) + } + if n < 1 { + t.Errorf("CountByCompany: got %d, want >= 1", n) + } + + // offset past the end must return nothing, not fall back to page 1 + rest, err := av2.QueryByCompany(companyID, "root_domain", 10, n) + if err != nil { + t.Fatal(err) + } + if len(rest) != 0 { + t.Errorf("QueryByCompany offset=%d: got %d rows, want 0", n, len(rest)) + } +} + +// 企业资产列表同样按页取,不被固定条数截断。 +func TestQueryByCompanyPaging(t *testing.T) { + d, av2, cs := testSetup(t) + defer d.Close() + + companyID, _, err := cs.UpsertCompany("QueryByCompanyPagingCorp", "") + if err != nil { + t.Fatal(err) + } + defer cleanupCompany(d, companyID) + + if added, _, _, errs := cs.AddScope(companyID, []string{"qbc-paging.io"}, "test"); added != 1 { + t.Fatalf("AddScope: added=%d, errors=%v", added, errs) + } + + // UpsertSubdomain 会顺带建根域名资产,一并清掉 + defer d.Exec(`DELETE FROM assets WHERE root_domain = 'qbc-paging.io'`) + + const n = 60 + for i := 0; i < n; i++ { + if _, err := av2.UpsertSubdomain(UpsertSubdomainReq{ + Domain: fmt.Sprintf("paging-%02d.qbc-paging.io", i), + }); err != nil { + t.Fatal(err) + } + } + + total, err := av2.CountByCompany(companyID, "subdomain") + if err != nil { + t.Fatal(err) + } + if total != n { + t.Fatalf("CountByCompany: got %d, want %d", total, n) + } + + const size = 25 + seen := map[int64]bool{} + for offset := 0; offset < total; offset += size { + page, err := av2.QueryByCompany(companyID, "subdomain", size, offset) + if err != nil { + t.Fatal(err) + } + for _, a := range page { + if seen[a.ID] { + t.Errorf("offset=%d: asset %d returned twice", offset, a.ID) + } + seen[a.ID] = true + } + } + if len(seen) != n { + t.Errorf("paged through %d assets, want %d", len(seen), n) + } } func TestCountsByType(t *testing.T) { diff --git a/server/assets.go b/server/assets.go index 6712a8a..62a794e 100644 --- a/server/assets.go +++ b/server/assets.go @@ -217,6 +217,9 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) { // GET /api/assets // ===================================================================== +// defaultAssetPageSize is the page size used when the caller omits ?limit. +const defaultAssetPageSize = 50 + func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { as := s.assetStore() if as == nil { @@ -227,8 +230,10 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { typ := q.Get("type") limit, _ := strconv.Atoi(q.Get("limit")) offset, _ := strconv.Atoi(q.Get("offset")) + // limit/offset are page controls, not a cap: `total` always reports the full + // match count so callers can page through everything. if limit <= 0 { - limit = 50 + limit = defaultAssetPageSize } var assets []*db.Asset @@ -247,11 +252,15 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { taskID, _ := strconv.ParseInt(q.Get("task_id"), 10, 64) switch { case companyID > 0: - assets, err = as.QueryByCompany(companyID, typ, limit) - total = len(assets) + assets, err = as.QueryByCompany(companyID, typ, limit, offset) + if err == nil { + total, err = as.CountByCompany(companyID, typ) + } case taskID > 0: - assets, err = as.QueryByTask(taskID, typ, limit) - total = len(assets) + assets, err = as.QueryByTask(taskID, typ, limit, offset) + if err == nil { + total, err = as.CountByTask(taskID, typ) + } default: if typ == "" { typ = "root_domain" @@ -284,7 +293,13 @@ func (s *Server) assetCounts(w http.ResponseWriter, r *http.Request) { writeErr(w, 503, "database unavailable") return } - counts, err := as.CountsByType() + var counts map[string]int + var err error + if taskID, _ := strconv.ParseInt(r.URL.Query().Get("task_id"), 10, 64); taskID > 0 { + counts, err = as.CountsByTypeForTask(taskID) + } else { + counts, err = as.CountsByType() + } if err != nil { writeErr(w, 500, err.Error()) return diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx index ffb1d51..7d066c6 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx @@ -152,50 +152,51 @@ const TABS: { key: string; label: string; icon: LucideIcon }[] = [ ]; export function AssetsTab({ taskId }: { taskId: string }) { - const [assets, setAssets] = React.useState([]); + // 服务端分页:只取当前分类的当前页,分类总数走 /assets/counts?task_id=,不做数量截断 + const [rows, setRows] = React.useState([]); + const [total, setTotal] = React.useState(0); + const [counts, setCounts] = React.useState>({}); const [loaded, setLoaded] = React.useState(false); const [tab, setTab] = React.useState("root_domain"); const [page, setPage] = React.useState(0); const [size, setSize] = React.useState(50); - React.useEffect(() => setPage(0), [tab, size]); + // 换分类/换页大小:回到第一页,并清掉旧行(各分类列不同,不能串用) + React.useEffect(() => { + setPage(0); + setRows([]); + }, [tab, size]); React.useEffect(() => { let alive = true; const load = async () => { try { - const results = await Promise.all( - TABS.map((t) => api.taskAssets(taskId, t.key)), - ); + const [cur, c] = await Promise.all([ + api.taskAssets(taskId, tab, size, page * size), + api.assetCounts(taskId), + ]); if (!alive) return; - setAssets(results.flat()); - setLoaded(true); + setRows(cur.assets); + setTotal(cur.total); + setCounts(c ?? {}); } catch { + // 保留上一次的数据 + } finally { if (alive) setLoaded(true); } }; load(); const timer = setInterval(load, 10_000); return () => { alive = false; clearInterval(timer); }; - }, [taskId]); + }, [taskId, tab, page, size]); - const byType = React.useMemo(() => { - const m: Record = { root_domain: [], ip: [], subdomain: [], service: [], endpoint: [] }; - for (const a of assets) { - if (a.type in m) m[a.type].push(a); - } - return m; - }, [assets]); + // 轮询期间资产变少时,别停在越界的空页上 + React.useEffect(() => { + const maxPage = Math.max(0, Math.ceil(total / size) - 1); + if (page > maxPage) setPage(maxPage); + }, [total, size, page]); - const counts: Record = { - root_domain: byType.root_domain.length, - ip: byType.ip.length, - subdomain: byType.subdomain.length, - service: byType.service.length, - endpoint: byType.endpoint.length, - }; - - const totalAll = assets.length; + const totalAll = TABS.reduce((n, t) => n + (counts[t.key] ?? 0), 0); if (!loaded) { return
加载中…
; @@ -209,9 +210,6 @@ export function AssetsTab({ taskId }: { taskId: string }) { ); } - const cur = byType[tab] ?? []; - const slice = cur.slice(page * size, page * size + size); - return ( {/* left nav */} @@ -238,8 +236,8 @@ export function AssetsTab({ taskId }: { taskId: string }) {
- - {slice.map((a) => ( + + {rows.map((a) => ( {a.domain} {a.icp || "—"} @@ -249,8 +247,8 @@ export function AssetsTab({ taskId }: { taskId: string }) { - - {slice.map((a) => ( + + {rows.map((a) => ( {a.ip} {a.c_segment || "—"} @@ -269,8 +267,8 @@ export function AssetsTab({ taskId }: { taskId: string }) { - - {slice.map((a) => ( + + {rows.map((a) => ( {a.domain} {a.root_domain || "—"} @@ -284,13 +282,13 @@ export function AssetsTab({ taskId }: { taskId: string }) { - {slice.map((a) => { + {rows.map((a) => { const isHttp = a.service_type === "http"; const addr = isHttp ? (a.url || "") @@ -330,8 +328,8 @@ export function AssetsTab({ taskId }: { taskId: string }) { - - {slice.map((a) => ( + + {rows.map((a) => ( {a.url || "—"} diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index e14ae4d..d2eb7ec 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -202,13 +202,15 @@ export const api = { get<{ count: number; total: number; assets: Asset[] }>( `/assets?dsl=${encodeURIComponent(dsl)}${type ? `&type=${encodeURIComponent(type)}` : ""}&limit=${limit}&offset=${offset}`, ).then((r) => ({ assets: r?.assets ?? [], total: r?.total ?? r?.count ?? 0 })), - assetCounts: () => get>("/assets/counts"), + assetCounts: (taskId = "") => + get>(`/assets/counts${taskId ? `?task_id=${encodeURIComponent(taskId)}` : ""}`), deleteAssets: (ids: number[]) => http<{ deleted: number }>("/assets", { method: "DELETE", body: JSON.stringify({ ids }) }), - // legacy — kept for task-specific views; hits the same endpoint with task_id filter - taskAssets: (taskId: string, type = "") => - get<{ count: number; assets: Asset[] }>(`/assets?task_id=${taskId}&type=${type}`) - .then((r) => r?.assets ?? []), + // task-scoped view of the same endpoint — server-side paginated like `assets` + taskAssets: (taskId: string, type = "", limit = 50, offset = 0) => + get<{ count: number; total: number; assets: Asset[] }>( + `/assets?task_id=${encodeURIComponent(taskId)}&type=${encodeURIComponent(type)}&limit=${limit}&offset=${offset}`, + ).then((r) => ({ assets: r?.assets ?? [], total: r?.total ?? r?.count ?? 0 })), // ---- companies (企业 + 资产范围;归属唯一来源) ---- companies: () => get("/companies").then(arr), From a6bd18050cdfc8652b5fd5530fdf2711ceb694cb Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:20:10 -0400 Subject: [PATCH 03/70] =?UTF-8?q?feat(agent):=20=E8=AF=8A=E6=96=AD=20run?= =?UTF-8?q?=20=E5=8F=96=E6=B6=88/=E7=BB=88=E6=AD=A2=E5=8E=9F=E5=9B=A0,?= =?UTF-8?q?=E5=8F=96=E4=BB=A3=E7=AC=BC=E7=BB=9F=E7=9A=84=E4=B8=AD=E6=96=AD?= =?UTF-8?q?=E6=8F=90=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 AbortCause(cancelcause.go):各取消点用 WithCancelCause/WithTimeoutCause 附上原因(暂停/删除/kill_work/目标达成/drain 兜底/聊天停止/进程关闭/硬超时…), captureRun 用 context.Cause 还原出到底是谁、为何取消。 - 新增 terminalreason.go:把终端状态转成可读的一行摘要 + 多行详情(取消原因、 终止 reason、轮数、卡住的工具),替换原来固定那句猜测式提示。 - engine/worker/planner/capture 及 conversations/orchestration/server 接入; cmd/artex/main.go、前端 llm/main-content 相应调整。 - 配套测试 terminalreason_test.go / engine_cancelcause_test.go。 --- agent/cancelcause.go | 116 ++++++++++++ agent/capture.go | 41 ++--- agent/planner.go | 2 +- agent/terminalreason.go | 160 ++++++++++++++++ agent/terminalreason_test.go | 172 ++++++++++++++++++ agent/worker.go | 2 +- cmd/artex/main.go | 11 +- server/conversations.go | 7 +- server/engine.go | 40 ++-- server/engine_cancelcause_test.go | 96 ++++++++++ server/engine_timeout.go | 2 +- server/orchestration.go | 2 +- server/server.go | 14 +- server/server_mgmt.go | 2 +- .../app/(main)/_components/main-content.tsx | 2 +- web/src/app/(main)/system/llm/page.tsx | 15 +- 16 files changed, 624 insertions(+), 60 deletions(-) create mode 100644 agent/cancelcause.go create mode 100644 agent/terminalreason.go create mode 100644 agent/terminalreason_test.go create mode 100644 server/engine_cancelcause_test.go diff --git a/agent/cancelcause.go b/agent/cancelcause.go new file mode 100644 index 0000000..cfda073 --- /dev/null +++ b/agent/cancelcause.go @@ -0,0 +1,116 @@ +package agent + +import ( + "context" + "errors" + "fmt" +) + +// AbortCause names WHY a planner/worker run's context was cancelled. +// +// Background: when a run's ctx is cancelled mid-flight the harness finishes with +// ReasonAbortedStreaming / ReasonAbortedTools and Terminal.Err = context.Canceled. +// ARTEX cancels that ctx from a dozen different places (pause, delete, kill_work, +// goal-met, settle drain backstop, chat stop, process shutdown, hard run timeout, +// …) but by the time it reaches captureRun they all collapse into one bare +// context.Canceled — so the trace could only print a guess-list, not a diagnosis. +// +// Fix: every cancel site uses context.WithCancelCause / WithTimeoutCause and +// attaches one of the causes below; captureRun recovers it via context.Cause and +// prints exactly who cancelled the run, why, and what happens to the intent. +// +// When you add a new cancel site, add its cause HERE — a site that cancels +// without one shows up in the trace as "canceled_no_cause", by design. +type AbortCause struct { + Code string // stable machine-readable id (logs / future filtering) + Short string // 一句话,进折叠摘要行 + Text string // 完整说明:谁取消的、为什么、这条意图接下来会怎样 +} + +func (c *AbortCause) Error() string { return c.Text } + +func cause(code, short, text string) *AbortCause { + return &AbortCause{Code: code, Short: short, Text: text} +} + +// Causef builds an ad-hoc cause with runtime detail (e.g. which intent, which +// deadline). Use the predeclared vars below when there is nothing dynamic to add. +func Causef(code, short, format string, args ...any) *AbortCause { + return &AbortCause{Code: code, Short: short, Text: fmt.Sprintf(format, args...)} +} + +// The full enumeration of cancel sites, grouped by which context gets cancelled. +var ( + // ---- 任务级 exec ctx(Engine.cancelExec:取消该任务在跑的 planner + 全部 worker)---- + + AbortPausedByUser = cause("paused_by_user", "用户暂停了任务", + "用户暂停了任务(前端「暂停」按钮 / POST /api/tasks/{id}/state action=pause)。本次运行被主动取消,该意图已退回 frontier(open),恢复任务后会被重新领取、从头重跑") + + AbortPausedByOrchestrator = cause("paused_by_orchestrator", "编排 agent 调用 pause_task 暂停了任务", + "编排 agent 调用了 pause_task 工具暂停本任务。本次运行被主动取消,该意图已退回 frontier(open),恢复后重跑") + + AbortTaskDeleted = cause("task_deleted", "任务被删除", + "任务被删除(DELETE /api/tasks/{id})——删除前先取消在跑的 planner/worker。本次运行的结果不会再被使用") + + AbortPausedOnReload = cause("paused_on_reload", "后端启动时恢复了「已暂停」状态", + "后端启动时按库里持久化的「已暂停」状态恢复暂停。本次运行被取消(正常情况下此时并没有在跑的 run)") + + AbortGoalMet = cause("goal_met", "规划者判定目标已达成,任务落终态 done", + "规划者判定任务目标已达成、任务落终态 done,随即取消所有在跑的 worker——它们手头的意图跑完也没意义了。该意图标记为 stopped,不是失败") + + AbortSettleDrainTimeout = cause("settle_drain_timeout", "任务级超时收尾:drain 宽限用尽后强制取消", + "任务级超时收尾的硬兜底:到达任务 timeout 后等在跑 worker 优雅收尾,超过 90s drain 宽限仍未结束 → 强制取消。该意图标记为 exhausted,收尾阶段通常已把事实/资产写回") + + // ---- 单个 work 的 ctx(Engine.workCancel[intentID])---- + + AbortKilledByPlanner = cause("killed_by_planner", "规划者用 kill_work 终止了这条意图", + "规划者调用 kill_work 主动终止了这条意图(判定方向跑偏/无意义,止损)。该意图标记为 stopped,不会被自动重领") + + AbortWorkFinished = cause("work_finished", "work 已正常结束,引擎释放其 context", + "work 已正常结束,引擎在 unregisterWork 里释放其 context 资源。这不是中断——若它出现在中断消息里,说明取消与收场事件存在竞态") + + AbortPausedRaceGuard = cause("paused_race_guard", "任务处于暂停态,引擎拒绝启动新 run", + "任务处于暂停态时请求执行上下文,引擎直接返回一个已取消的 context(守卫 claim→Execute 的竞态,避免暂停后仍启动新 run)。该意图会退回 frontier") + + // ---- 对话 / 主 Agent 的单轮 ctx(Server.chatCancel[busyKey])---- + + AbortChatStoppedByUser = cause("chat_stopped_by_user", "用户点了「停止」中止本轮对话", + "用户点了「停止」,主动中止本轮对话运行(POST .../chat/stop 或 .../conversations/{id}/stop)。已产生的步骤都保留,可以直接发下一条消息") + + AbortChatTurnFinished = cause("chat_turn_finished", "本轮对话已正常结束,服务端释放其 context", + "本轮对话已正常结束,服务端在 defer 里释放其 context 资源。这不是中断——若它出现在中断消息里,说明取消与收场事件存在竞态") + + // ---- 进程级 ---- + + AbortShutdown = cause("shutdown", "后端进程正在关停(重启 / 更新 / Ctrl-C)", + "后端进程收到 SIGINT/SIGTERM 正在关停(重启 / 更新 / 手动 Ctrl-C)。所有在跑的 planner/worker 一并取消;重启后残留的 running 意图会被重置为 open 重跑") + + // ---- 单次 run 级(worker.Execute / planner.Plan 的硬兜底)---- + + AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底触发:某一步卡死了", + "单次运行的硬超时兜底被触发:正常情况下软墙钟预算(MaxDuration)会在回合边界优雅收尾,只有当某一步(模型请求或某个工具)自身卡死、久久不返回、导致回合边界检查根本跑不到时,才会走到「软预算 + 90s 宽限」的硬 cancel。请重点看中断前最后一个 tool_use——多半是它挂住了") +) + +// AbortReason resolves the cause attached to a cancelled run context. ok=false +// means the ctx carries no cause at all (nothing was cancelled). +func AbortReason(ctx context.Context) (code, short, text string, ok bool) { + c := context.Cause(ctx) + if c == nil { + return "", "", "", false + } + if ac, ok := errors.AsType[*AbortCause](c); ok { + return ac.Code, ac.Short, ac.Text, true + } + switch { + case errors.Is(c, context.DeadlineExceeded): + return "deadline_exceeded", "上游 context 到达 deadline(未挂具名原因)", + "上游 context 到达 deadline,但设置方没有用 WithTimeoutCause 挂具名原因: " + c.Error(), true + case errors.Is(c, context.Canceled): + // a plain Canceled means some cancel site still lacks a cause — say so + // plainly instead of inventing one. + return "canceled_no_cause", "取消方未挂具名原因(漏网的取消点)", + "上游 context 被取消,但取消方没有用 context.WithCancelCause 挂上具名原因。" + + "这属于漏网的取消点——请在 agent/cancelcause.go 补一条枚举并接到该取消处", true + } + return "other", firstLine(c.Error(), 80), c.Error(), true +} diff --git a/agent/capture.go b/agent/capture.go index 419e855..bbede4f 100644 --- a/agent/capture.go +++ b/agent/capture.go @@ -3,6 +3,8 @@ package agent import ( "context" "strings" + "time" + "unicode/utf8" "github.com/Autumn-27/norma/agentcore" "github.com/Autumn-27/norma/harness" @@ -64,13 +66,18 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tbuf.WriteString(text) } + // breadcrumbs for the terminal message: when a run is aborted or times out, + // "which tool was still running, for how long" is the whole diagnosis. + lastTool := &runTrace{startedAt: time.Now()} + var finalText string var rerr error for ev, err := range s.Prompt(ctx, input) { if err != nil { flush() - if ctx.Err() != nil { // ctx cancelled = manual stop, not a failure - rec(db.Activity{Kind: "result", Summary: "(已手动停止本次运行)"}) + if ctx.Err() != nil { // ctx cancelled = stopped by the engine, not a failure + sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool) + rec(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}) return finalText, reason, ctx.Err() } rec(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()}) @@ -84,6 +91,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, flush() toolNames[ev.ToolUse.ID] = ev.ToolUse.Name in := string(ev.ToolUse.Input) + lastTool.start(ev.ToolUse.ID, ev.ToolUse.Name, in) rec(db.Activity{Kind: "tool_use", Tool: ev.ToolUse.Name, ToolUseID: ev.ToolUse.ID, Summary: ev.ToolUse.Name + " " + firstLine(in, 200), Detail: in}) case harness.KindToolResult: @@ -92,6 +100,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, } flush() out := blocksText(ev.ToolResult.Content) + lastTool.done(ev.ToolResult.ToolUseID) rec(db.Activity{Kind: "tool_result", Tool: toolNames[ev.ToolResult.ToolUseID], ToolUseID: ev.ToolResult.ToolUseID, IsError: ev.ToolResult.IsError, Summary: firstLine(out, 200), Detail: out}) case harness.KindText: @@ -120,25 +129,13 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tkind = "" } flush() // flush any trailing thinking / non-final text - sum := ev.Terminal.Text + sum, detail := ev.Terminal.Text, ev.Terminal.Text if sum == "" { - switch ev.Terminal.Reason { - case harness.ReasonAbortedTools, harness.ReasonAbortedStreaming: - // interruption (task paused / worker killed / backend restart), - // not a real completion — the intent is re-queued (open) or - // stopped, so the session icon reflects that, not this run. - sum = "(运行被中断:任务暂停 / 终止 / 重启所致;未完成)" - case harness.ReasonTimeout, harness.ReasonMaxTurns: - // budget hit → settlement ran (facts/assets were written back); - // only the trailing text summary is missing, not a failure. - sum = "(达运行预算上限,已收尾写回事实;本次无文字总结,终态: " + string(ev.Terminal.Reason) + ")" - default: - sum = "(无总结,终态: " + string(ev.Terminal.Reason) + ")" - } + sum, detail = terminalText(ctx, ev.Terminal, lastTool) } u := ev.Terminal.Usage // cumulative token usage for this session rec(db.Activity{Kind: "result", IsError: ev.Terminal.Err != nil, - Summary: firstLine(sum, 400), Detail: sum, + Summary: firstLine(sum, 400), Detail: detail, InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens, CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens}) if ev.Terminal.Err != nil { @@ -166,13 +163,15 @@ func blocksText(blocks []llm.ContentBlock) string { } // firstLine returns a single-line, length-capped preview for the summary column. +// max counts RUNES, not bytes: these previews are mostly 中文, and byte-slicing +// used to cut a multi-byte rune in half and render as "恢�…". func firstLine(s string, max int) string { s = strings.TrimSpace(s) - if i := strings.IndexByte(s, '\n'); i >= 0 { - s = s[:i] + if before, _, found := strings.Cut(s, "\n"); found { + s = before } - if len(s) > max { - s = s[:max] + "…" + if utf8.RuneCountInString(s) > max { + s = string([]rune(s)[:max]) + "…" } return s } diff --git a/agent/planner.go b/agent/planner.go index 9b46c30..b567d30 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -334,7 +334,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts runCtx := ctx if maxDur > 0 { var cancel context.CancelFunc - runCtx, cancel = context.WithTimeout(ctx, maxDur+settleHardGrace) + runCtx, cancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout) defer cancel() } _, _, err = captureRun(runCtx, opts, input, diff --git a/agent/terminalreason.go b/agent/terminalreason.go new file mode 100644 index 0000000..90228b2 --- /dev/null +++ b/agent/terminalreason.go @@ -0,0 +1,160 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/Autumn-27/norma/harness" +) + +// This file turns a run's terminal state into something a human can act on. +// +// Before: a cancelled run produced one fixed sentence — "(运行被中断:任务暂停 / +// 终止 / 重启所致;未完成)" — a guess-list, not a diagnosis. The information to do +// better was already there (the cancel cause, the terminal reason, the turn +// count, the tool that never returned); it was just being thrown away. +// +// terminalText produces a self-contained ONE-LINE summary (the transcript's +// collapsed row, cut at the first newline by firstLine) plus a multi-line detail +// block rendered in full when the row is expanded. + +// runTrace records which tool call is in flight, so an aborted / timed-out run can +// name the step it died on — the single most useful fact when a run hangs. +type runTrace struct { + startedAt time.Time + id string // tool_use id of the most recent call + name string // its tool name + input string // its raw input JSON + at time.Time + pending bool // true = no tool_result arrived for it yet +} + +func (t *runTrace) start(id, name, input string) { + t.id, t.name, t.input, t.at, t.pending = id, name, input, time.Now(), true +} + +func (t *runTrace) done(id string) { + if id == t.id { + t.pending = false + } +} + +// reasonHint explains each harness terminal reason in plain 中文: what it means, +// whether results were still written back, and what the engine does next. Covers +// every TerminalReason the harness can emit. +var reasonHint = map[harness.TerminalReason]string{ + harness.ReasonCompleted: "模型主动结束了本轮(正常收场),但没有留下文字总结。事实/资产以工具调用为准,看上面的 tool_use 记录", + harness.ReasonMaxTurns: "撞到步数上限(MaxTurns):SDK 已自动跑了一轮收尾(settlement)把事实/资产写回,只是缺一段文字总结。" + + "意图标记为 exhausted——表示「这个方向试过但没做完」,规划者会据此换角度,不是失败", + harness.ReasonTimeout: "撞到单次运行的软墙钟预算(MaxDuration):SDK 在回合边界优雅收尾并写回事实/资产,只是缺文字总结。" + + "意图标记为 exhausted,不是失败", + harness.ReasonModelError: "模型/API 调用失败(网络、鉴权、限流、供应商 5xx 等)。引擎会退避后自动重跑几次;重试用尽仍失败则意图标记为 blocked", + harness.ReasonBlockingLimit: "上下文长度撞硬上限,请求在发出前就被拦下。通常是单次运行塞进了过多工具输出——考虑收窄意图粒度或压缩工具返回", + harness.ReasonPromptTooLong: "提示词过长且压缩(compaction)重试已用尽,无法继续。同样指向单次运行累积的上下文过大", + harness.ReasonImageError: "多模态内容不被当前模型支持(截图/图片块)。换支持视觉的模型,或让工具别回图", + harness.ReasonStopHookPrevented: "Stop 钩子阻止了本轮结束(守卫判定还不能收场),且随后没能继续下去。" + + "检查任务的 Guard 规则是否过严", + harness.ReasonHookStopped: "某个工具/钩子主动叫停了继续执行(守卫拦截,如越界目标、禁用命令)。看上面最后一条 tool_result 里的拦截说明", + harness.ReasonAbortedStreaming: "运行在【模型输出流式生成阶段】被取消", + harness.ReasonAbortedTools: "运行在【工具执行阶段】被取消", +} + +// terminalText renders a terminal event with no final text into (summary, detail). +// summary is one self-contained line; detail is the expandable breakdown. +func terminalText(ctx context.Context, term *harness.Terminal, tr *runTrace) (string, string) { + reason := term.Reason + aborted := reason == harness.ReasonAbortedStreaming || reason == harness.ReasonAbortedTools + + // ---- one-line summary: cause first (that's what the reader is after), then + // just enough run shape to judge it without expanding. ---- + var sum string + switch { + case aborted: + _, short, _, ok := AbortReason(ctx) + if !ok { + short = "取不到取消原因(run 的 context 上没有挂 cause)" + } + stage := "工具执行阶段" + if reason == harness.ReasonAbortedStreaming { + stage = "模型输出阶段" + } + sum = "(运行被中断:" + short + ";停在" + stage + progressSuffix(term, tr) + ",未完成)" + case reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout: + sum = "(达运行预算上限(" + string(reason) + "),已收尾写回事实" + progressSuffix(term, tr) + ";本次无文字总结)" + default: + hint := reasonHint[reason] + if hint == "" { + hint = "未知终态(harness 新增了终态但这里还没登记,请补 reasonHint)" + } + sum = "(无文字总结,终态 " + string(reason) + ":" + firstLine(hint, 80) + ")" + } + + // ---- expandable detail ---- + var b strings.Builder + b.WriteString(sum) + b.WriteString("\n\n") + fmt.Fprintf(&b, "- **终态**: `%s` — %s\n", reason, reasonHint[reason]) + if aborted { + code, _, why, ok := AbortReason(ctx) + if ok { + fmt.Fprintf(&b, "- **中断原因** (`%s`): %s\n", code, why) + } else { + b.WriteString("- **中断原因**: 取不到——取消方没有用 context.WithCancelCause 挂上具名原因。" + + "这属于漏网的取消点,请补 agent/cancelcause.go 里的枚举\n") + } + } + if term.Err != nil { + fmt.Fprintf(&b, "- **底层错误**: `%v`\n", term.Err) + } + if term.Turns > 0 { + fmt.Fprintf(&b, "- **已执行**: %d 轮模型回合\n", term.Turns) + } + if !tr.startedAt.IsZero() { + fmt.Fprintf(&b, "- **本次运行耗时**: %s\n", roundDur(time.Since(tr.startedAt))) + } + if u := term.Usage; u.InputTokens+u.OutputTokens+u.CacheReadTokens+u.CacheWriteTokens > 0 { + fmt.Fprintf(&b, "- **累计 token**: 输入 %d / 输出 %d / 缓存读 %d / 缓存写 %d\n", + u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheWriteTokens) + } + if tr.name != "" { + if tr.pending { + fmt.Fprintf(&b, "- **中断时正在执行的工具**: `%s`(已跑 %s,**未返回结果**) — 若是硬超时/长时间无进展,问题多半就在这一步\n ```\n %s\n ```\n", + tr.name, roundDur(time.Since(tr.at)), firstLine(tr.input, 300)) + } else { + fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`(已正常返回)\n", tr.name) + } + } else { + b.WriteString("- **工具调用**: 本次运行还没发出任何工具调用就结束了\n") + } + return sum, b.String() +} + +// progressSuffix adds "已跑 N 轮 / 12m30s" to a summary line — how far the run got +// before it died, which decides whether an abort cost anything worth re-running. +func progressSuffix(term *harness.Terminal, tr *runTrace) string { + var parts []string + if term.Turns > 0 { + parts = append(parts, fmt.Sprintf("%d 轮", term.Turns)) + } + if !tr.startedAt.IsZero() { + parts = append(parts, roundDur(time.Since(tr.startedAt))) + } + if len(parts) == 0 { + return "" + } + return ",已跑 " + strings.Join(parts, " / ") +} + +// roundDur formats a duration at a readable granularity (no 3.400213s noise). +func roundDur(d time.Duration) string { + switch { + case d < time.Minute: + return d.Round(100 * time.Millisecond).String() + case d < time.Hour: + return d.Round(time.Second).String() + default: + return d.Round(time.Minute).String() + } +} diff --git a/agent/terminalreason_test.go b/agent/terminalreason_test.go new file mode 100644 index 0000000..a3ecd20 --- /dev/null +++ b/agent/terminalreason_test.go @@ -0,0 +1,172 @@ +package agent + +import ( + "context" + "errors" + "strings" + "testing" + "time" + "unicode/utf8" + + "github.com/Autumn-27/norma/harness" + "github.com/Autumn-27/norma/llm" +) + +// every TerminalReason the harness can emit must have an explanation — otherwise +// a run silently falls back to "未知终态", which is the bug this file exists to fix. +func TestReasonHintCoversEveryTerminalReason(t *testing.T) { + all := []harness.TerminalReason{ + harness.ReasonCompleted, harness.ReasonBlockingLimit, harness.ReasonImageError, + harness.ReasonModelError, harness.ReasonAbortedStreaming, harness.ReasonAbortedTools, + harness.ReasonPromptTooLong, harness.ReasonStopHookPrevented, harness.ReasonHookStopped, + harness.ReasonMaxTurns, harness.ReasonTimeout, + } + for _, r := range all { + if strings.TrimSpace(reasonHint[r]) == "" { + t.Errorf("terminal reason %q 没有中文说明", r) + } + } +} + +// the cause attached at cancel time must survive the ctx chain the engine builds: +// task exec ctx -> per-work child -> value ctx -> the run's hard-timeout backstop. +func TestAbortCausePropagatesThroughRunContextChain(t *testing.T) { + type key struct{} + exec, cancelExec := context.WithCancelCause(context.Background()) + work, cancelWork := context.WithCancelCause(exec) + defer cancelWork(nil) + valued := context.WithValue(work, key{}, "task-1") + run, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) + defer cancelRun() + + cancelExec(AbortPausedByUser) + + code, _, text, ok := AbortReason(run) + if !ok || code != "paused_by_user" { + t.Fatalf("code=%q ok=%v,期望 paused_by_user", code, ok) + } + if !strings.Contains(text, "退回 frontier") { + t.Errorf("原因文案没传下来: %q", text) + } +} + +func TestAbortReasonFallbacks(t *testing.T) { + if _, _, _, ok := AbortReason(context.Background()); ok { + t.Error("未取消的 ctx 不应报出中断原因") + } + // a cancel site that forgot to attach a cause must be called out, not guessed at + plain, cancel := context.WithCancel(context.Background()) + cancel() + code, _, _, ok := AbortReason(plain) + if !ok || code != "canceled_no_cause" { + t.Errorf("code=%q ok=%v,期望 canceled_no_cause", code, ok) + } + + timed, cancelT := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancelT() + <-timed.Done() + if code, _, _, _ := AbortReason(timed); code != "deadline_exceeded" { + t.Errorf("code=%q,期望 deadline_exceeded", code) + } +} + +// an aborted run's one-line summary must name the actual cause, and the detail +// must pin down the tool that never came back. +func TestTerminalTextAbortedNamesCauseAndHangingTool(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortKilledByPlanner) + + tr := &runTrace{startedAt: time.Now().Add(-90 * time.Second)} + tr.start("tu_1", "Bash", `{"command":"nmap -p- 10.0.0.1"}`) + + term := &harness.Terminal{ + Reason: harness.ReasonAbortedTools, + Err: context.Canceled, + Turns: 7, + Usage: llm.Usage{InputTokens: 1200, OutputTokens: 340}, + } + sum, detail := terminalText(ctx, term, tr) + + if !strings.Contains(sum, "kill_work") { + t.Errorf("摘要没点出取消方: %q", sum) + } + if strings.Contains(sum, "\n") { + t.Errorf("摘要必须是单行(前端按首行截断): %q", sum) + } + for _, want := range []string{"killed_by_planner", "aborted_tools", "7 轮", "Bash", "未返回结果", "1200"} { + if !strings.Contains(detail, want) { + t.Errorf("详情缺少 %q:\n%s", want, detail) + } + } +} + +// a tool that DID return must not be reported as the hang site. +func TestTerminalTextCompletedToolNotBlamed(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortShutdown) + tr := &runTrace{startedAt: time.Now()} + tr.start("tu_1", "Read", `{"path":"/etc/hosts"}`) + tr.done("tu_1") + + _, detail := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonAbortedStreaming}, tr) + if strings.Contains(detail, "未返回结果") { + t.Errorf("已返回的工具被误报为卡住:\n%s", detail) + } + if !strings.Contains(detail, "已正常返回") { + t.Errorf("详情没记录最后一个工具:\n%s", detail) + } +} + +// non-abort terminals keep their own explanation instead of the interruption text. +func TestTerminalTextNonAbortReasons(t *testing.T) { + ctx := context.Background() + tr := &runTrace{startedAt: time.Now()} + + sum, _ := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonMaxTurns}, tr) + if !strings.Contains(sum, "达运行预算上限") || strings.Contains(sum, "中断") { + t.Errorf("max_turns 摘要不对: %q", sum) + } + + sum, detail := terminalText(ctx, &harness.Terminal{ + Reason: harness.ReasonModelError, Err: errors.New("429 rate limited")}, tr) + if !strings.Contains(sum, "model_error") { + t.Errorf("model_error 摘要不对: %q", sum) + } + if !strings.Contains(detail, "429 rate limited") { + t.Errorf("详情没带上底层错误:\n%s", detail) + } +} + +// every enumerated cause must carry both a compact headline (for the collapsed +// row) and a full explanation (for the expanded detail). +func TestAbortCausesAreWellFormed(t *testing.T) { + all := []*AbortCause{ + AbortPausedByUser, AbortPausedByOrchestrator, AbortTaskDeleted, AbortPausedOnReload, + AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkFinished, + AbortPausedRaceGuard, AbortChatStoppedByUser, AbortChatTurnFinished, + AbortShutdown, AbortRunHardTimeout, + } + seen := map[string]bool{} + for _, c := range all { + switch { + case c.Code == "" || seen[c.Code]: + t.Errorf("code 缺失或重复: %q", c.Code) + case c.Short == "" || utf8.RuneCountInString(c.Short) > 40: + t.Errorf("%s: Short 必须存在且 ≤40 字(进折叠摘要行): %q", c.Code, c.Short) + case utf8.RuneCountInString(c.Text) <= utf8.RuneCountInString(c.Short): + t.Errorf("%s: Text 应比 Short 更完整", c.Code) + } + seen[c.Code] = true + } +} + +// summaries are mostly 中文; capping by bytes used to slice a rune in half. +func TestFirstLineCapsByRuneNotByte(t *testing.T) { + got := firstLine(strings.Repeat("恢复", 10), 5) + if got != "恢复恢复恢…" { + t.Errorf("firstLine 按字节截断了: %q", got) + } + if got := firstLine("头\n尾", 100); got != "头" { + t.Errorf("多行没截到首行: %q", got) + } +} diff --git a/agent/worker.go b/agent/worker.go index d21d246..71e1eb9 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -361,7 +361,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. runCtx := ctx if maxDur > 0 { var runCancel context.CancelFunc - runCtx, runCancel = context.WithTimeout(ctx, maxDur+settleHardGrace) + runCtx, runCancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout) defer runCancel() } _, reason, err := captureRunSession(runCtx, s, input, emitWrap) diff --git a/cmd/artex/main.go b/cmd/artex/main.go index 30440d4..b5579dc 100644 --- a/cmd/artex/main.go +++ b/cmd/artex/main.go @@ -16,6 +16,7 @@ import ( "syscall" "time" + "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/config" "github.com/Autumn-27/artex/server" ) @@ -70,8 +71,16 @@ func main() { log.Printf("[config] 配置文件: %s (不存在 — 将仅尝试环境变量 ARTEX_PG_DSN)", cfgPath) } - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + sigCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() + // re-wrap with a named cause so a run cancelled by shutdown says "后端进程正在 + // 关停" in its trace instead of a bare context.Canceled (see agent/cancelcause.go). + ctx, shutdown := context.WithCancelCause(sigCtx) + defer shutdown(agent.AbortShutdown) + go func() { + <-sigCtx.Done() + shutdown(agent.AbortShutdown) + }() mgr, err := server.NewManager(*dataDir, *proxy) if err != nil { diff --git a/server/conversations.go b/server/conversations.go index 7b0c77e..69294b8 100644 --- a/server/conversations.go +++ b/server/conversations.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" "github.com/Autumn-27/artex/intercept" ) @@ -243,7 +244,7 @@ func (s *Server) pgStopConversation(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"status": "idle"}) // nothing running return } - cancel() + cancel(agent.AbortChatStoppedByUser) writeJSON(w, 200, map[string]any{"status": "stopping"}) } @@ -331,12 +332,12 @@ func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) { func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) { // Per-run cancellable context so a manual stop (pgStopConversation) can abort // just this session. Registered under chatMu so the stop handler can find it. - ctx, cancel := context.WithCancel(intercept.WithConvID(s.ctx, c.ID)) + ctx, cancel := context.WithCancelCause(intercept.WithConvID(s.ctx, c.ID)) s.chatMu.Lock() s.chatCancel[busyKey] = cancel s.chatMu.Unlock() defer func() { - cancel() + cancel(agent.AbortChatTurnFinished) s.chatMu.Lock() delete(s.chatBusy, busyKey) delete(s.chatCancel, busyKey) diff --git a/server/engine.go b/server/engine.go index 699b204..205f270 100644 --- a/server/engine.go +++ b/server/engine.go @@ -97,14 +97,17 @@ type Engine struct { // per-task execution context: each planner.Plan / worker.Execute runs under it, // so pausing can CANCEL an in-flight run (not just skip the next one). Recreated // on resume since cancelling is one-shot. + // CancelCauseFunc (not CancelFunc): every cancel site attaches a named + // *agent.AbortCause so the run's trace can say WHICH of the ~10 cancel + // paths fired instead of one catch-all "运行被中断" line. execMu sync.Mutex - execCancel map[string]context.CancelFunc + execCancel map[string]context.CancelCauseFunc execCtx map[string]context.Context // per-work (per running intent) cancel, so the planner's kill_work can stop a // single worker without pausing the whole task. Keyed by globally-unique intent id. workMu sync.Mutex - workCancel map[int64]context.CancelFunc + workCancel map[int64]context.CancelCauseFunc // steerBox queues planner course-corrections for a running work (keyed by intent // id). The worker's PreToolUse hook drains it before its next tool call and hands @@ -136,18 +139,21 @@ func (e *Engine) nextPlannerRound(taskID string) int { // Pause stops a task: marks it paused AND cancels any in-flight planner/worker run // for it (a long worker.Execute would otherwise keep going until it finishes). -func (e *Engine) Pause(taskID string) { +// cause says WHY (user button / orchestrator tool / delete / reload) — it surfaces +// verbatim in the cancelled run's trace, so pass the matching agent.Abort* value. +func (e *Engine) Pause(taskID string, cause error) { e.paused.Store(taskID, true) - e.cancelExec(taskID) + e.cancelExec(taskID, cause) } // cancelExec cancels a task's current per-task exec context (any in-flight // planner.Plan / worker.Execute), if present. Shared by Pause and the settle -// sequence's hard-drain backstop. -func (e *Engine) cancelExec(taskID string) { +// sequence's hard-drain backstop. cause is recovered by the aborted run via +// context.Cause and printed in its result record. +func (e *Engine) cancelExec(taskID string, cause error) { e.execMu.Lock() if cancel := e.execCancel[taskID]; cancel != nil { - cancel() + cancel(cause) } e.execMu.Unlock() } @@ -166,14 +172,14 @@ func (e *Engine) execContextFor(parent context.Context, taskID string) context.C defer e.execMu.Unlock() if e.IsPaused(taskID) { // never hand out a live context while paused (guards the claim→Execute race) - c, cancel := context.WithCancel(parent) - cancel() + c, cancel := context.WithCancelCause(parent) + cancel(agent.AbortPausedRaceGuard) return c } if c := e.execCtx[taskID]; c != nil && c.Err() == nil { return c } - c, cancel := context.WithCancel(parent) + c, cancel := context.WithCancelCause(parent) e.execCtx[taskID] = c e.execCancel[taskID] = cancel return c @@ -204,12 +210,12 @@ func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix( func NewEngine(m *Manager) *Engine { return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(), - execCancel: map[string]context.CancelFunc{}, execCtx: map[string]context.Context{}, - workCancel: map[int64]context.CancelFunc{}, steerBox: map[int64][]string{}} + execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{}, + workCancel: map[int64]context.CancelCauseFunc{}, steerBox: map[int64][]string{}} } // registerWork records the cancel for the work currently running intentID. -func (e *Engine) registerWork(intentID int64, cancel context.CancelFunc) { +func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) { e.workMu.Lock() e.workCancel[intentID] = cancel e.workMu.Unlock() @@ -220,7 +226,7 @@ func (e *Engine) unregisterWork(intentID int64) { e.workMu.Lock() if c := e.workCancel[intentID]; c != nil { delete(e.workCancel, intentID) - c() // release context resources (no-op if already cancelled) + c(agent.AbortWorkFinished) // release context resources (no-op if already cancelled) } e.workMu.Unlock() e.steerMu.Lock() @@ -306,7 +312,7 @@ func (e *Engine) KillWork(intentID int64) error { if c == nil { return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID) } - c() + c(agent.AbortKilledByPlanner) return nil } @@ -530,7 +536,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { // 任务已判完成 → 立刻取消在跑的 worker:它们手头的意图跑出来也没意义了。 // 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支 // 归为 stopped(而非 blocked)。 - e.cancelExec(t.ID) + e.cancelExec(t.ID, agent.AbortGoalMet) default: log.Printf("[planner] task %s 规划完成", t.ID) } @@ -597,7 +603,7 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { emit := func(r db.Activity) { e.emitActivity(t, r) } ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline // per-work child context so the planner's kill_work can stop just this work. - workCtx, workCancel := context.WithCancel(ectx) + workCtx, workCancel := context.WithCancelCause(ectx) e.registerWork(intent.ID, workCancel) // wrap the guard hooks so steer_work can inject a mid-run course-correction // for THIS intent (drained before the worker's next tool call). diff --git a/server/engine_cancelcause_test.go b/server/engine_cancelcause_test.go new file mode 100644 index 0000000..0669679 --- /dev/null +++ b/server/engine_cancelcause_test.go @@ -0,0 +1,96 @@ +package server + +import ( + "context" + "testing" + + "github.com/Autumn-27/artex/agent" +) + +// The engine's cancel paths must hand a NAMED cause to the run being cancelled — +// that cause is the whole content of the "运行被中断" line in the trace. These +// exercise the context machinery only (no Manager / DB needed). + +func TestPauseCancelsWithNamedCause(t *testing.T) { + e := NewEngine(nil) + ectx := e.execContextFor(context.Background(), "t1") + // the per-work child the worker loop actually runs under + workCtx, workCancel := context.WithCancelCause(ectx) + defer workCancel(nil) + + e.Pause("t1", agent.AbortPausedByUser) + + if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "paused_by_user" { + t.Fatalf("work ctx 拿到的原因 code=%q ok=%v,期望 paused_by_user", code, ok) + } +} + +func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) { + e := NewEngine(nil) + ectx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(ectx) + e.registerWork(42, workCancel) + + if err := e.KillWork(42); err != nil { + t.Fatalf("KillWork: %v", err) + } + if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "killed_by_planner" { + t.Fatalf("code=%q ok=%v,期望 killed_by_planner", code, ok) + } + // the task ctx must survive — kill_work stops one work, not the whole task + if ectx.Err() != nil { + t.Error("kill_work 不该取消整个任务的 exec ctx") + } + if err := e.KillWork(999); err == nil { + t.Error("对没有在跑的意图 KillWork 应该报错") + } +} + +func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) { + for _, tc := range []struct { + name string + cause *agent.AbortCause + want string + }{ + {"settle", agent.AbortSettleDrainTimeout, "settle_drain_timeout"}, + {"goalMet", agent.AbortGoalMet, "goal_met"}, + } { + t.Run(tc.name, func(t *testing.T) { + e := NewEngine(nil) + ectx := e.execContextFor(context.Background(), "t1") + e.cancelExec("t1", tc.cause) + if code, _, _, _ := agent.AbortReason(ectx); code != tc.want { + t.Fatalf("code=%q,期望 %q", code, tc.want) + } + }) + } +} + +// asking for an exec ctx while paused must return an already-cancelled ctx that +// still SAYS why — this is the claim→Execute race guard, easily mistaken for a +// mystery abort. +func TestExecContextWhilePausedCarriesRaceGuardCause(t *testing.T) { + e := NewEngine(nil) + e.Pause("t1", agent.AbortPausedByUser) + ctx := e.execContextFor(context.Background(), "t1") + if ctx.Err() == nil { + t.Fatal("暂停期间不该发出可用的 exec ctx") + } + if code, _, _, _ := agent.AbortReason(ctx); code != "paused_race_guard" { + t.Fatalf("code=%q,期望 paused_race_guard", code) + } +} + +// Resume must hand out a live context again (cancelling is one-shot). +func TestResumeYieldsFreshContext(t *testing.T) { + e := NewEngine(nil) + old := e.execContextFor(context.Background(), "t1") + e.Pause("t1", agent.AbortPausedByUser) + if old.Err() == nil { + t.Fatal("Pause 没取消旧 ctx") + } + e.paused.Delete("t1") // Resume() without the Task/Notify plumbing + if fresh := e.execContextFor(context.Background(), "t1"); fresh.Err() != nil { + t.Fatal("恢复后应拿到全新可用的 exec ctx") + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 14fa3fe..751c61d 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -159,7 +159,7 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { for e.inflightCount(t.ID) > 0 { if time.Now().After(hardStop) { log.Printf("[deadline] task %s drain 超时(%s),硬取消在跑 run", t.ID, settleDrainGrace) - e.cancelExec(t.ID) + e.cancelExec(t.ID, agent.AbortSettleDrainTimeout) _ = sleepCtx(ctx, 3*time.Second) // 给 worker 分支一点时间落库/归类 break } diff --git a/server/orchestration.go b/server/orchestration.go index 7d3b280..011b0f6 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -304,7 +304,7 @@ func (s *Server) toolPauseTask() actool.CoreTool { if _, ok := s.m.Task(a.TaskID); !ok { return actool.Errorf("task 不存在: " + a.TaskID), nil } - s.engine.Pause(a.TaskID) + s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator) return actool.Text("task paused: " + a.TaskID), nil }) } diff --git a/server/server.go b/server/server.go index 5e8418d..55c40f6 100644 --- a/server/server.go +++ b/server/server.go @@ -60,7 +60,7 @@ type Server struct { // convBusyKey), so a manual stop can abort JUST that session's agent run. Set/ // cleared alongside chatBusy under chatMu. Aborting a run does NOT touch the P3 // trigger queue — the drain goroutine simply proceeds to the next queued fire. - chatCancel map[string]context.CancelFunc + chatCancel map[string]context.CancelCauseFunc // triggerQ buffers P3 trigger fires PER AGENT. A per-agent "pump" launches runs up // to a concurrency limit derived from the agent's策略: serial → limit 1 (+ optional @@ -123,7 +123,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi log.Fatalf("[auth] JWT key: %v", err) } s := &Server{m: m, engine: NewEngine(m), ctx: ctx, skillDir: skillDir, jwtKey: key, chatBusy: map[string]bool{}, - chatCancel: map[string]context.CancelFunc{}, triggerQ: map[string][]triggeredRun{}, + chatCancel: map[string]context.CancelCauseFunc{}, triggerQ: map[string][]triggeredRun{}, triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{}, profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{}, provByProfile: map[int64]*provEntry{}} @@ -217,7 +217,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi log.Printf("[engine] task %s 重置 %d 个残留 running 意图为 open", t.ID, n) } if t.Paused { - s.engine.Pause(t.ID) + s.engine.Pause(t.ID, agent.AbortPausedOnReload) } // 任务级超时:为每个未终态、带 timeout 的任务起 deadline 协调器,独立于 planner/worker // loop——非活跃任务重启后也能在到点后被收尾(deadline 已过则立即走收尾时序)。 @@ -902,7 +902,7 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { } switch req.Action { case "pause": - s.engine.Pause(t.ID) + s.engine.Pause(t.ID, agent.AbortPausedByUser) case "resume": s.engine.Run(s.ctx, t) // ensure loops are alive, then un-pause + nudge s.engine.Resume(t) @@ -2330,7 +2330,7 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { } // Cancellable context so stopChat can abort this turn without affecting // the server's root context (mirrors the conversation stop pattern). - ctx, cancel := context.WithCancel(s.ctx) + ctx, cancel := context.WithCancelCause(s.ctx) s.chatBusy[t.ID] = true s.chatCancel[t.ID] = cancel s.chatMu.Unlock() @@ -2342,7 +2342,7 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // handler returns immediately and the browser never needs to hold the request. go func() { defer func() { - cancel() + cancel(agent.AbortChatTurnFinished) s.chatMu.Lock() delete(s.chatBusy, t.ID) delete(s.chatCancel, t.ID) @@ -2386,7 +2386,7 @@ func (s *Server) stopChat(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"status": "idle"}) return } - cancel() + cancel(agent.AbortChatStoppedByUser) writeJSON(w, 200, map[string]any{"status": "stopping"}) } diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 04ece0a..f755afd 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -68,7 +68,7 @@ func decode(r *http.Request, v any) error { return json.NewDecoder(r.Body).Decod func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) { id := r.PathValue("id") - s.engine.Pause(id) // cancel running loops before removing + s.engine.Pause(id, agent.AbortTaskDeleted) // cancel running loops before removing if err := s.m.DeleteTask(id); err != nil { writeErr(w, 500, err.Error()) return diff --git a/web/src/app/(main)/_components/main-content.tsx b/web/src/app/(main)/_components/main-content.tsx index 8c08bab..3c5a39c 100644 --- a/web/src/app/(main)/_components/main-content.tsx +++ b/web/src/app/(main)/_components/main-content.tsx @@ -62,7 +62,7 @@ export function MainContent({ children }: { children: ReactNode }) {
{version && ( - 版本·v{version} + 版本 · {version} )} diff --git a/web/src/app/(main)/system/llm/page.tsx b/web/src/app/(main)/system/llm/page.tsx index 8e5e491..f01ac53 100644 --- a/web/src/app/(main)/system/llm/page.tsx +++ b/web/src/app/(main)/system/llm/page.tsx @@ -183,7 +183,12 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { value={model} onChange={(e) => setModel(e.target.value)} /> - + {/* modal: this Popover lives inside a Dialog, and its content is + portaled to — outside the Dialog's react-remove-scroll + lock, which then cancels every wheel event over it (the list + rendered but refused to scroll). modal makes the Popover own + the topmost scroll lock, so its own list scrolls again. */} + {models.length > 0 && ( - + {models.map((m) => ( {models.length > 0 && ( - + {models.map((m) => (
@@ -423,6 +437,99 @@ const TaskRow = React.memo(function TaskRow({ ); }); +// ConcurrencySettingsDialog is the 齿轮 button left of 新建任务: toggle the +// simultaneous-running-task cap and pick the limit. Default off; default limit 5. +// Persisted globally via /api/settings — applies to tasks created afterwards. +function ConcurrencySettingsDialog() { + const [open, setOpen] = React.useState(false); + const [enabled, setEnabled] = React.useState(false); + const [limit, setLimit] = React.useState("5"); + const [loading, setLoading] = React.useState(false); + const [saving, setSaving] = React.useState(false); + + // 每次打开都拉取当前值,避免显示陈旧状态。 + React.useEffect(() => { + if (!open) return; + setLoading(true); + api + .settings() + .then((s) => { + setEnabled(!!s.task_concurrency_enabled); + setLimit(String(s.task_concurrency_limit ?? 5)); + }) + .catch(() => {}) + .finally(() => setLoading(false)); + }, [open]); + + async function save() { + const n = Math.max(1, Math.floor(Number(limit) || 5)); + setSaving(true); + try { + await api.setSettings({ task_concurrency_enabled: enabled, task_concurrency_limit: n }); + toast.success(enabled ? `已开启并发限制:最多同时运行 ${n} 个任务` : "已关闭任务并发限制"); + setOpen(false); + } catch (e) { + toast.error("保存失败:" + (e as Error).message); + } finally { + setSaving(false); + } + } + + return ( + + + + + + + 任务并发限制 + + 限制同时「运行中」的任务数量。关闭时不限制;开启后新建任务若已达上限,会进入「排队中」,有空位时自动开始。 + + + +
+
+
+ + 默认关闭(不限制并发) +
+ +
+ + {enabled && ( +
+ + setLimit(e.target.value)} + disabled={loading} + /> + 开启后默认 5,最小 1。 +
+ )} +
+ + + + + + + +
+
+ ); +} + // CreateTaskSheet is the 新建任务 drawer. Its form state lives HERE, not in TasksPage: with // 描述/目标 held by the page component every keystroke re-rendered the whole task table // behind the drawer (plus its sticky column and 20 AlertDialog trees), which showed up as @@ -499,7 +606,7 @@ function CreateTaskSheet({ onCreated }: { onCreated: () => void }) { return ( - diff --git a/web/src/lib/status.ts b/web/src/lib/status.ts index 07d4d97..ab3245e 100644 --- a/web/src/lib/status.ts +++ b/web/src/lib/status.ts @@ -54,6 +54,7 @@ const intent: Record = { const task: Record = { created: { label: "已创建", tone: "slate" }, + queued: { label: "排队中", tone: "amber" }, running: { label: "运行中", tone: "blue" }, paused: { label: "已暂停", tone: "amber" }, done: { label: "已完成", tone: "green" }, diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 3444edd..50b578b 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -1,7 +1,7 @@ // ARTEX domain model — types used across the UI. // Derived from the functional spec (section 7: 关键数据形状). -export type TaskStatus = "created" | "running" | "paused" | "done" | "failed" | "timeout"; +export type TaskStatus = "created" | "queued" | "running" | "paused" | "done" | "failed" | "timeout"; export type EngineMode = "exploring" | "paused" | "stalled" | "idle"; export interface Task { @@ -498,6 +498,9 @@ export interface Settings { web_search_proxy?: string; python_interpreter?: string; // 自定义脚本工具的 python 解释器路径(空=运行时检测) workers?: number; // 并发工作 agent 数(默认3);对之后启动的任务生效 + // 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。 + task_concurrency_enabled?: boolean; // 默认 false + task_concurrency_limit?: number; // 开启后默认 5 } // ---- LLM config ---- From 90030051cee68e8ea8f3541fdce3e87f9c81f8e8 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:28:06 -0400 Subject: [PATCH 05/70] =?UTF-8?q?feat(findings):=20=E6=BC=8F=E6=B4=9E?= =?UTF-8?q?=E5=88=A0=E9=99=A4=E5=8A=9F=E8=83=BD(=E4=BA=8C=E6=AC=A1?= =?UTF-8?q?=E7=A1=AE=E8=AE=A4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - db: DeleteFinding(id) 删 findings 行 + 来源探索节点(kind='finding'), 级联清理 edges/node_assets、置空 activity 引用;返回受影响行数。 - server: DELETE /api/exploration/findings/{id}。 - web: api.deleteFinding;三处均可删且带 AlertDialog 二次确认—— 发现列表页独立「操作」列(每行垃圾桶按钮)、漏洞详情页头部删除按钮、 任务详情·发现 Tab 展开区删除按钮。两处列表同源,删一次全部消失。 - 测试 TestDeleteFinding:findings 行 + 来源节点均被删,重复删返回(0,nil)。 go build/vet、tsc 通过。 --- db/findings.go | 21 ++++++ db/findings_test.go | 48 ++++++++++++++ server/server.go | 21 ++++++ .../(main)/function/findings/detail/page.tsx | 47 +++++++++++++- web/src/app/(main)/function/findings/page.tsx | 64 ++++++++++++++++++- .../tasks/detail/_tabs/findings-tab.tsx | 55 ++++++++++++++-- web/src/lib/api.ts | 2 + 7 files changed, 250 insertions(+), 8 deletions(-) diff --git a/db/findings.go b/db/findings.go index e1e7afc..277f371 100644 --- a/db/findings.go +++ b/db/findings.go @@ -308,6 +308,27 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) { return f, nil } +// DeleteFinding removes a finding entirely: the standalone findings row and its +// originating exploration node (kind='finding'), so it disappears from the findings +// list, the per-task 发现 Tab, and the exploration graph alike. Deleting the node +// cascades its edges + node_assets and nulls any activity referencing it. Returns +// rows affected (0 = no finding with that id). +func (d *DB) DeleteFinding(id int64) (int64, error) { + var nodeID *int64 + err := d.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID) + if err == sql.ErrNoRows { + return 0, nil + } + if err != nil { + return 0, err + } + if nodeID != nil { + // best-effort: the finding is already gone; a stray node must not fail the op. + _, _ = d.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, *nodeID) + } + return 1, nil +} + // SetFindingStatus updates one finding's triage state. Returns rows affected. func (d *DB) SetFindingStatus(id int64, status string) (int64, error) { res, err := d.Exec(`UPDATE findings SET status=$1 WHERE id=$2`, status, id) diff --git a/db/findings_test.go b/db/findings_test.go index a978701..4d67c09 100644 --- a/db/findings_test.go +++ b/db/findings_test.go @@ -5,6 +5,54 @@ import ( "testing" ) +// TestDeleteFinding verifies删除漏洞 removes both the findings row and its +// originating exploration node (kind='finding'). +func TestDeleteFinding(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + tk, err := d.CreateTask("删除漏洞测试", "目标", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer d.DeleteTask(tk.ID) + + // seed a finding node in the task's exploration graph, then a findings row on it + es := d.Exploration(tk.ExplorationID) + nodeID, err := es.AddNode(KindFinding, map[string]any{"summary": "x", "severity": "high"}, 5, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + fid, err := d.AddFinding(tk.ID, nodeID, "XSS", "反射型 XSS", "high", "summary", "poc", "worker", nil) + if err != nil { + t.Fatal(err) + } + + n, err := d.DeleteFinding(fid) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("DeleteFinding rows: want 1, got %d", n) + } + if f, _ := d.GetFinding(fid); f != nil { + t.Fatalf("finding row should be gone, got %+v", f) + } + var cnt int + d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE id=$1`, nodeID).Scan(&cnt) + if cnt != 0 { + t.Fatalf("originating finding node should be deleted, still %d", cnt) + } + + // deleting a non-existent finding is a no-op (0 rows), not an error + if n, err := d.DeleteFinding(fid); err != nil || n != 0 { + t.Fatalf("re-delete: want (0,nil), got (%d,%v)", n, err) + } +} + // TestFindingsPageAndStats exercises ListFindingsPage (filter/sort/paging) and // FindingStats against the live dev PG. It tags its rows with a unique vulnclass // so assertions are isolated from any pre-existing data, and cleans up after. diff --git a/server/server.go b/server/server.go index 99f8710..4c68585 100644 --- a/server/server.go +++ b/server/server.go @@ -616,6 +616,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding) mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage) mux.HandleFunc("PATCH /api/exploration/findings/{id}", s.patchFinding) + mux.HandleFunc("DELETE /api/exploration/findings/{id}", s.deleteFinding) mux.HandleFunc("GET /api/exploration/intents", s.intents) mux.HandleFunc("GET /api/exploration/graph", s.explorationGraph) mux.HandleFunc("GET /api/exploration/activity", s.activity) @@ -1623,6 +1624,26 @@ func (s *Server) patchFinding(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, findingFromDB(f, s.resolveAssetIDs(f.AssetIDs))) } +// deleteFinding removes a finding (findings row + originating exploration node). +// id is the standalone findings-table id (DTO finding_id). +func (s *Server) deleteFinding(w http.ResponseWriter, r *http.Request) { + id := int64(atoiDefault(r.PathValue("id"), 0)) + if id <= 0 { + writeErr(w, 400, "bad finding id") + return + } + n, err := s.m.pg.DeleteFinding(id) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if n == 0 { + writeErr(w, 404, "finding not found") + return + } + writeJSON(w, 200, map[string]any{"deleted": true, "id": id}) +} + func (s *Server) intents(w http.ResponseWriter, r *http.Request) { t := s.m.ResolveTask(r.URL.Query().Get("task")) if t == nil { diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index c456ee4..07876b3 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -2,12 +2,23 @@ import * as React from "react"; import Link from "next/link"; -import { useSearchParams } from "next/navigation"; +import { useRouter, useSearchParams } from "next/navigation"; import { toast } from "sonner"; -import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react"; +import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon, Trash2Icon } from "lucide-react"; import { SidebarTrigger } from "@/components/ui/sidebar"; import { Separator } from "@/components/ui/separator"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, + AlertDialogTrigger, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -61,6 +72,7 @@ function FieldRow({ } function FindingDetailInner() { + const router = useRouter(); const searchParams = useSearchParams(); const id = searchParams.get("id") ?? ""; const [finding, setFinding] = React.useState(null); @@ -116,6 +128,18 @@ function FindingDetailInner() { [finding, id], ); + // deleteFinding 删除当前漏洞(需二次确认),成功后返回发现列表。 + const deleteFinding = React.useCallback(async () => { + if (!id) return; + try { + await api.deleteFinding(id); + toast.success("已删除漏洞"); + router.push("/function/findings"); + } catch (e) { + toast.error("删除失败:" + (e as Error).message); + } + }, [id, router]); + if (!finding) { return (
@@ -160,6 +184,25 @@ function FindingDetailInner() { + + + + + + + 确认删除该漏洞? + + 「{title}」将被永久删除,同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 + + + + 取消 + 删除 + + +
概览 diff --git a/web/src/app/(main)/function/findings/page.tsx b/web/src/app/(main)/function/findings/page.tsx index 498714d..fdf4798 100644 --- a/web/src/app/(main)/function/findings/page.tsx +++ b/web/src/app/(main)/function/findings/page.tsx @@ -7,10 +7,22 @@ import { ShieldAlertIcon, ArrowUpRightIcon, FileTextIcon, + Trash2Icon, } from "lucide-react"; import { StatusBadge } from "@/components/status-badge"; import { Markdown } from "@/components/markdown"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, + AlertDialogTrigger, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; @@ -231,6 +243,24 @@ export default function FindingsPage() { [edit], ); + // deleteFinding 删除一个漏洞(需二次确认):删成功后从列表移除、收起行、刷新统计。 + const deleteFinding = React.useCallback( + async (f: Finding) => { + if (!f.finding_id) return; + try { + await api.deleteFinding(f.finding_id); + setFindings((cur) => cur.filter((x) => x.id !== f.id)); + setTotal((t) => Math.max(0, t - 1)); + setExpanded((cur) => (cur === f.id ? null : cur)); + toast.success("已删除漏洞"); + api.findingStats().then(setStats).catch(() => {}); + } catch (e) { + toast.error("删除失败:" + (e as Error).message); + } + }, + [], + ); + const statCards: { label: string; value: number; tone?: string }[] = [ { label: "发现总数", value: stats.total }, { label: "待处理", value: stats.pending, tone: "text-amber-500" }, @@ -375,6 +405,7 @@ export default function FindingsPage() { 状态 所属任务 时间 + 操作 @@ -490,12 +521,41 @@ export default function FindingsPage() { {fmtTime(f.ts)} + e.stopPropagation()}> + {f.finding_id && ( + + + + + + + 确认删除该漏洞? + + 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, + 同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 + + + + 取消 + deleteFinding(f)}>删除 + + + + )} + {open && ( {/* whitespace-normal 覆盖 TableCell 默认的 nowrap,否则展开区文字 被强制单行、直接溢出单元格。 */} - +
{/* 行内编辑:名称/类别/严重等级,可改并保存(仅独立 finding 行)。 */} {f.finding_id && edit && ( @@ -626,7 +686,7 @@ export default function FindingsPage() { {findings.length === 0 && ( 没有匹配的发现。 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 e68aed2..425825b 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 @@ -2,10 +2,22 @@ import * as React from "react"; import Link from "next/link"; -import { ArrowUpRightIcon, ChevronRightIcon } from "lucide-react"; +import { ArrowUpRightIcon, ChevronRightIcon, Trash2Icon } from "lucide-react"; import { cn } from "@/lib/utils"; import { Card, CardContent } from "@/components/ui/card"; import { StatusBadge } from "@/components/status-badge"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, + AlertDialogTrigger, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; import { Select, SelectContent, @@ -31,9 +43,11 @@ const FINDING_STATUSES: FindingStatus[] = [ function Row({ f, onStatus, + onDelete, }: { f: Finding; onStatus: (f: Finding, next: FindingStatus) => void; + onDelete: (f: Finding) => void; }) { const [open, setOpen] = React.useState(false); return ( @@ -116,8 +130,30 @@ function Row({
{open && (
-
- 证据 / PoC +
+ 证据 / PoC + {f.finding_id && ( + + + + + + + 确认删除该漏洞? + + 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, + 同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 + + + + 取消 + onDelete(f)}>删除 + + + + )}
             {f.evidence}
@@ -169,6 +205,17 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     [],
   );
 
+  const onDelete = React.useCallback(async (f: Finding) => {
+    if (!f.finding_id) return;
+    try {
+      await api.deleteFinding(f.finding_id);
+      setFindings((cur) => cur.filter((x) => x.id !== f.id));
+      toast.success("已删除漏洞");
+    } catch (e) {
+      toast.error("删除失败:" + (e as Error).message);
+    }
+  }, []);
+
   const items = findings
     .filter((f) => f.task_id === taskId)
     .sort((a, b) => {
@@ -180,7 +227,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     
       
         {items.map((f) => (
-          
+          
         ))}
         {items.length === 0 && (
           

diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index d2eb7ec..a8eb4db 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -261,6 +261,8 @@ export const api = { id: string, fields: { name?: string; vulnclass?: string; severity?: Severity; status?: FindingStatus }, ) => patch(`/exploration/findings/${id}`, fields), + // 删除漏洞:移除 findings 记录 + 来源探索节点(从发现列表/任务发现 Tab/探索图一并消失)。 + deleteFinding: (id: string) => del<{ deleted: boolean; id: number }>(`/exploration/findings/${id}`), intents: (task?: string) => get(`/exploration/intents${tq(task)}`).then(arr), tokenStats: (task?: string) => get<{ workers: TokenUsage[]; total: TokenTotal }>(`/exploration/tokens${tq(task)}`).then((r) => ({ From 2aeb2e49dafef39606b0b4127fd9158ac3c8cb80 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:30:56 -0400 Subject: [PATCH 06/70] =?UTF-8?q?fix(ui):=20=E5=88=A0=E9=99=A4=E7=A1=AE?= =?UTF-8?q?=E8=AE=A4=E5=BC=B9=E7=AA=97=E9=95=BF=E6=96=87=E6=9C=AC=E6=BA=A2?= =?UTF-8?q?=E5=87=BA(break-words=20+=20break-all)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 任务/漏洞名称含长 URL、路径等无空格断点的长串时,弹窗描述不换行撑出边界。 给 AlertDialogDescription 加 break-words,用户内容包 span break-all 强制断行。 覆盖任务删除 + 三处漏洞删除弹窗。 --- .../(main)/function/findings/detail/page.tsx | 4 ++-- web/src/app/(main)/function/findings/page.tsx | 4 ++-- .../tasks/detail/_tabs/findings-tab.tsx | 4 ++-- web/src/app/(main)/function/tasks/page.tsx | 18 ++++++++++++++---- 4 files changed, 20 insertions(+), 10 deletions(-) diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index 07876b3..19f1c2a 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -193,8 +193,8 @@ function FindingDetailInner() { 确认删除该漏洞? - - 「{title}」将被永久删除,同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 + + 「{title}」将被永久删除,同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 diff --git a/web/src/app/(main)/function/findings/page.tsx b/web/src/app/(main)/function/findings/page.tsx index fdf4798..38b9eda 100644 --- a/web/src/app/(main)/function/findings/page.tsx +++ b/web/src/app/(main)/function/findings/page.tsx @@ -537,8 +537,8 @@ export default function FindingsPage() { 确认删除该漏洞? - - 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, + + 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, 同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 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 425825b..d4faf24 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 @@ -142,8 +142,8 @@ function Row({ 确认删除该漏洞? - - 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, + + 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, 同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 diff --git a/web/src/app/(main)/function/tasks/page.tsx b/web/src/app/(main)/function/tasks/page.tsx index fa7fac9..4cdd004 100644 --- a/web/src/app/(main)/function/tasks/page.tsx +++ b/web/src/app/(main)/function/tasks/page.tsx @@ -416,10 +416,20 @@ const TaskRow = React.memo(function TaskRow({ 确认删除任务 #{task.id}? - - {task.description - ? `「${task.description.length > 80 ? task.description.slice(0, 80) + "…" : task.description}」` - : "该任务"} + + {task.description ? ( + <> + 「 + + {task.description.length > 80 + ? task.description.slice(0, 80) + "…" + : task.description} + + 」 + + ) : ( + "该任务" + )} 的执行记录、会话与产物将被删除,此操作不可撤销(全局资产图保留)。 From 4f7c1a199d26f6a17ed8f0a7f1078ce469ed1f0c Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:40:19 -0400 Subject: [PATCH 07/70] =?UTF-8?q?feat(tasks):=20=E5=88=A0=E9=99=A4?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1=E5=8F=AF=E9=80=89=E8=BF=9E=E5=B8=A6=E5=88=A0?= =?UTF-8?q?=E9=99=A4=E8=B5=84=E4=BA=A7/=E6=BC=8F=E6=B4=9E/=E6=B5=81?= =?UTF-8?q?=E9=87=8F/=E6=96=87=E4=BB=B6(=E9=BB=98=E8=AE=A4=E4=B8=8D?= =?UTF-8?q?=E9=80=89)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 删除弹窗新增四个默认不勾的复选框: - 删除关联的资产:仅删本任务独有资产,与其他任务共享的仅解除关联(保留)。 - 删除发现的漏洞:删本任务全部 findings。 - 删除测试的流量:按本任务资产的 host 清录制流量(host 无任务标签,会连带清 其他任务命中同一 host 的流量)。 - 删除测试过程中写的文件:删 /tasks/ 目录(worker 产物 + 上传)。 - db: AssetStore.DeleteByTaskID(独有删+共享解关联)/HostsByTask;DeleteFindingsByTask。 - server: pgDeleteTask 读可选 JSON body {assets,findings,traffic,files},先按 host 清流量(需在删资产前取 host)→删漏洞→删资产→删任务→删文件目录;空 body=只删任务。 - web: api.deleteTask(id,opts);删除弹窗四复选框(默认不选)+说明。 go build/vet、tsc 通过;新增 TestDeleteByTaskID(独有删/共享解关联/host 反查)。 --- db/assets.go | 39 +++++++++++++ db/assets_test.go | 54 ++++++++++++++++++ db/findings.go | 11 ++++ server/server_mgmt.go | 49 +++++++++++++++- web/src/app/(main)/function/tasks/page.tsx | 65 ++++++++++++++++++---- web/src/lib/api.ts | 11 +++- 6 files changed, 215 insertions(+), 14 deletions(-) diff --git a/db/assets.go b/db/assets.go index 3cf1c1b..583cc87 100644 --- a/db/assets.go +++ b/db/assets.go @@ -1053,6 +1053,45 @@ func (s *AssetStore) GetByIDs(ids []int64) ([]*Asset, error) { return scanAssets(rows) } +// DeleteByTaskID removes a task's assets: hard-deletes assets owned solely by this +// task (task_ids = {taskID}) and de-associates the rest (drops this task from their +// task_ids, keeping assets shared with other tasks intact). Returns rows deleted. +func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) { + res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil { + return n, err + } + return n, nil +} + +// HostsByTask returns the distinct hostnames/IPs the task's assets carry (domain, +// root_domain, ip) — used to purge host-keyed traffic recorded for that task. +func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) { + rows, err := s.db.Query(` +SELECT DISTINCT h FROM ( + SELECT NULLIF(domain,'') AS h FROM assets WHERE $1 = ANY(task_ids) + UNION SELECT NULLIF(root_domain,'') FROM assets WHERE $1 = ANY(task_ids) + UNION SELECT NULLIF(ip,'') FROM assets WHERE $1 = ANY(task_ids) +) q WHERE h IS NOT NULL AND h <> ''`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []string + for rows.Next() { + var h string + if err := rows.Scan(&h); err != nil { + return nil, err + } + out = append(out, h) + } + return out, rows.Err() +} + // DeleteByCompanyID hard-deletes all assets belonging to a company. Returns rows deleted. func (s *AssetStore) DeleteByCompanyID(companyID int64) (int64, error) { res, err := s.db.Exec(`DELETE FROM assets WHERE company_id = $1`, companyID) diff --git a/db/assets_test.go b/db/assets_test.go index 2245192..367f7db 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -2,6 +2,7 @@ package db import ( "fmt" + "slices" "testing" ) @@ -535,6 +536,59 @@ func TestQueryByType(t *testing.T) { } } +// TestDeleteByTaskID: 独有资产被删,与其他任务共享的资产仅解除关联(保留),host 反查正确。 +func TestDeleteByTaskID(t *testing.T) { + d, av2, _ := testSetup(t) + defer d.Close() + + const taskA = int64(90001) + const taskB = int64(90002) + // solo:仅属 taskA + solo, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "solo-del.test", TaskID: taskA}) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, solo) + // shared:先 taskA 再 taskB → task_ids={A,B} + shared, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "shared-del.test", TaskID: taskA}) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, shared) + if _, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "shared-del.test", TaskID: taskB}); err != nil { + t.Fatal(err) + } + + // host 反查(删资产前):应含两个域名 + hosts, err := av2.HostsByTask(taskA) + if err != nil { + t.Fatal(err) + } + if !slices.Contains(hosts, "solo-del.test") || !slices.Contains(hosts, "shared-del.test") { + t.Fatalf("HostsByTask 缺 host: %v", hosts) + } + + n, err := av2.DeleteByTaskID(taskA) + if err != nil { + t.Fatal(err) + } + if n != 1 { + t.Fatalf("DeleteByTaskID: 应删 1 个独有资产,实删 %d", n) + } + // solo 已删 + if a, _ := av2.GetByIDs([]int64{solo}); len(a) != 0 { + t.Fatalf("solo 资产应被删除") + } + // shared 保留,且 task_ids 只剩 taskB + sa, _ := av2.GetByIDs([]int64{shared}) + if len(sa) != 1 { + t.Fatalf("shared 资产应保留") + } + if slices.Contains(sa[0].TaskIDs, taskA) || !slices.Contains(sa[0].TaskIDs, taskB) { + t.Fatalf("shared task_ids 应解除 A 保留 B,得 %v", sa[0].TaskIDs) + } +} + func TestQueryByTask(t *testing.T) { d, av2, _ := testSetup(t) defer d.Close() diff --git a/db/findings.go b/db/findings.go index 277f371..6fc64db 100644 --- a/db/findings.go +++ b/db/findings.go @@ -329,6 +329,17 @@ func (d *DB) DeleteFinding(id int64) (int64, error) { return 1, nil } +// DeleteFindingsByTask removes all findings rows of a task. The originating +// exploration finding nodes are cascade-deleted separately when the task's +// exploration subgraph is dropped. Returns rows deleted. +func (d *DB) DeleteFindingsByTask(taskID int64) (int64, error) { + res, err := d.Exec(`DELETE FROM findings WHERE task_id=$1`, taskID) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + // SetFindingStatus updates one finding's triage state. Returns rows affected. func (d *DB) SetFindingStatus(id int64, status string) (int64, error) { res, err := d.Exec(`UPDATE findings SET status=$1 WHERE id=$2`, status, id) diff --git a/server/server_mgmt.go b/server/server_mgmt.go index f755afd..0595787 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -66,14 +66,61 @@ func decode(r *http.Request, v any) error { return json.NewDecoder(r.Body).Decod // ---------- tasks (delete) ---------- +// pgDeleteTask removes a task and, per opt-in flags in the (optional) JSON body, +// its associated data: assets (task-owned), findings, recorded traffic (by the +// task's hosts) and the per-task working directory. Empty body → all flags false → +// only the task + its exploration subgraph are removed (legacy behavior). func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) { id := r.PathValue("id") + var opts struct { + Assets bool `json:"assets"` + Findings bool `json:"findings"` + Traffic bool `json:"traffic"` + Files bool `json:"files"` + } + _ = json.NewDecoder(r.Body).Decode(&opts) // optional body; ignore EOF/parse errors + + result := map[string]any{"deleted": id} + nid, _ := strconv.ParseInt(id, 10, 64) + as := s.m.Assets() + + // Traffic first: it needs the task's hosts, which come from assets that a later + // asset-delete would remove. Traffic is host-keyed (no per-task tag), so this + // purges ALL recorded traffic for those hosts, not just this task's. + if opts.Traffic && nid > 0 && as != nil { + if tr := s.m.Traffic(); tr != nil { + if hosts, err := as.HostsByTask(nid); err == nil && len(hosts) > 0 { + if n, err := tr.DeleteHostsExact(hosts); err == nil { + result["traffic_deleted"] = n + } + } + } + } + if opts.Findings && nid > 0 { + if n, err := s.m.pg.DeleteFindingsByTask(nid); err == nil { + result["findings_deleted"] = n + } + } + if opts.Assets && nid > 0 && as != nil { + if n, err := as.DeleteByTaskID(nid); err == nil { + result["assets_deleted"] = n + } + } + s.engine.Pause(id, agent.AbortTaskDeleted) // cancel running loops before removing if err := s.m.DeleteTask(id); err != nil { writeErr(w, 500, err.Error()) return } - writeJSON(w, 200, map[string]any{"deleted": id}) + + // Files last: the per-task working dir /tasks/ (worker artifacts + + // uploads). Confined to the work dir; id is numeric so no traversal risk. + if opts.Files && nid > 0 { + if err := os.RemoveAll(filepath.Join(s.m.dir, "tasks", id)); err == nil { + result["files_deleted"] = true + } + } + writeJSON(w, 200, result) } // ---------- agents ---------- diff --git a/web/src/app/(main)/function/tasks/page.tsx b/web/src/app/(main)/function/tasks/page.tsx index 4cdd004..718b6de 100644 --- a/web/src/app/(main)/function/tasks/page.tsx +++ b/web/src/app/(main)/function/tasks/page.tsx @@ -227,15 +227,18 @@ export default function TasksPage() { return () => clearInterval(i); }, [hasRunning]); - const deleteTask = React.useCallback(async (id: string) => { - try { - await api.deleteTask(id); - toast.success("任务已删除(全局资产图保留)"); - load(); - } catch (e) { - toast.error("删除失败:" + (e as Error).message); - } - }, [load]); + const deleteTask = React.useCallback( + async (id: string, opts: DeleteOpts) => { + try { + await api.deleteTask(id, opts); + toast.success("任务已删除"); + load(); + } catch (e) { + toast.error("删除失败:" + (e as Error).message); + } + }, + [load], + ); return ( @@ -333,6 +336,15 @@ export default function TasksPage() { // the POLL_MS list refresh only re-render the rows whose data actually moved — a table page // is 20 rows × (StatusBadge + Link + a Radix AlertDialog tree), far too heavy to rebuild // wholesale on every parent render. +type DeleteOpts = { assets: boolean; findings: boolean; traffic: boolean; files: boolean }; + +const DELETE_EXTRAS: { key: keyof DeleteOpts; label: string; hint: string }[] = [ + { key: "assets", label: "删除关联的资产", hint: "仅本任务独有的资产;与其他任务共享的仅解除关联" }, + { key: "findings", label: "删除发现的漏洞", hint: "本任务记录的全部漏洞" }, + { key: "traffic", label: "删除测试的流量", hint: "按本任务涉及的 host 清除录制流量" }, + { key: "files", label: "删除测试过程中写的文件", hint: "本任务工作目录 tasks/(产物 + 上传)" }, +]; + const TaskRow = React.memo(function TaskRow({ task, nowSec, @@ -340,8 +352,14 @@ const TaskRow = React.memo(function TaskRow({ }: { task: Task; nowSec: number; - onDelete: (id: string) => void; + onDelete: (id: string, opts: DeleteOpts) => void; }) { + const [delOpts, setDelOpts] = React.useState({ + assets: false, + findings: false, + traffic: false, + files: false, + }); return ( @@ -430,12 +448,35 @@ const TaskRow = React.memo(function TaskRow({ ) : ( "该任务" )} - 的执行记录、会话与产物将被删除,此操作不可撤销(全局资产图保留)。 + 的执行记录、会话与产物将被删除,此操作不可撤销。 + + {/* 额外清理项:默认都不选,勾选后连带删除。全局资产图默认保留。 */} +

+ + 同时删除以下关联数据(默认不删): + + {DELETE_EXTRAS.map((opt) => ( + + ))} +
+ 取消 - onDelete(task.id)}> + onDelete(task.id, delOpts)}> 删除 diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index a8eb4db..fc30e9d 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -122,7 +122,16 @@ export const api = { seed_first_intent: seedFirstIntent ?? false, plan_heartbeat_seconds: planHeartbeatSeconds ?? 0, // 0 = 后端归一到默认 600(10min) }), - deleteTask: (id: string) => del<{ deleted: number }>(`/tasks/${id}`), + // 删除任务;opts 里的四个可选项默认不删:assets 关联资产、findings 发现漏洞、 + // traffic 测试流量(按任务 host 清)、files 测试过程写的文件(tasks/ 目录)。 + deleteTask: ( + id: string, + opts?: { assets?: boolean; findings?: boolean; traffic?: boolean; files?: boolean }, + ) => + del<{ deleted: string; assets_deleted?: number; findings_deleted?: number; traffic_deleted?: number; files_deleted?: boolean }>( + `/tasks/${id}`, + opts, + ), controlTask: (id: string, action: "pause" | "resume") => post<{ id: string; paused: boolean }>(`/tasks/${id}/control`, { action }), // 重跑一条没跑成功的意图(blocked/exhausted/stopped):置回 open,worker 会重新认领、从头再跑。 rerunIntent: (taskId: string, intentId: string) => post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/${intentId}/rerun`), From 89c5e9b26f94cf652d4ef550e4e64ed8c35f6774 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:05:09 -0400 Subject: [PATCH 08/70] =?UTF-8?q?fix(mock):=20intentsPage=20=E8=BF=94?= =?UTF-8?q?=E5=9B=9E=E5=88=86=E9=A1=B5=E6=A0=BC=E5=BC=8F,=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=20worker=20=E5=88=97=E8=A1=A8=E4=B8=BA=E7=A9=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit handler 对 /exploration/intents 只返回裸数组,intentsPage 带 ?page= 参数时期望 {items, has_more} 形状,r.items 取到 undefined 导致 mock 模式下 worker 会话列表始终为空。 --- web/src/lib/mock/handler.ts | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 8b9f313..1d59db2 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -148,7 +148,16 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re page_size: pageSize, }; } - if (path === "/exploration/intents") return D.intents; + if (path === "/exploration/intents") { + if (q.has("page")) { + const before = Number(q.get("before") ?? 0); + const limit = Number(q.get("limit") ?? 300); + let list = D.intents; + if (before > 0) list = list.filter((n) => Number(n.id.replace(/\D/g, "") || n.id) < before); + return { items: list.slice(0, limit), has_more: list.length > limit }; + } + return D.intents; + } if (path === "/exploration/tokens") return { workers: D.tokenWorkers, total: D.tokenTotal }; if (path === "/exploration/graph") return D.explorationGraph; if (path === "/exploration/activity" && seg.length === 2) { @@ -190,6 +199,9 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/llm/profiles" && m === "GET") return { profiles: D.llmProfiles }; if (path === "/llm/profiles" && m === "POST") return { id: Number(b.id) || 3 }; if (path === "/llm/profiles/active") return { ok: true }; + if (path === "/llm/pool" && m === "GET") return D.llmPool; + if (path === "/llm/pool/reset") + return { ...D.llmPool, chain: D.llmPool.chain.map((c) => ({ ...c, state: "ok", fails: 0, cooldown_secs: 0 })) }; if (seg[0] === "llm" && seg[1] === "profiles" && seg.length === 3 && m === "DELETE") return { deleted: Number(seg[2]) }; // ── agents ── From 0c653a213598593b34cb66d69b0b28889c2d3884 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:07:53 -0400 Subject: [PATCH 09/70] =?UTF-8?q?feat(llm):=20LLM=20=E9=85=8D=E7=BD=AE?= =?UTF-8?q?=E8=BD=AE=E8=AF=A2(=E6=95=85=E9=9A=9C=E8=BD=AC=E7=A7=BB),?= =?UTF-8?q?=E6=9F=90=E4=B8=AA=E9=85=8D=E7=BD=AE=E6=B2=A1=E9=A2=9D=E5=BA=A6?= =?UTF-8?q?/=E5=A4=B1=E6=95=88=E6=97=B6=E8=87=AA=E5=8A=A8=E5=88=87?= =?UTF-8?q?=E4=B8=8B=E4=B8=80=E4=B8=AA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 默认关闭。开启后「未指定模型」的 agent 在当前配置不可用(402 余额不足 / 401 key 失效 / 429 限流 / 5xx)时自动切到下一个配置;agent 绑定或任务 pin 了配置的默认不轮询(绑定即独占,可由 llm_pool_bind_fallback 打开兜底)。 链序 = 激活配置 + 其余按 priority DESC,排除 pool_exclude 的。只给 llm_profiles 加两列,不引入关联表,所以 agents/tasks/conversations 对 llm_profile_id 的四处引用一行未改。 核心约束:只在本次调用尚未 yield 出任何事件时切换成员。已经吐出 text_delta / tool_use_start 后中途断流一律原样上抛(交给 agentcore 的 resume),否则重发同一 request 会重复输出、污染会话历史。SDK 的 doStream 保证 4xx/5xx 在读 body 前返回,所以目标场景都落在这个安全窗口内。 context.Canceled / DeadlineExceeded 绝不触发切换——那是用户取消或任务 超时,误判会白烧备用 key 并污染 a6bd180 的终止原因诊断。400 也不切 (请求本身有问题,换模型结果一样)。 熔断:402/401/403/404 一次即熔断,429/5xx/网络错误连续 3 次;冷却 60s→5min→30min 退避,到期半开试探,成功清零。Registry 是进程级的,独立 于 provider 缓存,重建链路不丢状态,且全进程共享(A 任务撞见的 402,B 任务 立刻知道)。状态旁路落库,重启只恢复尚未到期的冷却窗口。 上下文窗口装不下当前请求的成员直接跳过(必然 400),不把主力模型的窗口 压到全链最小值。全员熔断时仍探链首,不让引擎停摆。 Pool 包在 llmrec 外面、每个成员自带 recorder,失败调用会以 status=error + 正确的 profile 名进 llm_records;切换/熔断走 log.Printf,自动出现在 系统日志页。 顺带修两个既有问题: - DecomposeGoals 原本自己 NewProvider(),绕过限速器/录制/轮询,且每个 任务新建一个实例——改为由 server 传入共享 provider - saveLLMConfig 会把 priority/pool_exclude 重置为零值 前端:LLM 页新增轮询卡片(总开关/兜底开关/实际链路顺序/健康徽章/冷却 倒计时/立即恢复),配置编辑区加优先级与「不参与轮询」。 只有一个配置时链长度为 1,poolChain 返回 nil 走裸 provider,行为与开关 打开前完全一致。设计见 docs/LLM轮询设计.md。 --- agent/goals.go | 11 +- db/config.go | 67 ++++- db/config_test.go | 48 ++++ db/llmhealth.go | 57 ++++ db/schema.sql | 20 ++ llmpool/health.go | 149 ++++++++++ llmpool/pool.go | 261 ++++++++++++++++++ llmpool/pool_test.go | 359 +++++++++++++++++++++++++ server/dto.go | 6 + server/goals.go | 23 +- server/llmpool.go | 239 ++++++++++++++++ server/llmpool_test.go | 41 +++ server/manager.go | 36 +++ server/server.go | 70 ++++- server/server_mgmt.go | 44 ++- web/src/app/(main)/system/llm/page.tsx | 306 ++++++++++++++++++++- web/src/lib/api.ts | 8 +- web/src/lib/mock/data.ts | 19 +- web/src/lib/types.ts | 36 +++ 19 files changed, 1750 insertions(+), 50 deletions(-) create mode 100644 db/llmhealth.go create mode 100644 llmpool/health.go create mode 100644 llmpool/pool.go create mode 100644 llmpool/pool_test.go create mode 100644 server/llmpool.go create mode 100644 server/llmpool_test.go diff --git a/agent/goals.go b/agent/goals.go index 1ce68fa..dbbae9f 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -7,6 +7,7 @@ import ( "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/agentcore" + "github.com/Autumn-27/norma/llm" acperm "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" ) @@ -63,6 +64,11 @@ type GoalSpec struct { // no provider is configured or the call yields nothing — the caller then falls // back to a rule-based split so goal nodes always exist. // +// prov is supplied by the caller (rather than built here from a Config) so goal +// decomposition rides the SAME provider instance as the rest of the engine — it +// 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 数量/交战说明等). // 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. @@ -77,9 +83,8 @@ type GoalSpec struct { // straight into it (the same managed tool the main agent uses to add goals at // runtime). The returned specs are read back from the store so callers can emit // per-goal activity and detect the "LLM produced nothing" case for their fallback. -func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { - prov, err := c.NewProvider() - if err != nil { +func DecomposeGoals(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { + if prov == nil { return nil } // worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link diff --git a/db/config.go b/db/config.go index df31ddf..67d9edc 100644 --- a/db/config.go +++ b/db/config.go @@ -27,10 +27,29 @@ type LLMProfile struct { // 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider. ReasoningEffort string `json:"reasoning_effort"` IsDefault bool `json:"is_default"` + // Priority orders the failover chain: higher goes first. The ACTIVE profile + // (IsDefault) always heads the chain regardless of this value. + Priority int `json:"priority"` + // PoolExclude=true keeps this profile out of the failover chain — it stays + // usable when an agent/task binds it explicitly, it just never gets picked up + // as a fallback target. + PoolExclude bool `json:"pool_exclude"` +} + +// profileCols is the read column list (hint variant, no api key) shared by the +// list query; profileColsKey is the same with api_key for the single-row loads. +const profileCols = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude` +const profileColsKey = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude` + +// scanProfile reads one row in the profileCols / profileColsKey column order. The +// 7th column lands in APIKeyHint or APIKey depending on which list the caller used. +func scanProfile(sc interface{ Scan(...any) error }, into *string, p *LLMProfile) error { + return sc.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, into, + &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault, &p.Priority, &p.PoolExclude) } func (d *DB) ListProfiles() ([]*LLMProfile, error) { - rows, err := d.Query(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles ORDER BY id`) + rows, err := d.Query(`SELECT ` + profileCols + ` FROM llm_profiles ORDER BY id`) if err != nil { return nil, err } @@ -38,7 +57,7 @@ func (d *DB) ListProfiles() ([]*LLMProfile, error) { var out []*LLMProfile for rows.Next() { var p LLMProfile - if err := rows.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKeyHint, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault); err != nil { + if err := scanProfile(rows, &p.APIKeyHint, &p); err != nil { return nil, err } out = append(out, &p) @@ -49,8 +68,7 @@ func (d *DB) ListProfiles() ([]*LLMProfile, error) { // ActiveProfile returns the default (active) profile with its api key, or nil. func (d *DB) ActiveProfile() (*LLMProfile, error) { var p LLMProfile - err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE is_default LIMIT 1`). - Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault) + err := scanProfile(d.QueryRow(`SELECT `+profileColsKey+` FROM llm_profiles WHERE is_default LIMIT 1`), &p.APIKey, &p) if err == sql.ErrNoRows { return nil, nil } @@ -61,14 +79,37 @@ func (d *DB) ActiveProfile() (*LLMProfile, error) { // Used to run a task on a specific (non-default) LLM profile. func (d *DB) ProfileByID(id int64) (*LLMProfile, error) { var p LLMProfile - err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE id=$1`, id). - Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault) + err := scanProfile(d.QueryRow(`SELECT `+profileColsKey+` FROM llm_profiles WHERE id=$1`, id), &p.APIKey, &p) if err == sql.ErrNoRows { return nil, nil } return &p, err } +// PoolProfiles returns the failover chain in run order, api keys included: the +// active profile first, then every other keyed profile that isn't excluded, by +// priority DESC (id ASC to stay stable). Profiles without an api key can't serve +// a request, so they never enter the chain. The ordering IS the policy — callers +// walk the slice front to back. +func (d *DB) PoolProfiles() ([]*LLMProfile, error) { + rows, err := d.Query(`SELECT ` + profileColsKey + ` FROM llm_profiles +WHERE COALESCE(api_key,'') <> '' AND (is_default OR NOT pool_exclude) +ORDER BY is_default DESC, priority DESC, id ASC`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*LLMProfile + for rows.Next() { + var p LLMProfile + if err := scanProfile(rows, &p.APIKey, &p); err != nil { + return nil, err + } + out = append(out, &p) + } + return out, rows.Err() +} + // SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing. func (d *DB) SaveProfile(p *LLMProfile) (int64, error) { hint := p.APIKeyHint @@ -77,18 +118,18 @@ func (d *DB) SaveProfile(p *LLMProfile) (int64, error) { } if p.ID == 0 { var id int64 - err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11) RETURNING id`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort).Scan(&id) + err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort,priority,pool_exclude) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13) RETURNING id`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude).Scan(&id) return id, err } if p.APIKey == "" { - _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9 WHERE id=$10`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID) + _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9,priority=$10,pool_exclude=$11 WHERE id=$12`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ID) return p.ID, err } - _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11 WHERE id=$12`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID) + _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11,priority=$12,pool_exclude=$13 WHERE id=$14`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ID) return p.ID, err } diff --git a/db/config_test.go b/db/config_test.go index cd4c642..6f1743e 100644 --- a/db/config_test.go +++ b/db/config_test.go @@ -5,6 +5,54 @@ import ( "testing" ) +// TestPoolProfilesOrder pins the failover chain query: keyless profiles can't +// serve a request and excluded ones aren't fallback targets, so neither belongs +// in the chain; the rest come back by priority, highest first. +// Deliberately does NOT touch is_default — flipping the active profile would be a +// side effect on the shared dev database. +func TestPoolProfilesOrder(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + mk := func(name string, priority int, exclude bool, key string) int64 { + id, err := d.SaveProfile(&LLMProfile{ + Name: name, Format: "openai", Model: "m", APIKey: key, + Priority: priority, PoolExclude: exclude, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) }) + return id + } + lo := mk("t-pool-lo", 1, false, "k1") + hi := mk("t-pool-hi", 9, false, "k2") + mk("t-pool-excluded", 99, true, "k3") // excluded despite the top priority + mk("t-pool-nokey", 50, false, "") // no key → cannot serve anything + + chain, err := d.PoolProfiles() + if err != nil { + t.Fatal(err) + } + var got []int64 + for _, p := range chain { + switch p.Name { + case "t-pool-lo", "t-pool-hi": + got = append(got, p.ID) + case "t-pool-excluded": + t.Fatal("pool_exclude profile entered the failover chain") + case "t-pool-nokey": + t.Fatal("keyless profile entered the failover chain") + } + } + if len(got) != 2 || got[0] != hi || got[1] != lo { + t.Fatalf("chain order = %v, want [hi=%d lo=%d]", got, hi, lo) + } +} + func TestConfigStores(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { diff --git a/db/llmhealth.go b/db/llmhealth.go new file mode 100644 index 0000000..21f9f02 --- /dev/null +++ b/db/llmhealth.go @@ -0,0 +1,57 @@ +package db + +import "time" + +// LLM failover health: one row per profile tracking its circuit-breaker state. +// The authoritative copy lives in memory (llmpool.Registry); this table only +// survives a restart so a cooling-off window isn't silently reset by one. + +// LLMHealth is one profile's circuit-breaker state. +type LLMHealth struct { + ProfileID int64 `json:"profile_id"` + Fails int `json:"fails"` // consecutive failures; cleared on success + Trips int `json:"trips"` // total trips, drives the backoff ladder + OpenUntil *time.Time `json:"open_until"` // nil/past = closed (healthy) + LastError string `json:"last_error"` + LastAt time.Time `json:"last_at"` +} + +// LoadLLMHealth returns the profiles still in an UNEXPIRED cooling-off window. +// Expired rows are deliberately skipped: after a restart a profile that has +// finished cooling should be treated as healthy again and re-probed on its next +// call, not resurrected as broken. +func (d *DB) LoadLLMHealth() ([]LLMHealth, error) { + rows, err := d.Query(`SELECT profile_id,fails,trips,open_until,COALESCE(last_error,''),last_at +FROM llm_profile_health WHERE open_until IS NOT NULL AND open_until > now()`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []LLMHealth + for rows.Next() { + var h LLMHealth + if err := rows.Scan(&h.ProfileID, &h.Fails, &h.Trips, &h.OpenUntil, &h.LastError, &h.LastAt); err != nil { + return nil, err + } + out = append(out, h) + } + return out, rows.Err() +} + +// SaveLLMHealth upserts one profile's circuit-breaker state. +func (d *DB) SaveLLMHealth(h LLMHealth) error { + _, err := d.Exec(` +INSERT INTO llm_profile_health(profile_id,fails,trips,open_until,last_error,last_at) +VALUES ($1,$2,$3,$4,$5,now()) +ON CONFLICT (profile_id) DO UPDATE SET + fails=EXCLUDED.fails, trips=EXCLUDED.trips, open_until=EXCLUDED.open_until, + last_error=EXCLUDED.last_error, last_at=now()`, + h.ProfileID, h.Fails, h.Trips, h.OpenUntil, h.LastError) + return err +} + +// ClearLLMHealth drops one profile's state (manual "recover now" from the UI). +func (d *DB) ClearLLMHealth(profileID int64) error { + _, err := d.Exec(`DELETE FROM llm_profile_health WHERE profile_id=$1`, profileID) + return err +} diff --git a/db/schema.sql b/db/schema.sql index 053cfe5..32ba6f9 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -214,6 +214,11 @@ CREATE TABLE IF NOT EXISTS llm_profiles ( context_window_k INTEGER NOT NULL DEFAULT 0, reasoning_effort TEXT NOT NULL DEFAULT '', is_default BOOLEAN NOT NULL DEFAULT false, + -- 轮询(故障转移)参数,见 docs/LLM轮询设计.md: + -- priority 顺位,越大越先被选中;激活配置(is_default)永远排链首,与本值无关。 + -- pool_exclude true=不作为故障转移目标(仍可被 agent/任务显式绑定使用)。 + priority INTEGER NOT NULL DEFAULT 0, + pool_exclude BOOLEAN NOT NULL DEFAULT false, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); @@ -221,6 +226,21 @@ CREATE UNIQUE INDEX IF NOT EXISTS uq_llm_one_default ON llm_profiles(is_default) DROP TRIGGER IF EXISTS trg_llm_upd ON llm_profiles; CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles FOR EACH ROW EXECUTE FUNCTION set_updated_at(); +-- 轮询顺位/排除标记;补旧库。默认 0 / false = 全部配置都参与轮询。 +ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS priority INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS pool_exclude BOOLEAN NOT NULL DEFAULT false; + +-- LLM 轮询熔断状态:某个配置连续失败(余额不足/key 失效/限流)后进入冷却,冷却期内 +-- 轮询直接跳过它。内存态为准,这里落库只为重启后不丢冷却窗口——加载时只取尚未 +-- 到期的行(open_until > now),已到期的自然回到"正常",等下一次调用半开试探。 +CREATE TABLE IF NOT EXISTS llm_profile_health ( + profile_id BIGINT PRIMARY KEY REFERENCES llm_profiles(id) ON DELETE CASCADE, + fails INTEGER NOT NULL DEFAULT 0, -- 当前连续失败次数(成功即清零) + trips INTEGER NOT NULL DEFAULT 0, -- 累计熔断次数,用于冷却时间指数退避 + open_until TIMESTAMPTZ, -- 冷却截止;NULL/过期 = 未熔断 + last_error TEXT NOT NULL DEFAULT '', + last_at TIMESTAMPTZ NOT NULL DEFAULT now() +); -- ===================================================================== -- D. 任务层 diff --git a/llmpool/health.go b/llmpool/health.go new file mode 100644 index 0000000..80ce1f4 --- /dev/null +++ b/llmpool/health.go @@ -0,0 +1,149 @@ +// Package llmpool implements LLM failover ("轮询"): a Provider decorator that +// walks an ordered chain of LLM profiles and moves to the next one when the +// current one can't serve the request — out of credit, revoked key, rate-limited +// past the SDK's own retries, or down. +// +// The chain is built by the server from llm_profiles (active profile first, then +// by priority), so the whole process shares ONE chain and ONE circuit-breaker +// registry: when a task discovers a profile is out of credit, every other task +// skips it immediately. +package llmpool + +import ( + "sync" + "time" +) + +// backoff is the cooling-off ladder, indexed by trip count: the 1st trip cools +// for a minute, the 2nd for five, everything after that for half an hour. A +// profile that is merely rate-limited recovers quickly; one that keeps failing +// stops being probed every round. +var backoff = []time.Duration{time.Minute, 5 * time.Minute, 30 * time.Minute} + +// State is one profile's circuit-breaker state. +type State struct { + Fails int // consecutive failures; reset by a success + Trips int // total trips, indexes the backoff ladder + OpenUntil time.Time // zero / past = closed (usable) + LastError string + LastAt time.Time +} + +// Open reports whether the breaker is currently open (profile should be skipped). +func (s State) Open() bool { return time.Now().Before(s.OpenUntil) } + +// softTripAfter is how many consecutive TRANSIENT failures (429 / 5xx / network) +// trip the breaker. Deterministic failures (no credit, bad key) trip on the first. +const softTripAfter = 3 + +// Registry holds the circuit-breaker state of every profile, keyed by profile id. +// It is process-wide and outlives any single Pool instance, so rebuilding the pool +// (saving an unrelated profile, toggling a setting) never clears what we've learned +// about which backends are broken. +type Registry struct { + mu sync.Mutex + m map[int64]*State + + // persist / forget mirror state to the DB so a cooling-off window survives a + // restart. Both may be nil (no DB); both are called off the hot path. + persist func(id int64, st State) + forget func(id int64) +} + +// NewRegistry builds an empty registry. persist/forget may be nil. +func NewRegistry(persist func(id int64, st State), forget func(id int64)) *Registry { + return &Registry{m: map[int64]*State{}, persist: persist, forget: forget} +} + +// Restore seeds state loaded from the DB at startup (bypasses persistence). +func (r *Registry) Restore(id int64, st State) { + r.mu.Lock() + defer r.mu.Unlock() + cp := st + r.m[id] = &cp +} + +// Get returns a copy of one profile's state. +func (r *Registry) Get(id int64) State { + r.mu.Lock() + defer r.mu.Unlock() + if st := r.m[id]; st != nil { + return *st + } + return State{} +} + +// IsOpen reports whether a profile is in its cooling-off window. +func (r *Registry) IsOpen(id int64) bool { return r.Get(id).Open() } + +// Pass records a successful call: clears the failure counters so a profile that +// recovers is fully trusted again (and drops the persisted row). +func (r *Registry) Pass(id int64) { + r.mu.Lock() + st := r.m[id] + if st == nil || (st.Fails == 0 && st.Trips == 0 && st.OpenUntil.IsZero()) { + r.mu.Unlock() + return // already clean — nothing to write + } + delete(r.m, id) + forget := r.forget + r.mu.Unlock() + if forget != nil { + forget(id) + } +} + +// Trip records a failed call. hard=true marks a deterministic failure (no credit, +// invalid key, missing model) which opens the breaker immediately; hard=false is a +// transient one (429 / 5xx / network) that needs softTripAfter in a row. Returns +// true when this call is what opened the breaker, so the caller can log it once. +func (r *Registry) Trip(id int64, errMsg string, hard bool) (tripped bool) { + r.mu.Lock() + st := r.m[id] + if st == nil { + st = &State{} + r.m[id] = st + } + st.Fails++ + st.LastError = errMsg + st.LastAt = time.Now() + if hard || st.Fails >= softTripAfter { + st.Trips++ + d := backoff[len(backoff)-1] + if st.Trips-1 < len(backoff) { + d = backoff[st.Trips-1] + } + st.OpenUntil = time.Now().Add(d) + st.Fails = 0 // counted into this trip; start fresh for the half-open probe + tripped = true + } + snap := *st + persist := r.persist + r.mu.Unlock() + if persist != nil { + persist(id, snap) + } + return tripped +} + +// Reset clears one profile's state — the UI's "立即恢复" action. +func (r *Registry) Reset(id int64) { + r.mu.Lock() + delete(r.m, id) + forget := r.forget + r.mu.Unlock() + if forget != nil { + forget(id) + } +} + +// Snapshot returns a copy of all tracked state, for the status API. +func (r *Registry) Snapshot() map[int64]State { + r.mu.Lock() + defer r.mu.Unlock() + out := make(map[int64]State, len(r.m)) + for id, st := range r.m { + out[id] = *st + } + return out +} diff --git a/llmpool/pool.go b/llmpool/pool.go new file mode 100644 index 0000000..30b3665 --- /dev/null +++ b/llmpool/pool.go @@ -0,0 +1,261 @@ +package llmpool + +import ( + "context" + "errors" + "fmt" + "iter" + "log" + "regexp" + "strconv" + "strings" + "sync/atomic" + + "github.com/Autumn-27/norma/llm" +) + +// Member is one LLM profile in the chain, already built into a provider (wrapped +// with the recorder by the caller, so a failed attempt is still recorded under +// its own profile name). +type Member struct { + ID int64 // llm_profiles.id + Name string // profile name, for logs / UI + Model string + Format string // "anthropic" | "openai" + Priority int // the profile's configured priority (display only) + Active bool // is_default (display only) + // Rank is the ordering key the caller assigned: higher goes first, and members + // sharing a Rank take turns leading (load-spreading across duplicate keys). + // The caller encodes "active profile heads the chain" as a Rank above every + // user-settable priority, so this type needs no policy of its own. + Rank int + // WindowTokens is the profile's context window in tokens. A member whose + // window can't hold the request is skipped rather than made to fail on it. + WindowTokens int + Prov llm.Provider +} + +// RankActive is the Rank the caller gives the chain head (the active profile, or +// an explicitly bound one) so it always outranks any configured priority. +const RankActive = int(^uint(0)>>1) - 1 + +// ErrExhausted is returned when every member of the chain failed. +var ErrExhausted = errors.New("LLM 轮询:所有配置均不可用") + +// Pool is an llm.Provider that fails over across an ordered chain of members. +// It is safe for concurrent use: members are immutable after construction and +// all mutable state lives in the shared Registry. +type Pool struct { + members []*Member // in chain order (active first, then priority DESC) + health *Registry + rr atomic.Uint64 // rotates the starting point within an equal-priority group +} + +// New builds a Pool over members (already in chain order). Returns nil when the +// chain is empty. A single-member chain is still a valid Pool — it just behaves +// exactly like the bare provider. +func New(members []*Member, health *Registry) *Pool { + if len(members) == 0 { + return nil + } + if health == nil { + health = NewRegistry(nil, nil) + } + return &Pool{members: members, health: health} +} + +// Members returns the chain in order (read-only). +func (p *Pool) Members() []*Member { return p.members } + +// Head returns the first member of the chain. +func (p *Pool) Head() *Member { return p.members[0] } + +// Stream implements llm.Provider with failover. +// +// The one hard rule: a member may only be abandoned BEFORE it has yielded any +// event. Once text or a tool_use has reached the caller, re-sending the same +// request to another model would duplicate output and corrupt the conversation +// history — so a mid-stream failure is surfaced as-is and left to the agent +// harness's resume logic. Fortunately the failures this exists for (402 no +// credit, 401 bad key, 429, 5xx) all surface during request establishment, +// before the body is read, so they always land in the safe window. +func (p *Pool) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + order := p.order(req) + return func(yield func(llm.StreamEvent, error) bool) { + var lastErr error + for i, m := range order { + emitted := false + var failed error + for ev, err := range m.Prov.Stream(ctx, req) { + if err != nil && !emitted && shouldFailover(ctx, err) { + failed = err + break // safe window: nothing reached the caller yet + } + emitted = true + if !yield(ev, err) { + return // caller stopped consuming (cancel / early exit) + } + if err != nil { + return // terminal error already handed to the caller + } + } + if failed == nil { + p.health.Pass(m.ID) // completed (or failed in a non-failover way) + return + } + lastErr = failed + hard := isHardFailure(failed) + if p.health.Trip(m.ID, trimErr(failed), hard) { + log.Printf("[llmpool] 配置 %q(%s) 已熔断:%s", m.Name, m.Model, trimErr(failed)) + } + if i+1 < len(order) { + n := order[i+1] + log.Printf("[llmpool] LLM 故障转移:%q(%s) → %q(%s),原因:%s", + m.Name, m.Model, n.Name, n.Model, trimErr(failed)) + } + } + if lastErr == nil { + lastErr = ErrExhausted + } + log.Printf("[llmpool] 轮询链已耗尽(%d 个配置全部失败),最后错误:%s", len(order), trimErr(lastErr)) + yield(llm.StreamEvent{}, fmt.Errorf("%w:%v", ErrExhausted, lastErr)) + } +} + +// order picks the members to try, in order: skip those in a cooling-off window +// and those whose context window can't hold this request, then rotate within each +// equal-priority group so same-priority profiles share the load. Never returns an +// empty slice — if everything is filtered out, the head of the chain is tried +// anyway, since stalling the engine is worse than one more failed request. +func (p *Pool) order(req llm.CompletionRequest) []*Member { + est := estimateTokens(req) + var open []*Member + for _, m := range p.members { + if p.health.IsOpen(m.ID) { + continue + } + if m.WindowTokens > 0 && est > m.WindowTokens { + continue // would 400 on length — not a useful failover target + } + open = append(open, m) + } + if len(open) == 0 { + return p.members[:1] // last resort: probe the head rather than stall + } + return rotateGroups(open, p.rr.Add(1)-1) +} + +// rotateGroups rotates each run of equal-ranked members by n, so profiles sharing +// a priority take turns going first (free load-spreading across duplicate keys). +// The chain head has its own Rank and always stays at the front. +func rotateGroups(in []*Member, n uint64) []*Member { + out := make([]*Member, 0, len(in)) + for i := 0; i < len(in); { + j := i + 1 + for j < len(in) && in[j].Rank == in[i].Rank { + j++ + } + g := in[i:j] + if len(g) > 1 { + off := int(n % uint64(len(g))) + for k := range g { + out = append(out, g[(k+off)%len(g)]) + } + } else { + out = append(out, g...) + } + i = j + } + return out +} + +// statusRe pulls the HTTP status out of the SDK's error text, which is formatted +// as ": status : " (norma/llm/retry.go). The SDK exposes no +// typed error, so the string is what we have. +var statusRe = regexp.MustCompile(`status (\d{3})`) + +// statusOf returns the HTTP status carried by err, or 0 if it isn't one. +func statusOf(err error) int { + m := statusRe.FindStringSubmatch(err.Error()) + if m == nil { + return 0 + } + code, _ := strconv.Atoi(m[1]) + return code +} + +// shouldFailover reports whether err justifies trying the next profile. +// +// Never fails over on context cancellation — that's the user stopping a task or a +// task-level timeout, and burning a backup key on it would both waste credit and +// pollute the run's termination diagnosis. Never on 400 either: a malformed or +// over-long request fails identically everywhere. +func shouldFailover(ctx context.Context, err error) bool { + if err == nil { + return false + } + if ctx.Err() != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + switch code := statusOf(err); { + case code == 0: + return true // no status → transport-level failure (reset / DNS / timeout) + case code == 400: + return false // bad or over-long request: identical everywhere + case code == 401, code == 402, code == 403, code == 404, code == 408, code == 429: + return true + case code >= 500: + return true + } + return false +} + +// isHardFailure reports whether the failure is deterministic (the profile will +// keep failing until a human fixes it) rather than transient. Hard failures open +// the breaker on the first occurrence. +func isHardFailure(err error) bool { + switch statusOf(err) { + case 401, 402, 403, 404: + return true + } + return false +} + +// trimErr shortens an error for logs/UI — provider bodies can be long. +func trimErr(err error) string { + s := strings.TrimSpace(strings.ReplaceAll(err.Error(), "\n", " ")) + if len(s) > 300 { + s = s[:300] + "…" + } + return s +} + +// estimateTokens roughly sizes a request so members whose context window clearly +// can't hold it are skipped. Deliberately crude (~3.5 chars/token) and biased to +// over-estimate slightly; it only needs to tell "fits" from "nowhere near". +func estimateTokens(req llm.CompletionRequest) int { + n := 0 + for _, s := range req.System { + n += len(s) + } + for _, m := range req.Messages { + n += blocksLen(m.Content) + } + for _, t := range req.Tools { + // InputSchema is a decoded map, so its serialized size isn't available + // cheaply — charge a flat ~120 chars per top-level property instead. + n += len(t.Name) + len(t.Description) + len(t.InputSchema)*120 + } + return n * 2 / 7 // ≈ len/3.5 +} + +func blocksLen(bs []llm.ContentBlock) int { + n := 0 + for _, b := range bs { + n += len(b.Text) + len(b.Thinking) + len(b.Input) + len(b.Name) + if len(b.Content) > 0 { + n += blocksLen(b.Content) + } + } + return n +} diff --git a/llmpool/pool_test.go b/llmpool/pool_test.go new file mode 100644 index 0000000..b1e69c8 --- /dev/null +++ b/llmpool/pool_test.go @@ -0,0 +1,359 @@ +package llmpool + +import ( + "context" + "errors" + "fmt" + "iter" + "strings" + "testing" + "time" + + "github.com/Autumn-27/norma/llm" +) + +// fakeProv is a scripted provider: script[i] is what the i-th call yields — +// some events, then optionally an error. +type fakeProv struct { + name string + calls int + events [][]llm.StreamEvent // events emitted before the error, per call + errs []error // error to end each call with (nil = clean finish) +} + +// at returns script entry n, repeating the last one once the script runs out, so +// a provider defined as "always succeeds" / "always 402" keeps behaving that way +// across repeated calls. +func at[T any](s []T, n int) (T, bool) { + var zero T + if len(s) == 0 { + return zero, false + } + if n >= len(s) { + n = len(s) - 1 + } + return s[n], true +} + +func (f *fakeProv) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + n := f.calls + f.calls++ + return func(yield func(llm.StreamEvent, error) bool) { + if evs, ok := at(f.events, n); ok { + for _, e := range evs { + if !yield(e, nil) { + return + } + } + } + err, _ := at(f.errs, n) + if err != nil { + yield(llm.StreamEvent{}, err) + return + } + yield(llm.StreamEvent{Type: llm.SEMessageStop}, nil) + } +} + +// ok builds a provider that always succeeds with one text delta. +func okProv(name, text string) *fakeProv { + return &fakeProv{name: name, events: [][]llm.StreamEvent{{{Type: llm.SETextDelta, Text: text}}}} +} + +// failProv builds a provider that fails immediately (before any event) with status. +func failProv(name string, status int) *fakeProv { + return &fakeProv{name: name, errs: []error{fmt.Errorf("anthropic: status %d: nope", status)}} +} + +func member(id int64, name string, rank int, p llm.Provider) *Member { + return &Member{ID: id, Name: name, Model: "m" + name, Rank: rank, Prov: p} +} + +// drain consumes a stream, returning the concatenated text and terminal error. +func drain(seq iter.Seq2[llm.StreamEvent, error]) (string, error) { + var sb strings.Builder + for ev, err := range seq { + if err != nil { + return sb.String(), err + } + if ev.Type == llm.SETextDelta { + sb.WriteString(ev.Text) + } + } + return sb.String(), nil +} + +func TestFailoverOnNoCredit(t *testing.T) { + a, b := failProv("a", 402), okProv("b", "hello") + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil)) + + got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if err != nil { + t.Fatalf("expected failover to succeed, got %v", err) + } + if got != "hello" { + t.Fatalf("text = %q, want %q", got, "hello") + } + if a.calls != 1 || b.calls != 1 { + t.Fatalf("calls: a=%d b=%d, want 1/1", a.calls, b.calls) + } +} + +// A 402 is deterministic: one failure must open the breaker, so the NEXT request +// skips that profile entirely instead of paying for another round-trip. +func TestHardFailureTripsBreakerImmediately(t *testing.T) { + a, b := failProv("a", 402), okProv("b", "x") + reg := NewRegistry(nil, nil) + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg) + + _, _ = drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if !reg.IsOpen(1) { + t.Fatal("402 should have tripped the breaker on the first failure") + } + _, _ = drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if a.calls != 1 { + t.Fatalf("tripped profile was called again: calls=%d, want 1", a.calls) + } + if b.calls != 2 { + t.Fatalf("fallback calls=%d, want 2", b.calls) + } +} + +// 429 is transient: the SDK already retried, but we shouldn't write a profile off +// until it fails repeatedly. +func TestSoftFailureNeedsRepeats(t *testing.T) { + reg := NewRegistry(nil, nil) + for i := 1; i < softTripAfter; i++ { + if reg.Trip(1, "429", false) { + t.Fatalf("tripped after %d transient failures, want %d", i, softTripAfter) + } + } + if !reg.Trip(1, "429", false) { + t.Fatalf("should trip on failure #%d", softTripAfter) + } + if !reg.IsOpen(1) { + t.Fatal("breaker should be open") + } +} + +// A success must fully clear the counters, so an intermittent profile never +// accumulates its way to a trip. +func TestPassResetsCounters(t *testing.T) { + reg := NewRegistry(nil, nil) + reg.Trip(1, "429", false) + reg.Trip(1, "429", false) + reg.Pass(1) + if got := reg.Get(1).Fails; got != 0 { + t.Fatalf("fails=%d after Pass, want 0", got) + } + if reg.Trip(1, "429", false) { + t.Fatal("tripped immediately after a success — counters were not reset") + } +} + +func TestBackoffLadderGrows(t *testing.T) { + reg := NewRegistry(nil, nil) + var prev time.Duration + for i := range 4 { + reg.Reset(1) + st := State{Trips: i} + reg.Restore(1, st) + reg.Trip(1, "402", true) + d := time.Until(reg.Get(1).OpenUntil) + // Non-decreasing, with a second of slack: the ladder plateaus at its last + // rung, and each Trip stamps its own time.Now(). + if i > 0 && d < prev-time.Second { + t.Fatalf("trip #%d cools for %v, shorter than the previous %v", i+1, d, prev) + } + prev = d + } + if prev < 25*time.Minute { + t.Fatalf("ladder tops out at %v, want ~30m", prev) + } +} + +// The safety rule: once output has reached the caller, a mid-stream failure must +// NOT be retried on another model — that would duplicate the assistant turn. +func TestNoFailoverAfterEmit(t *testing.T) { + a := &fakeProv{ + name: "a", + events: [][]llm.StreamEvent{{{Type: llm.SETextDelta, Text: "partial"}}}, + errs: []error{fmt.Errorf("anthropic: status 500: mid-stream drop")}, + } + b := okProv("b", "full") + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil)) + + got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if err == nil { + t.Fatal("mid-stream error should surface, not be swallowed by a failover") + } + if got != "partial" { + t.Fatalf("text = %q, want the partial output %q", got, "partial") + } + if b.calls != 0 { + t.Fatalf("fell over to the backup after emitting output (calls=%d) — would duplicate the turn", b.calls) + } +} + +// Cancelling a task must not burn a backup key, and must not be diagnosed as an +// LLM fault. +func TestNoFailoverOnCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + a := &fakeProv{name: "a", errs: []error{context.Canceled}} + b := okProv("b", "x") + reg := NewRegistry(nil, nil) + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg) + + _, err := drain(p.Stream(ctx, llm.CompletionRequest{})) + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want context.Canceled", err) + } + if b.calls != 0 { + t.Fatalf("failed over on cancellation (backup calls=%d)", b.calls) + } + if reg.Get(1).Fails != 0 { + t.Fatal("cancellation counted as a profile failure") + } +} + +func TestShouldFailoverByStatus(t *testing.T) { + cases := []struct { + status int + want bool + }{ + {400, false}, // bad/over-long request fails identically everywhere + {401, true}, {402, true}, {403, true}, {404, true}, + {408, true}, {429, true}, {500, true}, {503, true}, + {200, false}, + } + for _, c := range cases { + err := fmt.Errorf("openai: status %d: body", c.status) + if got := shouldFailover(context.Background(), err); got != c.want { + t.Errorf("status %d: shouldFailover=%v, want %v", c.status, got, c.want) + } + } + // A transport error carries no status and must fail over. + if !shouldFailover(context.Background(), errors.New("dial tcp: connection reset by peer")) { + t.Error("network error should fail over") + } +} + +func TestHardVsSoftClassification(t *testing.T) { + for _, s := range []int{401, 402, 403, 404} { + if !isHardFailure(fmt.Errorf("status %d: x", s)) { + t.Errorf("status %d should be a hard failure", s) + } + } + for _, s := range []int{408, 429, 500, 502, 503} { + if isHardFailure(fmt.Errorf("status %d: x", s)) { + t.Errorf("status %d should be transient, not hard", s) + } + } +} + +// A member whose context window can't hold the request is a guaranteed 400 — +// skip it rather than spend a round-trip proving it. +func TestSkipsMembersTooSmallForRequest(t *testing.T) { + small, big := okProv("small", "s"), okProv("big", "b") + ms := member(1, "small", 10, small) + ms.WindowTokens = 1000 + mb := member(2, "big", 5, big) + mb.WindowTokens = 1_000_000 + p := New([]*Member{ms, mb}, NewRegistry(nil, nil)) + + req := llm.CompletionRequest{Messages: []llm.Message{llm.UserText(strings.Repeat("x", 100_000))}} + got, err := drain(p.Stream(context.Background(), req)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != "b" { + t.Fatalf("text=%q, want the large-window member to serve it", got) + } + if small.calls != 0 { + t.Fatalf("sent a request to a member that can't hold it (calls=%d)", small.calls) + } +} + +// Everything tripped: probing the head beats stalling the engine outright. +func TestAllTrippedStillProbesHead(t *testing.T) { + a, b := okProv("a", "a"), okProv("b", "b") + reg := NewRegistry(nil, nil) + reg.Trip(1, "402", true) + reg.Trip(2, "402", true) + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg) + + got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != "a" || a.calls != 1 { + t.Fatalf("text=%q a.calls=%d — want the head probed as a last resort", got, a.calls) + } +} + +func TestExhaustedChainReportsClearly(t *testing.T) { + a, b := failProv("a", 402), failProv("b", 402) + p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil)) + + _, err := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if !errors.Is(err, ErrExhausted) { + t.Fatalf("err = %v, want it to wrap ErrExhausted", err) + } +} + +// Equal-rank members take turns leading, so duplicate keys share the load. +func TestEqualRankRotates(t *testing.T) { + a, b := okProv("a", "a"), okProv("b", "b") + head := member(1, "head", RankActive, failProv("head", 402)) + p := New([]*Member{head, member(2, "a", 5, a), member(3, "b", 5, b)}, NewRegistry(nil, nil)) + + first := map[string]int{} + for range 4 { + txt, _ := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + first[txt]++ + } + if first["a"] == 0 || first["b"] == 0 { + t.Fatalf("same-rank members did not rotate: %v", first) + } +} + +// The head keeps its own rank and never joins a rotation group. +func TestHeadAlwaysFirst(t *testing.T) { + head := okProv("head", "H") + p := New([]*Member{ + member(1, "head", RankActive, head), + member(2, "a", 5, okProv("a", "a")), + member(3, "b", 5, okProv("b", "b")), + }, NewRegistry(nil, nil)) + for range 5 { + if txt, _ := drain(p.Stream(context.Background(), llm.CompletionRequest{})); txt != "H" { + t.Fatalf("head was not tried first, got %q", txt) + } + } +} + +// A single-member chain must behave exactly like the bare provider. +func TestSingleMemberPassthrough(t *testing.T) { + a := okProv("a", "solo") + p := New([]*Member{member(1, "a", RankActive, a)}, NewRegistry(nil, nil)) + got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{})) + if err != nil || got != "solo" { + t.Fatalf("got %q / %v, want solo / nil", got, err) + } +} + +func TestNewEmptyChainIsNil(t *testing.T) { + if New(nil, nil) != nil { + t.Fatal("empty chain should yield a nil pool so callers use the bare provider") + } +} + +// An expired cooling-off window must not keep a profile out of the chain. +func TestExpiredWindowIsClosed(t *testing.T) { + reg := NewRegistry(nil, nil) + reg.Restore(1, State{Trips: 1, OpenUntil: time.Now().Add(-time.Second)}) + if reg.IsOpen(1) { + t.Fatal("expired window should read as closed") + } +} diff --git a/server/dto.go b/server/dto.go index 56b056a..1f7f0cd 100644 --- a/server/dto.go +++ b/server/dto.go @@ -436,6 +436,10 @@ type LLMProfileDTO struct { ContextWindowK int `json:"context_window_k"` ReasoningEffort string `json:"reasoning_effort"` IsDefault bool `json:"is_default"` + // 轮询(故障转移)参数:priority 越大越先被选中(激活配置恒为链首); + // pool_exclude=true 则不作为故障转移目标,但仍可被 agent/任务显式绑定。 + Priority int `json:"priority"` + PoolExclude bool `json:"pool_exclude"` } func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO { @@ -452,6 +456,8 @@ func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO { ContextWindowK: p.ContextWindowK, ReasoningEffort: p.ReasoningEffort, IsDefault: p.IsDefault, + Priority: p.Priority, + PoolExclude: p.PoolExclude, } } diff --git a/server/goals.go b/server/goals.go index 2a72c60..41c4e99 100644 --- a/server/goals.go +++ b/server/goals.go @@ -10,6 +10,7 @@ import ( "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/norma/llm" ) type goalSpec struct { @@ -188,7 +189,7 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity // so FromEnv returned empty and every task silently fell back to the crude rule // splitter (which shredded URLs / made meaningless 2-way splits). var specs []goalSpec - if cfg, ok := s.goalsLLMConfig(t); ok { + if prov, ok := s.goalsProvider(t); ok { var as *db.AssetStore if s.m != nil { as = s.m.Assets() @@ -197,7 +198,7 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity // DecomposeGoals persists each decomposed goal via the set_goals tool straight // into t.Store; it returns the specs it wrote so we can emit activity + spot the // empty (LLM error / no provider) case below. No write-back needed here anymore. - for _, g := range agent.DecomposeGoals(ctx, cfg, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) { + for _, g := range agent.DecomposeGoals(ctx, prov, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) { if strings.TrimSpace(g.Text) != "" { specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) } @@ -220,18 +221,22 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity return specs } -// goalsLLMConfig resolves the LLM config for goal decomposition by the standard -// precedence: the goals agent's own binding → the task's pinned profile → the global -// active profile (same source the engine runs on). -func (s *Server) goalsLLMConfig(t *Task) (agent.Config, bool) { +// goalsProvider resolves the provider for goal decomposition by the standard +// precedence: the goals agent's own binding → the task's pinned profile → the +// global active provider (the very instance the engine runs on, so decomposition +// shares its rate limiter, gets recorded, and follows LLM failover). +func (s *Server) goalsProvider(t *Task) (llm.Provider, bool) { var pin *int64 if t != nil { pin = t.LLMProfileID } if eff := s.effectiveProfileForAgent("goals", pin); eff != nil { - if cfg, ok := s.loadProfileConfig(*eff); ok { - return cfg, true + if prov, cfg, ok := s.providerForProfile(*eff); ok { + return s.poolForBinding(*eff, prov, cfg), true } } - return s.loadLLMConfig() + s.cfgMu.Lock() + prov, on := s.llmProv, s.llmOn + s.cfgMu.Unlock() + return prov, on && prov != nil } diff --git a/server/llmpool.go b/server/llmpool.go new file mode 100644 index 0000000..527e09c --- /dev/null +++ b/server/llmpool.go @@ -0,0 +1,239 @@ +package server + +import ( + "log" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/llmpool" + "github.com/Autumn-27/norma/llm" +) + +// LLM 轮询(故障转移)的服务端接线。设计见 docs/LLM轮询设计.md: +// - 全局激活配置这条路径(agent 未绑定、任务未 pin)才轮询; +// - 绑定/pin 的路径默认独占该配置,失败即失败(可由 llm_pool_bind_fallback 打开兜底); +// - 链序 = 激活配置 → 其余按 priority DESC,排除 pool_exclude 的; +// - 熔断状态进程级共享(s.llmHealth),重建 pool 不清空。 + +// newLLMHealthRegistry builds the process-wide circuit-breaker registry, mirroring +// state into PG so a cooling-off window survives a restart. Writes are async and +// best-effort — the in-memory copy is authoritative. +func newLLMHealthRegistry(pg *db.DB) *llmpool.Registry { + if pg == nil { + return llmpool.NewRegistry(nil, nil) + } + persist := func(id int64, st llmpool.State) { + h := db.LLMHealth{ProfileID: id, Fails: st.Fails, Trips: st.Trips, LastError: st.LastError} + if !st.OpenUntil.IsZero() { + t := st.OpenUntil + h.OpenUntil = &t + } + go func() { + if err := pg.SaveLLMHealth(h); err != nil { + log.Printf("[llmpool] 熔断状态落库失败: %v", err) + } + }() + } + forget := func(id int64) { + go func() { _ = pg.ClearLLMHealth(id) }() + } + reg := llmpool.NewRegistry(persist, forget) + // Restore only windows that haven't expired (LoadLLMHealth filters), so a + // profile that finished cooling while we were down comes back healthy. + if rows, err := pg.LoadLLMHealth(); err == nil { + for _, h := range rows { + st := llmpool.State{Fails: h.Fails, Trips: h.Trips, LastError: h.LastError, LastAt: h.LastAt} + if h.OpenUntil != nil { + st.OpenUntil = *h.OpenUntil + } + reg.Restore(h.ProfileID, st) + log.Printf("[llmpool] 恢复熔断状态: 配置 #%d 冷却至 %s", h.ProfileID, st.OpenUntil.Format(time.RFC3339)) + } + } + return reg +} + +// poolMember builds one chain member from a profile, reusing the per-profile +// provider cache so every agent pointing at the same profile shares one provider +// (hence one rate limiter). Returns nil when the profile can't be built. +func (s *Server) poolMember(p *db.LLMProfile, rank int) *llmpool.Member { + prov, cfg, ok := s.providerForProfile(p.ID) + if !ok { + return nil + } + return &llmpool.Member{ + ID: p.ID, Name: p.Name, Model: p.Model, Format: p.Format, + Priority: p.Priority, Active: p.IsDefault, Rank: rank, + WindowTokens: cfg.CompactionWindow(), Prov: prov, + } +} + +// poolChain reads the failover chain from the DB. headID/headProv/headCfg describe +// the member that must lead the chain — the globally active profile on the normal +// path, or an explicitly bound one when bound-profile fallback is on. The head is +// passed in (rather than resolved through the cache) so applyLLM's already-built +// provider is reused instead of duplicated. +// +// Returns nil when failover is off or the chain has fewer than two usable members +// — callers then use the bare provider, which is byte-for-byte the old behavior. +func (s *Server) poolChain(headID int64, headProv llm.Provider, headCfg agent.Config) *llmpool.Pool { + if s.m == nil || s.m.pg == nil || !s.m.LLMPoolEnabled() { + return nil + } + profs, err := s.m.pg.PoolProfiles() + if err != nil { + log.Printf("[llmpool] 读取轮询链失败: %v", err) + return nil + } + var head *db.LLMProfile + for _, p := range profs { + if p.ID == headID { + head = p + break + } + } + if head == nil { // head excluded from the chain (bound + pool_exclude) — still leads + if p, err := s.m.pg.ProfileByID(headID); err == nil && p != nil { + head = p + } else { + return nil + } + } + members := []*llmpool.Member{{ + ID: head.ID, Name: head.Name, Model: head.Model, Format: head.Format, + Priority: head.Priority, Active: head.IsDefault, Rank: llmpool.RankActive, + WindowTokens: headCfg.CompactionWindow(), Prov: headProv, + }} + for _, p := range profs { + if p.ID == headID { + continue + } + // The globally active profile outranks the others when it isn't the head + // (bound-fallback chains), so it stays the first fallback tried. + rank := p.Priority + if p.IsDefault { + rank = llmpool.RankActive - 1 + } + if m := s.poolMember(p, rank); m != nil { + members = append(members, m) + } + } + if len(members) < 2 { + return nil // nothing to fail over to + } + return llmpool.New(members, s.llmHealth) +} + +// poolForActive wraps the globally-active provider in the failover chain. Returns +// prov unchanged when failover is off or there's nothing to fall back to. +func (s *Server) poolForActive(activeID int64, prov llm.Provider, cfg agent.Config) llm.Provider { + pool := s.poolChain(activeID, prov, cfg) + if pool == nil { + return prov + } + names := make([]string, 0, len(pool.Members())) + for _, m := range pool.Members() { + names = append(names, m.Name+"/"+m.Model) + } + log.Printf("[llmpool] LLM 轮询已启用,链路(%d): %v", len(names), names) + return pool +} + +// poolForBinding wraps a BOUND profile's provider so it falls back to the chain. +// Only active when both the failover master switch and the bound-fallback switch +// are on; otherwise a bound profile stays exclusive (fails hard), which is the +// default and the documented precedence. +func (s *Server) poolForBinding(id int64, prov llm.Provider, cfg agent.Config) llm.Provider { + if s.m == nil || !s.m.LLMPoolEnabled() || !s.m.LLMPoolBindFallback() { + return prov + } + if pool := s.poolChain(id, prov, cfg); pool != nil { + return pool + } + return prov +} + +// LLMPoolMemberStatus is one chain entry as shown in the UI. +type LLMPoolMemberStatus struct { + ProfileID string `json:"profile_id"` + Name string `json:"name"` + Model string `json:"model"` + Format string `json:"format"` + Priority int `json:"priority"` + Active bool `json:"active"` // is the globally-activated profile + Excluded bool `json:"excluded"` // pool_exclude — not a failover target + // Health: state is "ok" | "degraded" (failing but not tripped) | "tripped". + State string `json:"state"` + Fails int `json:"fails"` + Trips int `json:"trips"` + CooldownSecs int `json:"cooldown_secs"` // remaining cooling-off seconds; 0 = none + LastError string `json:"last_error,omitempty"` + LastAt string `json:"last_at,omitempty"` +} + +// llmPoolStatus reports the whole picture for the LLM page: whether failover is +// on, the resolved chain order, and each profile's breaker state. Excluded +// profiles are listed too (flagged), so the user can see why one isn't in line. +func (s *Server) llmPoolStatus() map[string]any { + out := map[string]any{ + "enabled": false, + "bind_fallback": false, + "chain": []LLMPoolMemberStatus{}, + } + if s.m == nil || s.m.pg == nil { + return out + } + out["enabled"] = s.m.LLMPoolEnabled() + out["bind_fallback"] = s.m.LLMPoolBindFallback() + + all, err := s.m.pg.ListProfiles() + if err != nil { + return out + } + health := s.llmHealth.Snapshot() + now := time.Now() + // Chain order: active first, then priority DESC / id ASC — the same ordering + // PoolProfiles applies, recomputed here so excluded ones can be shown in place. + chain := make([]LLMPoolMemberStatus, 0, len(all)) + for _, p := range all { + st := health[p.ID] + m := LLMPoolMemberStatus{ + ProfileID: i64s(p.ID), Name: p.Name, Model: p.Model, Format: p.Format, + Priority: p.Priority, Active: p.IsDefault, Excluded: p.PoolExclude, + State: "ok", Fails: st.Fails, Trips: st.Trips, + LastError: st.LastError, + } + if st.Open() { + m.State = "tripped" + m.CooldownSecs = int(st.OpenUntil.Sub(now).Seconds()) + 1 + } else if st.Fails > 0 { + m.State = "degraded" + } + if !st.LastAt.IsZero() { + m.LastAt = st.LastAt.Format(time.RFC3339) + } + chain = append(chain, m) + } + sortPoolStatus(chain) + out["chain"] = chain + return out +} + +// sortPoolStatus orders the status rows exactly like the live chain: active first, +// then priority DESC, then id ASC (the incoming slice is already id-ordered, so a +// stable insertion sort on the first two keys is enough). +func sortPoolStatus(in []LLMPoolMemberStatus) { + for i := 1; i < len(in); i++ { + for j := i; j > 0 && poolLess(in[j], in[j-1]); j-- { + in[j], in[j-1] = in[j-1], in[j] + } + } +} + +func poolLess(a, b LLMPoolMemberStatus) bool { + if a.Active != b.Active { + return a.Active + } + return a.Priority > b.Priority +} diff --git a/server/llmpool_test.go b/server/llmpool_test.go new file mode 100644 index 0000000..0927b86 --- /dev/null +++ b/server/llmpool_test.go @@ -0,0 +1,41 @@ +package server + +import "testing" + +// The status list drives the "轮询顺序" strip in the UI, so it has to match the +// order PoolProfiles actually runs: active first, then priority DESC, then id ASC +// (input arrives id-ordered, so equal priorities must keep their relative order). +func TestSortPoolStatusMatchesChainOrder(t *testing.T) { + in := []LLMPoolMemberStatus{ + {ProfileID: "1", Name: "low", Priority: 0}, + {ProfileID: "2", Name: "active", Priority: 0, Active: true}, + {ProfileID: "3", Name: "high", Priority: 10}, + {ProfileID: "4", Name: "mid-a", Priority: 5}, + {ProfileID: "5", Name: "mid-b", Priority: 5}, + } + sortPoolStatus(in) + + want := []string{"active", "high", "mid-a", "mid-b", "low"} + for i, w := range want { + if in[i].Name != w { + got := make([]string, len(in)) + for j, m := range in { + got[j] = m.Name + } + t.Fatalf("order = %v, want %v", got, want) + } + } +} + +// The active profile heads the chain no matter how low its own priority is — +// that's the documented precedence, and the UI must not imply otherwise. +func TestActiveProfileHeadsStatusList(t *testing.T) { + in := []LLMPoolMemberStatus{ + {ProfileID: "1", Name: "loud", Priority: 999}, + {ProfileID: "2", Name: "active", Priority: -5, Active: true}, + } + sortPoolStatus(in) + if in[0].Name != "active" { + t.Fatalf("head = %q, want the active profile regardless of priority", in[0].Name) + } +} diff --git a/server/manager.go b/server/manager.go index 454b229..78e403d 100644 --- a/server/manager.go +++ b/server/manager.go @@ -86,6 +86,12 @@ const ( settingWebSearchProxy = "web_search_proxy" settingWorkers = "workers" settingLLMRecord = "llm_record" + // LLM 轮询(故障转移)。默认关闭——开启后走「全局激活配置」的 agent 在当前配置 + // 不可用(余额不足/key 失效/限流/服务异常)时自动切到下一个配置。 + // settingLLMPoolBindFallback 仅在轮询开启时有意义:默认关闭,即 agent/任务显式 + // 绑定了某个配置就只用它、失败即失败;开启后绑定的配置失败也会回落到轮询链。 + settingLLMPoolOn = "llm_pool_enabled" + settingLLMPoolBindFallback = "llm_pool_bind_fallback" // 任务并发上限:开关 + 上限数。默认关闭;开启后默认上限 5(见 defaultConcurrencyLimit)。 settingConcurrencyOn = "task_concurrency_enabled" settingConcurrencyLimit = "task_concurrency_limit" @@ -266,6 +272,36 @@ func (m *Manager) SetLLMRecordEnabled(on bool) error { return nil } +// LLMPoolEnabled reports whether LLM failover ("轮询") is on (默认关; +// settings.llm_pool_enabled). Read when the provider chain is built (applyLLM), +// so a change requires a rebuild — putSettings does that. +func (m *Manager) LLMPoolEnabled() bool { + if m.pg == nil { + return false + } + return m.pg.GetBool(settingLLMPoolOn, false) +} + +// SetLLMPoolEnabled persists the failover toggle. Callers rebuild agents +// (applyLLM) afterwards so it takes effect. +func (m *Manager) SetLLMPoolEnabled(on bool) error { return m.pg.SetBool(settingLLMPoolOn, on) } + +// LLMPoolBindFallback reports whether an agent/task that is BOUND to a specific +// profile still falls back to the chain when that profile fails (默认关:绑定即 +// 独占,失败即失败). Only meaningful while LLMPoolEnabled. +func (m *Manager) LLMPoolBindFallback() bool { + if m.pg == nil { + return false + } + return m.pg.GetBool(settingLLMPoolBindFallback, false) +} + +// SetLLMPoolBindFallback persists the bound-profile fallback toggle. Callers +// rebuild agents (applyLLM) afterwards. +func (m *Manager) SetLLMPoolBindFallback(on bool) error { + return m.pg.SetBool(settingLLMPoolBindFallback, on) +} + // WebSearch returns the current web-search config: whether it is enabled, the // backend ("ddgs" | "brave-free" | "tavily"), the Brave API key, the Tavily API // key (each empty unless set), and the dedicated egress proxy (empty = direct). diff --git a/server/server.go b/server/server.go index 4c68585..2fe0d1d 100644 --- a/server/server.go +++ b/server/server.go @@ -19,6 +19,7 @@ import ( "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/llmpool" "github.com/Autumn-27/artex/llmrec" "github.com/Autumn-27/artex/report" "github.com/Autumn-27/norma/llm" @@ -53,6 +54,11 @@ type Server struct { llmCfg agent.Config // current LLM config (key not exposed) llmOn bool llmProf string // active LLM profile name (for llmrec tagging) + // llmProv is the fully-decorated global provider (recorder + failover chain) + // applyLLM installed. Kept so one-off callers (goal decomposition) reuse the + // same instance instead of building a private one that skips the rate limiter, + // the recorder and failover. + llmProv llm.Provider // chatBusy guards the per-task main-agent run: the chat handler launches the // agent on the server's background ctx (not the request ctx) and returns @@ -97,6 +103,12 @@ type Server struct { // leaving it unset), just not deduped. Cleared alongside profAgents on profile edits. provCacheMu sync.Mutex provByProfile map[int64]*provEntry + + // llmHealth is the process-wide circuit-breaker state for LLM failover (轮询). + // It deliberately lives OUTSIDE the provider caches: rebuilding the chain + // (saving an unrelated profile, flipping a setting) must not erase what we + // learned about which backends are out of credit / rate-limited. + llmHealth *llmpool.Registry } // profBundle is a planner/worker pair built from one LLM profile. @@ -131,7 +143,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi chatCancel: map[string]context.CancelCauseFunc{}, triggerQ: map[string][]triggeredRun{}, triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{}, profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{}, - provByProfile: map[int64]*provEntry{}} + provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg)} // per-task LLM: a task pinned to a specific profile runs on that profile's // dedicated planner/worker; unpinned tasks fall back to the global active pair. s.engine.SetAgentResolver(func(t *Task) (*agent.Planner, *agent.Worker) { @@ -291,10 +303,15 @@ func (s *Server) saveLLMConfig(cfg agent.Config) { format = "openai" } var id int64 + // agent.Config carries no failover fields, so carry the stored ones forward — + // otherwise this legacy endpoint would silently reset the profile's priority + // and pool_exclude to zero values on every save. + var priority int + var poolExclude bool if profs, _ := s.m.pg.ListProfiles(); profs != nil { for _, p := range profs { if p.Name == "default" { - id = p.ID + id, priority, poolExclude = p.ID, p.Priority, p.PoolExclude break } } @@ -303,6 +320,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) { ID: id, Name: "default", Format: format, Model: cfg.Model, BaseURL: cfg.BaseURL, Proxy: cfg.Proxy, APIKey: cfg.APIKey, RatePerSecond: cfg.RatePerSecond, RatePerMinute: cfg.RatePerMinute, ContextWindowK: cfg.ContextWindowK, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, + Priority: priority, PoolExclude: poolExclude, }) if err == nil { _ = s.m.pg.SetActiveProfile(newID) @@ -371,6 +389,12 @@ func (s *Server) applyLLM(cfg agent.Config) error { profName := s.llmProf s.cfgMu.Unlock() prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) + // LLM 轮询(默认关):把激活配置包进故障转移链,当前配置不可用时自动切下一个。 + // 只影响「走全局激活配置」的这条路径——agent 绑定 / 任务 pin 的走 providerForProfile, + // 默认仍然独占该配置(见 poolForBinding)。关闭或无备选时返回原 provider,行为不变。 + if act, err := s.m.pg.ActiveProfile(); err == nil && act != nil { + prov = s.poolForActive(act.ID, prov, cfg) + } pl, wk := s.buildPlannerWorker(nil, prov, cfg) s.engine.UseLLM(pl, wk) @@ -391,6 +415,7 @@ func (s *Server) applyLLM(cfg agent.Config) error { s.chatAgent.SetWebSearch(s.m.WebSearchOpts()) s.chatAgent.SetGuard(s.chatGuard()) s.llmCfg = cfg + s.llmProv = prov s.llmOn = true s.cfgMu.Unlock() @@ -468,10 +493,14 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) // providerForAgent returns the provider+cfg one agent should run on: its bound/pinned // profile if that resolves, else the passed global fallback (gProv/gCfg). +// A bound/pinned profile is EXCLUSIVE by default: it does not fail over, so an +// agent deliberately put on a specific model never silently runs on another one. +// Turning on llm_pool_bind_fallback relaxes that — poolForBinding then appends the +// failover chain behind it. func (s *Server) providerForAgent(agentKey string, pinID *int64, gProv llm.Provider, gCfg agent.Config) (llm.Provider, agent.Config) { if eff := s.effectiveProfileForAgent(agentKey, pinID); eff != nil { if prov, cfg, ok := s.providerForProfile(*eff); ok { - return prov, cfg + return s.poolForBinding(*eff, prov, cfg), cfg } } return gProv, gCfg @@ -492,7 +521,9 @@ func (s *Server) agentsForProfile(id int64) (*agent.Planner, *agent.Worker) { if !ok { return nil, nil } - pl, wk := s.buildPlannerWorker(&id, prov, cfg) + // The pin is the fallback for agents without their own binding — same + // exclusive-by-default rule, so route it through poolForBinding too. + pl, wk := s.buildPlannerWorker(&id, s.poolForBinding(id, prov, cfg), cfg) s.profMu.Lock() if ex := s.profAgents[id]; ex != nil { // lost the race → keep the winner pl, wk = ex.pl, ex.wk @@ -517,7 +548,7 @@ func (s *Server) chatAgentForProfile(id int64) *agent.ChatAgent { return nil } tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) - ca := agent.NewChatAgent(prov, cfg.Model, s.m.dir, tx, cfg.CompactionWindow()) + ca := agent.NewChatAgent(s.poolForBinding(id, prov, cfg), cfg.Model, s.m.dir, tx, cfg.CompactionWindow()) ca.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) ca.SetWebSearch(s.m.WebSearchOpts()) ca.SetGuard(s.chatGuard()) @@ -726,6 +757,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /api/llm/profiles", s.pgSaveProfile) mux.HandleFunc("DELETE /api/llm/profiles/{id}", s.pgDeleteProfile) mux.HandleFunc("POST /api/llm/profiles/active", s.pgActivateProfile) + mux.HandleFunc("GET /api/llm/pool", s.pgLLMPoolStatus) + mux.HandleFunc("POST /api/llm/pool/reset", s.pgLLMPoolReset) mux.HandleFunc("POST /api/llm/models", s.pgListModels) // 拦截规则管理 @@ -2191,6 +2224,10 @@ func (s *Server) settingsPayload() map[string]any { "workers": s.m.Workers(), // 并发工作 agent 数(默认3);对之后启动的任务生效 "task_concurrency_enabled": concOn, // 任务并发上限开关(默认关) "task_concurrency_limit": concLimit, // 同时运行任务上限(开启后默认5) + // LLM 轮询(故障转移)。默认关;开启后走全局激活配置的 agent 在当前配置不可用时 + // 自动切到下一个配置。bind_fallback 仅在轮询开启时有意义(默认关)。 + "llm_pool_enabled": s.m.LLMPoolEnabled(), + "llm_pool_bind_fallback": s.m.LLMPoolBindFallback(), } } @@ -2227,6 +2264,10 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { // 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。 ConcurrencyEnabled *bool `json:"task_concurrency_enabled"` ConcurrencyLimit *int `json:"task_concurrency_limit"` + // LLM 轮询(故障转移)开关 + 「绑定配置失败也兜底回轮询链」开关。两者都需要 + // 重建 provider 链才生效,走下面的 changed → applyLLM 路径。 + LLMPoolEnabled *bool `json:"llm_pool_enabled"` + LLMPoolBindFallback *bool `json:"llm_pool_bind_fallback"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { writeErr(w, 400, err.Error()) @@ -2272,6 +2313,25 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { } } changed := false + if req.LLMPoolEnabled != nil { + if err := s.m.SetLLMPoolEnabled(*req.LLMPoolEnabled); err != nil { + writeErr(w, 500, err.Error()) + return + } + changed = true // the provider chain itself changes shape → rebuild + } + if req.LLMPoolBindFallback != nil { + if err := s.m.SetLLMPoolBindFallback(*req.LLMPoolBindFallback); err != nil { + writeErr(w, 500, err.Error()) + return + } + changed = true + } + if req.LLMPoolEnabled != nil || req.LLMPoolBindFallback != nil { + // Pinned tasks' planner/worker and per-profile chat agents hold providers + // built under the OLD switch state — drop them so they pick up the new one. + s.invalidateProfileAgents() + } if req.TrafficCapture != nil { if err := s.m.SetTrafficEnabled(*req.TrafficCapture); err != nil { writeErr(w, 500, err.Error()) diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 0595787..eef78bb 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -1609,10 +1609,10 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) { } // Editing a profile rebuilds any task pinned to it on its next round. s.invalidateProfileAgents() - // If we just edited the active profile, hot-apply so changes take effect now. - if act, _ := pg.ActiveProfile(); act != nil && act.ID == id { - s.reapplyActiveProfile() - } + // Hot-apply. Not just when the ACTIVE profile changed: with failover on, any + // profile's key/model/priority/pool_exclude edit reshapes the chain, so the + // engine's provider has to be rebuilt either way. + s.reapplyActiveProfile() writeJSON(w, 200, map[string]any{"id": id}) } @@ -1627,6 +1627,8 @@ func (s *Server) pgDeleteProfile(w http.ResponseWriter, r *http.Request) { return } s.invalidateProfileAgents() // drop cached agents for the removed profile + s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away) + s.reapplyActiveProfile() // the deleted profile may have been in the failover chain writeJSON(w, 200, map[string]any{"deleted": id}) } @@ -1651,6 +1653,40 @@ func (s *Server) pgActivateProfile(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"ok": true}) } +// pgLLMPoolStatus reports the failover ("轮询") switches, the resolved chain order +// and every profile's circuit-breaker state — what the LLM page renders as the +// "轮询顺序" strip and the per-card health badges. +func (s *Server) pgLLMPoolStatus(w http.ResponseWriter, r *http.Request) { + if s.pg(w) == nil { + return + } + writeJSON(w, 200, s.llmPoolStatus()) +} + +// pgLLMPoolReset clears a tripped profile's circuit breaker so the next call +// tries it again immediately ("立即恢复"). id=0 clears every profile. +func (s *Server) pgLLMPoolReset(w http.ResponseWriter, r *http.Request) { + pg := s.pg(w) + if pg == nil { + return + } + var body struct { + ID int64 `json:"id"` // 0 / omitted = all + } + if err := decode(r, &body); err != nil { + writeErr(w, 400, err.Error()) + return + } + if body.ID > 0 { + s.llmHealth.Reset(body.ID) + } else { + for id := range s.llmHealth.Snapshot() { + s.llmHealth.Reset(id) + } + } + writeJSON(w, 200, s.llmPoolStatus()) +} + // pgListModels fetches available models from the provider's API endpoint. // Supports OpenAI-format (GET /models) and Anthropic-format (GET /v1/models); for // Anthropic-compatible third parties (e.g. DeepSeek) whose model list lives only on diff --git a/web/src/app/(main)/system/llm/page.tsx b/web/src/app/(main)/system/llm/page.tsx index f01ac53..6de5007 100644 --- a/web/src/app/(main)/system/llm/page.tsx +++ b/web/src/app/(main)/system/llm/page.tsx @@ -2,7 +2,17 @@ import * as React from "react"; -import { Loader2Icon, PlugZapIcon, PlusIcon, RefreshCwIcon, SaveIcon, StarIcon, Trash2Icon } from "lucide-react"; +import { + Loader2Icon, + PlugZapIcon, + PlusIcon, + RefreshCwIcon, + RotateCcwIcon, + SaveIcon, + StarIcon, + Trash2Icon, + ZapIcon, +} from "lucide-react"; import { toast } from "sonner"; import { Badge } from "@/components/ui/badge"; @@ -23,8 +33,9 @@ import { Label } from "@/components/ui/label"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; +import { Switch } from "@/components/ui/switch"; import { api } from "@/lib/api"; -import type { LLMProfile } from "@/lib/types"; +import type { LLMPoolStatus, LLMProfile } from "@/lib/types"; import { cn } from "@/lib/utils"; // 思考强度档位(仅在开关打开时生效)。 @@ -48,6 +59,203 @@ const toStore = (mode: ThinkMode, effort: string) => (mode === "on" ? effort : m const modeFromStore = (v?: string): ThinkMode => (!v ? "none" : v === "off" ? "off" : "on"); const effortFromStore = (v?: string) => (v && v !== "off" ? v : "high"); +// 熔断状态 → 徽章样式。degraded 是「在失败但还没到熔断阈值」的中间态。 +const STATE_BADGE: Record = { + ok: { label: "正常", cls: "border-emerald-500/50 text-emerald-600 dark:text-emerald-400" }, + degraded: { label: "异常", cls: "border-amber-500/50 text-amber-600 dark:text-amber-400" }, + tripped: { label: "已熔断", cls: "border-destructive/50 text-destructive" }, +}; + +function cooldownText(secs: number) { + if (secs <= 0) return ""; + if (secs < 60) return `${secs}s`; + return `${Math.ceil(secs / 60)}min`; +} + +// PoolCard 是轮询功能的总控:主开关、绑定兜底开关,以及后端算出来的**实际**链路顺序。 +// 顺序预览很关键——「激活 + 优先级 + 不参与轮询」三个因素叠加后,光看单张配置卡片 +// 是猜不出真实顺序的。 +function PoolCard({ refreshKey }: { refreshKey: number }) { + const [pool, setPool] = React.useState(null); + const [busy, setBusy] = React.useState(false); + + const load = React.useCallback(async () => { + try { + setPool(await api.llmPool()); + } catch { + /* ignore */ + } + }, []); + // biome-ignore lint/correctness/useExhaustiveDependencies: refreshKey is a prop — bumping it is exactly how the parent asks for a refetch after a profile edit + React.useEffect(() => { + void load(); + }, [load, refreshKey]); + // 冷却倒计时是后端算的剩余秒数,定时拉一次让它走起来(只在有熔断时才轮询)。 + React.useEffect(() => { + if (!pool?.enabled || !pool.chain.some((m) => m.state !== "ok")) return; + const t = setInterval(() => void load(), 10_000); + return () => clearInterval(t); + }, [pool, load]); + + async function toggle(patch: { llm_pool_enabled?: boolean; llm_pool_bind_fallback?: boolean }) { + if (busy) return; + setBusy(true); + try { + await api.setSettings(patch); + await load(); + if (patch.llm_pool_enabled !== undefined) { + toast.success(patch.llm_pool_enabled ? "已开启 LLM 轮询" : "已关闭 LLM 轮询"); + } else { + toast.success("已更新绑定兜底设置"); + } + } catch (e) { + toast.error(`设置失败:${(e as Error).message}`); + } finally { + setBusy(false); + } + } + + async function recover(id?: string) { + try { + setPool(await api.resetLLMPool(id)); + toast.success(id ? "已恢复该配置" : "已恢复全部配置"); + } catch (e) { + toast.error(`恢复失败:${(e as Error).message}`); + } + } + + const enabled = pool?.enabled ?? false; + // 链路顺序 = 后端返回的顺序;不参与轮询的配置仍会列出(灰显并标注), + // 这样用户能看见「它为什么不在队列里」,而不是凭空消失。 + const chain = pool?.chain ?? []; + const inChain = chain.filter((m) => m.active || !m.excluded); + const tripped = chain.filter((m) => m.state === "tripped"); + + return ( + + + + LLM 轮询 · 故障转移 + + + 开启后,未指定模型的 Agent 在当前配置不可用(余额不足 / Key 失效 / 限流 / + 服务异常)时自动切换到下一个配置。 指定了模型的 Agent 与任务默认不参与轮询。 + + + void toggle({ llm_pool_enabled: v })} + aria-label="LLM 轮询开关" + /> + + + {enabled && ( + +
+
+ +

+ 默认关闭:Agent 或任务指定了某个配置就只用它,失败即失败(不会悄悄换成别的模型)。 + 开启后,指定的配置失败时也会回落到下面的轮询链。 +

+
+ void toggle({ llm_pool_bind_fallback: v })} + aria-label="绑定配置失败兜底开关" + /> +
+ +
+
+ + {tripped.length > 0 && ( + + )} +
+ {inChain.length < 2 && ( +

+ 当前只有 {inChain.length} 个可用配置,轮询不会生效——至少需要 2 个已填 API Key 且参与轮询的配置。 +

+ )} +
+ {chain.map((m) => { + const badge = STATE_BADGE[m.state] ?? STATE_BADGE.ok; + const order = + m.excluded && !m.active ? null : inChain.findIndex((x) => x.profile_id === m.profile_id) + 1; + return ( +
+ + {order ?? "—"} + + {m.name} + {m.model} + {m.active && ( + + 激活 + + )} + {m.excluded && !m.active && 不参与轮询} + {!m.active && 优先级 {m.priority}} +
+ {m.state === "tripped" && m.cooldown_secs > 0 && ( + 冷却 {cooldownText(m.cooldown_secs)} + )} + {m.state === "degraded" && ( + 连续失败 {m.fails} 次 + )} + + {badge.label} + + {m.state !== "ok" && ( + + )} +
+ {m.last_error && ( +

+ {m.last_error} +

+ )} +
+ ); + })} + {chain.length === 0 && ( +
+ 暂无配置 +
+ )} +
+

+ 激活配置恒为第 1 顺位,其余按优先级从高到低。某个配置失败后进入冷却(60s → 5min → 30min), + 冷却期内被跳过;恢复后自动切回。上下文窗口装不下当前请求的配置会被跳过。 +

+
+
+ )} +
+ ); +} + function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { const [open, setOpen] = React.useState(false); const [name, setName] = React.useState(""); @@ -61,6 +269,8 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { const [cw, setCw] = React.useState("0"); const [thinkMode, setThinkMode] = React.useState("none"); const [effort, setEffort] = React.useState("high"); + const [priority, setPriority] = React.useState("0"); + const [poolExclude, setPoolExclude] = React.useState(false); const [models, setModels] = React.useState([]); const [loadingModels, setLoadingModels] = React.useState(false); const [modelsOpen, setModelsOpen] = React.useState(false); @@ -77,6 +287,8 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { setCw("0"); setThinkMode("none"); setEffort("high"); + setPriority("0"); + setPoolExclude(false); setModels([]); setModelsOpen(false); } @@ -118,6 +330,8 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { rate_per_minute: Number(rpm) || 0, context_window_k: Number(cw) || 0, reasoning_effort: toStore(thinkMode, effort), + priority: Number(priority) || 0, + pool_exclude: poolExclude, }); toast.success(`已新建 Profile:${name.trim()}(在列表中「设为激活」以启用)`); reset(); @@ -275,12 +489,39 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { />
+
+
+
+ +

数字越大越先被选中;激活配置恒为第 1 顺位。

+
+ setPriority(e.target.value)} + /> +
+
+
+ +

+ 开启后不会被当作故障转移目标(仍可被 Agent / 任务显式指定使用)。 +

+
+ +
+

- 不发送=不带思考字段(兼容 MiniMax 等不支持该字段的模型);显式关闭=发 disabled(默认开思考的模型才需要);开启=先推理再作答。 + 不发送=不带思考字段(兼容 MiniMax 等不支持该字段的模型);显式关闭=发 + disabled(默认开思考的模型才需要);开启=先推理再作答。

setPriority(e.target.value)} + /> +
+
+
+ +

+ 开启后不会被当作故障转移目标(仍可被 Agent / 任务显式指定使用)。 适合「只给某个 Agent + 专用、不希望别人失败时烧掉」的昂贵配置。 +

+
+ +
+
+

- 不发送=不带思考字段(兼容 MiniMax 等不支持该字段的模型);显式关闭=发 disabled(默认开思考的模型才需要);开启=先推理再作答。 + 不发送=不带思考字段(兼容 MiniMax 等不支持该字段的模型);显式关闭=发 + disabled(默认开思考的模型才需要);开启=先推理再作答。

{uploading ? : } - setInput(e.target.value)} onKeyDown={(e) => { - if (e.key === "Enter" && !mainLive) send(); + // Enter 换行;只有 Ctrl/Cmd+Enter 才发送。 + if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) { + e.preventDefault(); + if (!mainLive) send(); + } }} disabled={mainLive} /> From c436b9d7d8d68d4f99bb20b9e0f392e8c81fc96f Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:12:26 -0400 Subject: [PATCH 11/70] =?UTF-8?q?fix(update):=20Docker=20=E6=9B=B4?= =?UTF-8?q?=E6=96=B0=E5=8F=AA=E6=8B=89=20artex=20=E9=95=9C=E5=83=8F,?= =?UTF-8?q?=E4=B8=8D=E5=8A=A8=20postgres?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit postgres 固定 16-alpine 无需跟着升级,pull/up 带上服务名即可;artex 声明了 depends_on postgres,pg 没起会自动拉起,已在跑的不重建。 --- update.sh | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/update.sh b/update.sh index 3a870c6..10cfbbc 100755 --- a/update.sh +++ b/update.sh @@ -38,10 +38,13 @@ update_docker(){ ok "已将 ARTEX_TAG 设为 ${tag}" fi - info "拉取新镜像…" - docker compose pull + # 只动 artex:postgres 是固定的 16-alpine,不需要跟着升级(拉它纯属浪费带宽, + # 且大版本变动还会有兼容风险)。artex 声明了 depends_on postgres,所以带服务名 + # up 时若 pg 没起会自动拉起,已在跑的则原样保留、不重建。 + info "拉取新镜像(仅 artex)…" + docker compose pull artex info "重建并启动(artex 重启时自动迁移 schema)…" - docker compose up -d + docker compose up -d artex ok "更新完成 → http://localhost:8787" info "查看日志:docker compose logs -f artex" info "清理旧镜像(可选):docker image prune -f" From 7090b034f8141719efa8134f7d8f0186da60f7c0 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:25:45 -0400 Subject: [PATCH 12/70] =?UTF-8?q?refactor(ui):=20LLM=20=E9=A1=B5=E6=94=B9?= =?UTF-8?q?=E4=B8=BA=E9=85=8D=E7=BD=AE=E5=8D=A1=E7=89=87=E5=88=97=E8=A1=A8?= =?UTF-8?q?=20+=20=E6=8A=BD=E5=B1=89=E7=BC=96=E8=BE=91,=E6=96=B0=E5=A2=9E?= =?UTF-8?q?=E8=BD=AE=E8=AF=A2=E9=85=8D=E7=BD=AE=E5=85=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 原来是「上边轮询卡片 + 下边左列表右详情」的分栏布局,详情区常驻占掉三分之二 版面。改为: - 页面主体 = 模型配置卡片网格,点卡片打开右侧抽屉编辑/测试/保存 - 卡片右上角显示该配置是否正常:未配置 Key / 正常 / 异常·失败 N 次 / 已熔断·冷却 xx,悬停出具体错误 - 工具栏「新建」左边加「轮询配置」入口,抽屉里放开关、兜底开关、实际链路 顺序、熔断状态与恢复操作;开启时按钮上带「已开启」徽章 新建与编辑合用同一个抽屉(profile === null 即新建)。原先新建是 Dialog、编辑 是常驻表单,两份几乎完全相同的表单代码约 150 行,加字段得改两处——现在只有 一份 ProfileSheet。 抽屉的开关状态与内容状态分开存(editOpen / editing):否则关闭动画那 200ms 里标题会从「编辑 XX」闪成「新建模型配置」。 卡片上的优先级/不参与轮询只在轮询开启时才显示——关着时这两个字段没有意义, 不占版面。 注:卡片的「正常」是排除法得出的弱信号(填了 Key 且轮询没记录到它失败过), 不是主动探测——从未被调用过的错误配置也会显示正常。 --- web/src/app/(main)/system/llm/page.tsx | 1295 ++++++++++-------------- 1 file changed, 521 insertions(+), 774 deletions(-) diff --git a/web/src/app/(main)/system/llm/page.tsx b/web/src/app/(main)/system/llm/page.tsx index 6de5007..fb37607 100644 --- a/web/src/app/(main)/system/llm/page.tsx +++ b/web/src/app/(main)/system/llm/page.tsx @@ -17,25 +17,16 @@ import { toast } from "sonner"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; -import { - Dialog, - DialogClose, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, - DialogTrigger, -} from "@/components/ui/dialog"; +import { Card, CardContent } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; +import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; import { Switch } from "@/components/ui/switch"; import { api } from "@/lib/api"; -import type { LLMPoolStatus, LLMProfile } from "@/lib/types"; +import type { LLMPoolMember, LLMPoolStatus, LLMProfile } from "@/lib/types"; import { cn } from "@/lib/utils"; // 思考强度档位(仅在开关打开时生效)。 @@ -59,54 +50,74 @@ const toStore = (mode: ThinkMode, effort: string) => (mode === "on" ? effort : m const modeFromStore = (v?: string): ThinkMode => (!v ? "none" : v === "off" ? "off" : "on"); const effortFromStore = (v?: string) => (v && v !== "off" ? v : "high"); -// 熔断状态 → 徽章样式。degraded 是「在失败但还没到熔断阈值」的中间态。 -const STATE_BADGE: Record = { - ok: { label: "正常", cls: "border-emerald-500/50 text-emerald-600 dark:text-emerald-400" }, - degraded: { label: "异常", cls: "border-amber-500/50 text-amber-600 dark:text-amber-400" }, - tripped: { label: "已熔断", cls: "border-destructive/50 text-destructive" }, -}; - function cooldownText(secs: number) { if (secs <= 0) return ""; if (secs < 60) return `${secs}s`; return `${Math.ceil(secs / 60)}min`; } -// PoolCard 是轮询功能的总控:主开关、绑定兜底开关,以及后端算出来的**实际**链路顺序。 -// 顺序预览很关键——「激活 + 优先级 + 不参与轮询」三个因素叠加后,光看单张配置卡片 -// 是猜不出真实顺序的。 -function PoolCard({ refreshKey }: { refreshKey: number }) { - const [pool, setPool] = React.useState(null); +// 一个配置在卡片上显示的「是否正常」。没填 Key 的配置根本发不出请求,比熔断更该先说; +// 其余状态来自轮询的熔断记录(轮询关着时不会产生新记录,此时「正常」= 没有已知故障)。 +type Health = { label: string; cls: string; hint?: string }; +function healthOf(p: LLMProfile, m?: LLMPoolMember): Health { + if (!p.api_key_hint) { + return { + label: "未配置 Key", + cls: "border-muted-foreground/40 text-muted-foreground", + hint: "未填 API Key,无法调用", + }; + } + if (m?.state === "tripped") { + return { + label: m.cooldown_secs > 0 ? `已熔断 · ${cooldownText(m.cooldown_secs)}` : "已熔断", + cls: "border-destructive/50 text-destructive", + hint: m.last_error, + }; + } + if (m?.state === "degraded") { + return { + label: `异常 · 失败 ${m.fails} 次`, + cls: "border-amber-500/50 text-amber-600 dark:text-amber-400", + hint: m.last_error, + }; + } + return { label: "正常", cls: "border-emerald-500/50 text-emerald-600 dark:text-emerald-400" }; +} + +// ───────────────────────────────────────────────────────────────────────────── +// 轮询配置抽屉 +// ───────────────────────────────────────────────────────────────────────────── + +function PoolSheet({ + open, + onOpenChange, + pool, + onReload, +}: { + open: boolean; + onOpenChange: (o: boolean) => void; + pool: LLMPoolStatus | null; + onReload: () => Promise; +}) { const [busy, setBusy] = React.useState(false); - const load = React.useCallback(async () => { - try { - setPool(await api.llmPool()); - } catch { - /* ignore */ - } - }, []); - // biome-ignore lint/correctness/useExhaustiveDependencies: refreshKey is a prop — bumping it is exactly how the parent asks for a refetch after a profile edit + // 冷却倒计时是后端算出的剩余秒数——抽屉开着且有配置不正常时才定时拉,让它走起来。 React.useEffect(() => { - void load(); - }, [load, refreshKey]); - // 冷却倒计时是后端算的剩余秒数,定时拉一次让它走起来(只在有熔断时才轮询)。 - React.useEffect(() => { - if (!pool?.enabled || !pool.chain.some((m) => m.state !== "ok")) return; - const t = setInterval(() => void load(), 10_000); + if (!open || !pool?.enabled || !pool.chain.some((m) => m.state !== "ok")) return; + const t = setInterval(() => void onReload(), 10_000); return () => clearInterval(t); - }, [pool, load]); + }, [open, pool, onReload]); async function toggle(patch: { llm_pool_enabled?: boolean; llm_pool_bind_fallback?: boolean }) { if (busy) return; setBusy(true); try { await api.setSettings(patch); - await load(); + await onReload(); if (patch.llm_pool_enabled !== undefined) { toast.success(patch.llm_pool_enabled ? "已开启 LLM 轮询" : "已关闭 LLM 轮询"); } else { - toast.success("已更新绑定兜底设置"); + toast.success("已更新兜底设置"); } } catch (e) { toast.error(`设置失败:${(e as Error).message}`); @@ -117,7 +128,8 @@ function PoolCard({ refreshKey }: { refreshKey: number }) { async function recover(id?: string) { try { - setPool(await api.resetLLMPool(id)); + await api.resetLLMPool(id); + await onReload(); toast.success(id ? "已恢复该配置" : "已恢复全部配置"); } catch (e) { toast.error(`恢复失败:${(e as Error).message}`); @@ -125,180 +137,214 @@ function PoolCard({ refreshKey }: { refreshKey: number }) { } const enabled = pool?.enabled ?? false; - // 链路顺序 = 后端返回的顺序;不参与轮询的配置仍会列出(灰显并标注), - // 这样用户能看见「它为什么不在队列里」,而不是凭空消失。 const chain = pool?.chain ?? []; + // 参与轮询的成员(排除被标记「不参与轮询」的),顺序即后端实际的尝试顺序。 const inChain = chain.filter((m) => m.active || !m.excluded); const tripped = chain.filter((m) => m.state === "tripped"); return ( - - - - LLM 轮询 · 故障转移 - - - 开启后,未指定模型的 Agent 在当前配置不可用(余额不足 / Key 失效 / 限流 / - 服务异常)时自动切换到下一个配置。 指定了模型的 Agent 与任务默认不参与轮询。 - - - void toggle({ llm_pool_enabled: v })} - aria-label="LLM 轮询开关" - /> - - - {enabled && ( - + + + + + LLM 轮询 · 故障转移 + + + 开启后,未指定模型的 Agent 在当前配置不可用(余额不足 / Key 失效 / 限流 / + 服务异常)时自动切到下一个配置。 + + + +
- -

- 默认关闭:Agent 或任务指定了某个配置就只用它,失败即失败(不会悄悄换成别的模型)。 - 开启后,指定的配置失败时也会回落到下面的轮询链。 -

+ +

默认关闭。关闭时始终只用激活配置,失败即失败。

void toggle({ llm_pool_bind_fallback: v })} - aria-label="绑定配置失败兜底开关" + onCheckedChange={(v) => void toggle({ llm_pool_enabled: v })} + aria-label="LLM 轮询开关" />
-
-
- - {tripped.length > 0 && ( - - )} -
- {inChain.length < 2 && ( -

- 当前只有 {inChain.length} 个可用配置,轮询不会生效——至少需要 2 个已填 API Key 且参与轮询的配置。 -

- )} -
- {chain.map((m) => { - const badge = STATE_BADGE[m.state] ?? STATE_BADGE.ok; - const order = - m.excluded && !m.active ? null : inChain.findIndex((x) => x.profile_id === m.profile_id) + 1; - return ( -
- - {order ?? "—"} - - {m.name} - {m.model} - {m.active && ( - - 激活 - - )} - {m.excluded && !m.active && 不参与轮询} - {!m.active && 优先级 {m.priority}} -
- {m.state === "tripped" && m.cooldown_secs > 0 && ( - 冷却 {cooldownText(m.cooldown_secs)} + {enabled && ( + <> +
+
+ +

+ 默认关闭:Agent 或任务指定了某个配置就只用它,失败即失败(不会悄悄换成别的模型)。 + 开启后,指定的配置失败时也会回落到下面的轮询链。 +

+
+ void toggle({ llm_pool_bind_fallback: v })} + aria-label="绑定配置失败兜底开关" + /> +
+ + + +
+
+ + {tripped.length > 0 && ( + + )} +
+ {inChain.length < 2 && ( +

+ 当前只有 {inChain.length} 个可用配置,轮询不会生效——至少需要 2 个已填 API Key 且参与轮询的配置。 +

+ )} + {chain.map((m) => { + const excluded = m.excluded && !m.active; + const order = excluded ? null : inChain.findIndex((x) => x.profile_id === m.profile_id) + 1; + return ( +
连续失败 {m.fails} 次 - )} - - {badge.label} - - {m.state !== "ok" && ( - + > +
+ + {order ?? "—"} + + {m.name} + {m.active && ( + + 激活 + + )} + {excluded && 不参与轮询} +
+ {m.state === "tripped" && m.cooldown_secs > 0 && ( + 冷却 {cooldownText(m.cooldown_secs)} + )} + {m.state === "degraded" && ( + 连续失败 {m.fails} 次 + )} + {m.state !== "ok" && ( + + )} +
+
+
+ {m.model} + {!m.active && 优先级 {m.priority}} +
+ {m.last_error && ( +

+ {m.last_error} +

)}
- {m.last_error && ( -

- {m.last_error} -

- )} + ); + })} + {chain.length === 0 && ( +
+ 暂无配置
- ); - })} - {chain.length === 0 && ( -
- 暂无配置 -
- )} -
-

- 激活配置恒为第 1 顺位,其余按优先级从高到低。某个配置失败后进入冷却(60s → 5min → 30min), - 冷却期内被跳过;恢复后自动切回。上下文窗口装不下当前请求的配置会被跳过。 -

-
- - )} - + )} +
+ +
+ 激活配置恒为第 1 顺位,其余按优先级从高到低(在各配置里设置)。某个配置失败后进入冷却 (60s → 5min → + 30min),冷却期内被跳过,恢复后自动切回。上下文窗口装不下当前请求的配置会被跳过。 指定了模型的 Agent + 与任务默认不参与轮询。 +
+ + )} +
+ + ); } -function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { - const [open, setOpen] = React.useState(false); +// ───────────────────────────────────────────────────────────────────────────── +// 模型配置抽屉(新建 / 编辑共用同一套表单) +// ───────────────────────────────────────────────────────────────────────────── + +function ProfileSheet({ + profile, + open, + onOpenChange, + onSaved, +}: { + profile: LLMProfile | null; // null = 新建 + open: boolean; + onOpenChange: (o: boolean) => void; + onSaved: (id: string) => void; +}) { + const isNew = !profile; const [name, setName] = React.useState(""); const [format, setFormat] = React.useState<"anthropic" | "openai">("anthropic"); const [model, setModel] = React.useState(""); const [baseUrl, setBaseUrl] = React.useState(""); const [proxy, setProxy] = React.useState(""); const [apiKey, setApiKey] = React.useState(""); + const [keyHint, setKeyHint] = React.useState(""); const [rps, setRps] = React.useState("0"); const [rpm, setRpm] = React.useState("0"); - const [cw, setCw] = React.useState("0"); + const [cw, setCw] = React.useState("0"); // 上下文窗口(K tokens);0=默认200K const [thinkMode, setThinkMode] = React.useState("none"); const [effort, setEffort] = React.useState("high"); - const [priority, setPriority] = React.useState("0"); + const [priority, setPriority] = React.useState("0"); // 轮询顺位;越大越先 const [poolExclude, setPoolExclude] = React.useState(false); + const [testing, setTesting] = React.useState(false); + const [saving, setSaving] = React.useState(false); const [models, setModels] = React.useState([]); const [loadingModels, setLoadingModels] = React.useState(false); const [modelsOpen, setModelsOpen] = React.useState(false); - function reset() { - setName(""); - setFormat("anthropic"); - setModel(""); - setBaseUrl(""); - setProxy(""); + // 每次打开时从传入的 profile 灌一遍表单(新建则重置为默认值)。抽屉关掉再打开 + // 就是一次干净的开始,不会留下上一个配置的残影。 + React.useEffect(() => { + if (!open) return; + setName(profile?.name ?? ""); + setFormat(profile?.format === "openai" ? "openai" : "anthropic"); + setModel(profile?.model ?? ""); + setBaseUrl(profile?.base_url ?? ""); + setProxy(profile?.proxy ?? ""); + setRps(String(profile?.rate_per_second ?? 0)); + setRpm(String(profile?.rate_per_minute ?? 0)); + setCw(String(profile?.context_window_k ?? 0)); + setThinkMode(modeFromStore(profile?.reasoning_effort)); + setEffort(effortFromStore(profile?.reasoning_effort)); + setPriority(String(profile?.priority ?? 0)); + setPoolExclude(profile?.pool_exclude ?? false); setApiKey(""); - setRps("0"); - setRpm("0"); - setCw("0"); - setThinkMode("none"); - setEffort("high"); - setPriority("0"); - setPoolExclude(false); + setKeyHint(profile?.api_key_hint ?? ""); setModels([]); setModelsOpen(false); - } + }, [open, profile]); + + const profileId = profile ? Number(profile.id) : undefined; async function loadModels() { if (loadingModels) return; setLoadingModels(true); setModels([]); try { - const r = await api.fetchLLMModels(format, baseUrl, apiKey, proxy); + const r = await api.fetchLLMModels(format, baseUrl, apiKey, proxy, profileId); if (r.ok && r.models && r.models.length > 0) { setModels(r.models); setModelsOpen(true); @@ -313,13 +359,32 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { } } - async function create() { + async function testConnection() { + if (testing) return; + setTesting(true); + try { + // 用配置实际会跑的思考参数来测,这样不支持该字段的模型在这里就失败, + // 而不是等到跑任务时才炸。传 profile id:Key 输入框留空时用已存的 Key。 + const r = await api.testLLM(format, model, baseUrl, apiKey, proxy, toStore(thinkMode, effort), profileId); + if (r.ok) toast.success(`连接成功 · ${r.latency_ms ?? "?"}ms · ${r.model ?? model}`); + else toast.error(`连接失败:${r.error ?? "未知"}`); + } catch (e) { + toast.error(`测试出错:${(e as Error).message}`); + } finally { + setTesting(false); + } + } + + async function save() { if (!name.trim() || !model.trim()) { toast.error("请填写名称与模型"); return; } + if (saving) return; + setSaving(true); try { const { id } = await api.saveLLMProfile({ + ...(profile ? { id: Number(profile.id) } : {}), name: name.trim(), format, model: model.trim(), @@ -333,48 +398,52 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) { priority: Number(priority) || 0, pool_exclude: poolExclude, }); - toast.success(`已新建 Profile:${name.trim()}(在列表中「设为激活」以启用)`); - reset(); - setOpen(false); - onCreated(String(id)); + if (isNew) toast.success(`已新建:${name.trim()}(在卡片上「设为激活」以启用)`); + else toast.success(profile?.is_default ? "已保存,激活配置即时生效,无需重启" : "已保存"); + onSaved(String(id)); + onOpenChange(false); } catch (e) { - toast.error(`新建失败:${(e as Error).message}`); + toast.error(`保存失败:${(e as Error).message}`); + } finally { + setSaving(false); } } return ( - { - setOpen(o); - if (!o) reset(); - }} - > - - - - + e.preventDefault()} > - - 新建 Profile - 新建后不会自动激活,请在列表中「设为激活」以启用。 - -
-
- - setName(e.target.value)} - /> -
+ + + {isNew ? "新建模型配置" : `编辑:${profile?.name}`} + {profile?.is_default && ( + + 激活中 + + )} + + + {isNew + ? "新建后不会自动激活,请在卡片上「设为激活」以启用。" + : "修改后点击保存;激活配置保存后对全部 Agent 立即生效。"} + + + +
+
+ + setName(e.target.value)} + /> +
-
- -
- setModel(e.target.value)} - /> - {/* modal: this Popover lives inside a Dialog, and its content is - portaled to — outside the Dialog's react-remove-scroll - lock, which then cancels every wheel event over it (the list - rendered but refused to scroll). modal makes the Popover own - the topmost scroll lock, so its own list scrolls again. */} - - - - - {models.length > 0 && ( - - {models.map((m) => ( - - ))} - - )} - -
+
+ +
+ +
+ setModel(e.target.value)} + /> + {/* modal: 这个 Popover 的内容被 portal 到 ,在 Sheet 的滚动锁之外, + 不加 modal 时列表能渲染却滚不动。modal 让它自己持有最上层滚动锁。 */} + + + + + {models.length > 0 && ( + + {models.map((m) => ( + + ))} + + )} +
+
- + setBaseUrl(e.target.value)} />
+
- + setProxy(e.target.value)} /> +

+ 仅 LLM 出站请求走此代理,支持 http/https/socks5;留空则用环境变量(HTTP_PROXY/HTTPS_PROXY)。 +

+
- + setApiKey(e.target.value)} />
+
- - setRps(e.target.value)} /> + + setRps(e.target.value)} />
- - setRpm(e.target.value)} /> + + setRpm(e.target.value)} />
- + void }) { />
+

+ 限速 0 = 不限,全 Agent 共享。上下文窗口单位 K(千 token),0 = 默认 200K,上限 1000(即 + 1M);设太高会导致压缩不触发。 +

+
-
void }) {

- 开启后不会被当作故障转移目标(仍可被 Agent / 任务显式指定使用)。 + 开启后不会被当作故障转移目标(仍可被 Agent / 任务显式指定使用)。 适合「只给某个 Agent + 专用、不希望别人失败时烧掉」的昂贵配置。

+
@@ -541,7 +624,9 @@ function NewProfileDialog({ onCreated }: { onCreated: (id: string) => void }) {
-

推理投入档位,越高越深入。

+

+ 推理投入档位,越高越深入(OpenAI reasoning_effort / Anthropic output_config.effort)。 +

setName(e.target.value)} - /> -
-
- - -
-
- -
- -
- setModel(e.target.value)} - /> - - - - - {models.length > 0 && ( - - {models.map((m) => ( - - ))} - - )} - -
-
- -
- - setBaseUrl(e.target.value)} - /> -
- -
- - setProxy(e.target.value)} - /> -

- 仅 LLM 出站请求走此代理,支持 http/https/socks5;留空则用环境变量(HTTP_PROXY/HTTPS_PROXY)。 -

-
- -
- - setApiKey(e.target.value)} - /> -
- -
-
- - setRps(e.target.value)} /> -
-
- - setRpm(e.target.value)} /> -

0 = 不限

-
-
- -
- - setCw(e.target.value)} - placeholder="200" - /> -

- 模型上下文窗口,用于压缩阈值。单位 K(千 token);0/留空 = 默认 200K;上限 1000(即 - 1M)。设太高会导致压缩不触发。 -

-
- -
-
-
- -

- 数字越大越先被选中;激活配置恒为第 1 - 顺位,与本值无关。相同优先级的配置会轮流打头,天然分摊额度。 -

-
- setPriority(e.target.value)} - /> -
-
-
- -

- 开启后不会被当作故障转移目标(仍可被 Agent / 任务显式指定使用)。 适合「只给某个 Agent - 专用、不希望别人失败时烧掉」的昂贵配置。 -

-
- -
-
- -
-
-
- -

- 不发送=不带思考字段(兼容 MiniMax 等不支持该字段的模型);显式关闭=发 - disabled(默认开思考的模型才需要);开启=先推理再作答。 -

-
- -
- {thinkMode === "on" && ( -
-
- -

- 推理投入档位,越高越深入(OpenAI reasoning_effort / Anthropic output_config.effort)。 -

-
- -
- )} -
- - -
- - -
- - ) : ( - -
- 还没有 Profile,点击左侧「新建」创建第一个。 -
-
+
+
+

LLM

+

+ 全 Agent 共享的格式 / 模型 / 限速配置。点击卡片编辑,星标为当前激活配置。 +

+
+
+ - -
-
- ))} - {profiles.length === 0 && ( -
- 暂无 Profile -
- )} - - + +
+ +
+ {profiles.map((p) => { + const h = healthOf(p, health.get(p.id)); + return ( + // biome-ignore lint/a11y/useSemanticElements: 卡片内含自己的操作按钮,用原生 + +
+ + + ); + })} + {profiles.length === 0 && ( +
+ 还没有模型配置,点击右上角「新建」创建第一个。 +
+ )} +
+ + void load()} /> +
); } From 9793b522d93e14180dd52a90df21682f0652362d Mon Sep 17 00:00:00 2001 From: pphua Date: Thu, 20 Aug 2026 18:19:10 +0800 Subject: [PATCH 13/70] feat(tasks): add inherited context and resilient LLM execution Add read-only task context inheritance, ordered task-level LLM failover, live Worker controls, cascading task cleanup, FIFO task concurrency, cancellation diagnostics, UI updates, tests, build script, and changelog. --- CHANGELOG.md | 37 + agent/blackboard_inheritance_test.go | 198 +++++ agent/cancelcause.go | 124 ++- agent/capture.go | 13 +- agent/goals.go | 10 + agent/mainagent.go | 22 +- agent/planner.go | 21 +- agent/provider.go | 107 ++- agent/provider_quota_test.go | 131 +++ agent/terminalreason.go | 135 +-- agent/terminalreason_test.go | 161 ++-- agent/tools.go | 697 ++++++++++++---- agent/tools_insert.go | 9 +- agent/tools_overview_test.go | 31 + agent/worker.go | 14 +- build.sh | 65 ++ cmd/artex/main.go | 23 +- cmd/artex/main_test.go | 19 + db/assets.go | 196 +++-- db/assets_test.go | 10 +- db/config.go | 281 ++++++- db/config_test.go | 14 + db/conversation.go | 51 +- db/exploration.go | 375 ++++++++- db/exploration_sources.go | 350 ++++++++ db/exploration_sources_test.go | 431 ++++++++++ db/exploration_test.go | 136 ++- db/findings.go | 3 +- db/schema.sql | 80 +- db/task_assets_context.go | 215 +++++ db/task_context.go | 388 +++++++++ db/task_context_lock_test.go | 527 ++++++++++++ db/task_context_unit_test.go | 27 + db/task_delete_concurrency_test.go | 210 +++++ db/task_scope.go | 14 +- db/task_source_limit_test.go | 20 + db/tasks.go | 296 ++++++- db/tasks_test.go | 410 +++++++++- llmrec/llmrec.go | 41 +- llmrec/llmrec_test.go | 29 + server/assets.go | 7 +- server/chatupload.go | 13 + server/core_test.go | 78 +- server/dto.go | 162 ++-- server/engine.go | 564 +++++++++---- server/engine_cancelcause_test.go | 95 ++- server/engine_llm_calls_test.go | 43 + server/engine_timeout.go | 49 +- server/goals.go | 150 ++-- server/goals_test.go | 12 + server/inheritance_api_test.go | 100 +++ server/inheritance_dto_test.go | 117 +++ server/manager.go | 394 ++++++++- server/manager_delete_files_test.go | 280 +++++++ server/orchestration.go | 10 +- server/server.go | 774 +++++++++++++++--- server/server_mgmt.go | 182 ++-- server/task_delete_barrier_test.go | 220 +++++ server/task_llm.go | 462 +++++++++++ server/task_llm_test.go | 507 ++++++++++++ traffic/traffic.go | 214 ++++- traffic/traffic_test.go | 146 ++++ .../app/(main)/_components/main-content.tsx | 4 +- .../(main)/function/findings/detail/page.tsx | 147 ++-- .../tasks/detail/_tabs/assets-tab.tsx | 10 +- .../tasks/detail/_tabs/coverage-graph-tab.tsx | 21 +- .../tasks/detail/_tabs/findings-tab.tsx | 100 +-- .../tasks/detail/_tabs/report-tab.tsx | 46 +- .../tasks/detail/_tabs/sessions-tab.tsx | 422 ++++++++-- .../app/(main)/function/tasks/detail/page.tsx | 282 +++++-- web/src/app/(main)/function/tasks/page.tsx | 644 ++++++++++----- web/src/components/agent-editor.tsx | 7 +- web/src/components/exploration-graph.tsx | 21 +- web/src/components/task-llm-profile-chain.tsx | 223 +++++ web/src/components/transcript.tsx | 7 +- web/src/components/ui/checkbox.tsx | 8 +- web/src/components/ui/combobox.tsx | 7 +- web/src/lib/api.ts | 74 +- web/src/lib/mock/data.ts | 34 +- web/src/lib/mock/handler.ts | 189 ++++- web/src/lib/status.ts | 1 + web/src/lib/types.ts | 59 +- 82 files changed, 11094 insertions(+), 1712 deletions(-) create mode 100644 CHANGELOG.md create mode 100644 agent/blackboard_inheritance_test.go create mode 100644 agent/provider_quota_test.go create mode 100644 agent/tools_overview_test.go create mode 100755 build.sh create mode 100644 cmd/artex/main_test.go create mode 100644 db/exploration_sources.go create mode 100644 db/exploration_sources_test.go create mode 100644 db/task_assets_context.go create mode 100644 db/task_context.go create mode 100644 db/task_context_lock_test.go create mode 100644 db/task_context_unit_test.go create mode 100644 db/task_delete_concurrency_test.go create mode 100644 db/task_source_limit_test.go create mode 100644 llmrec/llmrec_test.go create mode 100644 server/engine_llm_calls_test.go create mode 100644 server/goals_test.go create mode 100644 server/inheritance_api_test.go create mode 100644 server/inheritance_dto_test.go create mode 100644 server/manager_delete_files_test.go create mode 100644 server/task_delete_barrier_test.go create mode 100644 server/task_llm.go create mode 100644 server/task_llm_test.go create mode 100644 web/src/components/task-llm-profile-chain.tsx diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..4c53e7e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,37 @@ +# 更新日志 + +本项目的重要变更记录在此文件中,格式参考 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/)。 + +## [Unreleased] - 2026-08-20 + +### Added + +- 新建任务支持关联多个已有任务,实时、只读继承直接来源任务的事实、漏洞、已完成意图、资产范围和黑板上下文;新任务仍使用独立的执行队列、工作目录和会话历史。 +- 任务支持有序 LLM 配置链。明确识别到供应商额度不足时自动切换到下一配置,并持久化当前配置、耗尽状态和结构化审计活动。 +- 运行中和暂停任务支持编辑 LLM 配置链、调整顺序及手动切换当前配置;自动切换、手动切换和全链耗尽会在任务详情页提示。 +- 运行中的 Worker 支持暂停、恢复和取消。暂停保留黑板数据,取消会在 Worker 停止写入后事务性清理该意图及其直接产出的事实、漏洞和执行记录。 +- 删除任务可选择同时清理关联资产、流量、漏洞和任务落地文件,并增加删除屏障、并发保护和可审计的删除统计。 +- 增加可选的任务并发上限;达到上限的新任务按 FIFO 排队,并在运行槽位释放后自动启动。 +- 增加任务资产分页、公司资产分页、任务分类统计及分页意图 Mock 契约。 +- 增加 `build.sh` 单二进制构建脚本,支持静态前端导出、资源嵌入、跨平台目标和构建版本注入。 + +### Changed + +- 任务级显式 LLM 配置链与现有全局 LLM Pool 并存;显式链继续采用严格的额度故障转移语义,无显式链时沿用 Agent 绑定和全局配置规则。 +- 主 Agent 输入框支持自动增长的多行输入,`Enter` 发送、`Shift+Enter` 换行,并避免中文输入法组合态误发送。 +- Planner、Worker 和主 Agent 共享任务级 LLM runtime;任务链使用候选模型中最小上下文窗口作为安全压缩阈值。 +- 任务详情页按真实 LLM 调用状态展示运行、空闲和暂停状态;来源任务的事实、漏洞、意图、资产引用和图谱节点统一标记来源并保持只读。 +- Agent 预览、任务报告和相关详情视图统一使用共享 Markdown 渲染组件。 +- LLM 模型配置采用卡片和抽屉交互,模型列表在抽屉内支持独立滚动。 + +### Fixed + +- 修复暂停任务后主 Agent 会话仍继续运行,以及编排 Agent 暂停任务时未同步停止主 Agent 的问题。 +- 修复任务完成、Planner、Worker 或主 Agent 实际运行时,任务状态徽标与操作按钮状态不同步的问题。 +- 修复 Worker 取消、任务删除和并发写入之间可能产生迟到黑板写入或残留文件的问题。 +- 修复 Markdown 预览失效、删除确认长文本溢出、移动端宽度和部分任务详情按钮未对齐的问题。 +- 为所有 Agent 取消路径增加具名终止原因,活动详情可显示取消方、终态、运行轮数、耗时、Token 使用量和未返回的工具调用。 +- 修复后端关闭时父 context 可能抢先覆盖具名 `shutdown` 原因的竞态,以及中文活动摘要按字节截断导致乱码的问题。 +- 修复带部分流式输出的取消事件丢失真实终止原因的问题,同时保留取消前已生成的内容。 + +[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.1...HEAD diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go new file mode 100644 index 0000000..978c8bf --- /dev/null +++ b/agent/blackboard_inheritance_test.go @@ -0,0 +1,198 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "github.com/Autumn-27/artex/db" + actool "github.com/Autumn-27/norma/tool" +) + +func callReadJSON(t *testing.T, tool actool.CoreTool, input string) any { + t.Helper() + result, err := tool.Call(context.Background(), json.RawMessage(input), nil) + if err != nil { + t.Fatalf("tool call: %v", err) + } + var out any + if err := json.Unmarshal([]byte(result.Flatten()), &out); err != nil { + t.Fatalf("decode tool result: %v; raw=%s", err, result.Flatten()) + } + return out +} + +func TestBlackboardToolsReadDirectSources(t *testing.T) { + d := testDB(t) + defer d.Close() + + grand, err := d.CreateTask("grand", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("source", "source goal", db.TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}}) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + }) + + grandStore := d.Exploration(grand.ExplorationID) + sourceStore := d.Exploration(source.ExplorationID) + currentStore := d.Exploration(current.ExplorationID) + grandFact, err := grandStore.AddNode(db.KindFact, map[string]any{"summary": "indirect-only"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFact, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "shared fact", "confidence": "verified"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "shared work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil { + t.Fatal(err) + } + sourceFinding, err := sourceStore.AddNode(db.KindFinding, map[string]any{ + "summary": "shared finding", "vulnclass": "idor", "severity": "high", + }, 9, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + if err := sourceStore.Link(sourceIntent, db.RelYields, sourceFinding); err != nil { + t.Fatal(err) + } + stepID, err := sourceStore.AppendActivity(db.Activity{ + NodeID: &sourceIntent, Worker: "source-worker", Kind: "result", Tool: "HTTP", + Summary: "shared trace marker", Detail: "shared trace full detail", + }) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 301; i++ { + if _, err := sourceStore.AddIntent(map[string]any{"summary": fmt.Sprintf("newer live source work %d", i)}, 1, nil, "planner"); err != nil { + t.Fatal(err) + } + } + assets := d.Assets() + host := fmt.Sprintf("overview-host-%d.invalid", source.ID) + assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) }) + if _, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "source host anchor"}, 0, "confirmed", "worker", []int64{assetID}); err != nil { + t.Fatal(err) + } + + tools := NewToolSet(currentStore, "worker") + tools.SetTaskID(current.ID) + tools.SetAssetStore(assets, assets.Companies()) + parentID, _ := json.Marshal(sourceFact) + derivedID, err := tools.addOneIntent(intentItem{ + Summary: "derive locally from shared fact", ParentIDs: []json.RawMessage{parentID}, + }) + if err != nil { + t.Fatalf("derive from inherited fact: %v", err) + } + if derived, err := currentStore.GetNode(derivedID); err != nil || derived == nil { + t.Fatalf("derived intent must be local: node=%+v err=%v", derived, err) + } + beforeFacts, _ := currentStore.ListByKind(db.KindFact, 100) + sourceIntentID, _ := json.Marshal(sourceIntent) + if _, err := tools.recordOneFact(factItem{Summary: "must not attach to inherited intent", IntentID: json.RawMessage(sourceIntentID)}, sourceIntent); err == nil { + t.Fatal("record_fact must reject inherited intent") + } + afterFacts, _ := currentStore.ListByKind(db.KindFact, 100) + if len(afterFacts) != len(beforeFacts) { + t.Fatalf("rejected inherited write persisted a fact: before=%d after=%d", len(beforeFacts), len(afterFacts)) + } + overview := tools.graphOverviewData() + related, ok := overview["related_tasks"].([]map[string]any) + if !ok || len(related) != 1 { + t.Fatalf("related_tasks: %#v", overview["related_tasks"]) + } + if related[0]["source_task_id"] != source.ID || related[0]["inherited"] != true { + t.Fatalf("related source provenance: %#v", related[0]) + } + if facts, ok := related[0]["recent_facts"].([]map[string]any); !ok || len(facts) == 0 || facts[0]["id"] != sourceFact { + t.Fatalf("related fact summary: %#v", related[0]["recent_facts"]) + } + if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding { + t.Fatalf("related finding summary: %#v", related[0]["recent_findings"]) + } + results, ok := related[0]["recent_intent_results"].([]map[string]any) + if !ok || len(results) == 0 || results[0]["id"] != sourceIntent { + t.Fatalf("terminal source intent was starved by newer live work: %#v", related[0]["recent_intent_results"]) + } + if results[0]["result_summary"] != "shared trace marker" { + t.Fatalf("terminal source intent result was not distilled: %#v", results[0]) + } + coverage, ok := overview["coverage"].(map[string]any) + if !ok { + t.Fatalf("coverage missing from overview: %#v", overview["coverage"]) + } + hosts, ok := coverage["hosts"].([]string) + if !ok || len(hosts) != 1 || hosts[0] != host || coverage["host_count"] != 1 { + t.Fatalf("inherited host context missing: %#v", coverage) + } + + facts := callReadJSON(t, tools.listFacts(), `{}`).([]any) + seenSource, seenGrand := false, false + for _, raw := range facts { + item := raw.(map[string]any) + id := int64(item["id"].(float64)) + if id == sourceFact { + seenSource = item["inherited"] == true && int64(item["source_task_id"].(float64)) == source.ID + } + seenGrand = seenGrand || id == grandFact + } + if !seenSource || seenGrand { + t.Fatalf("list_facts direct-only: source=%v grand=%v payload=%#v", seenSource, seenGrand, facts) + } + + findings := callReadJSON(t, tools.listFindings(), `{}`).([]any) + if len(findings) != 1 { + t.Fatalf("list_findings: %#v", findings) + } + finding := findings[0].(map[string]any) + if finding["inherited"] != true || int64(finding["task_id"].(float64)) != source.ID || int64(finding["intent_id"].(float64)) != sourceIntent { + t.Fatalf("inherited finding provenance: %#v", finding) + } + + detailInput, _ := json.Marshal(map[string]any{"id": sourceFact}) + detail := callReadJSON(t, tools.nodeDetail(), string(detailInput)).(map[string]any) + if detail["inherited"] != true || int64(detail["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited node detail: %#v", detail) + } + + traceInput, _ := json.Marshal(map[string]any{"intent_id": sourceIntent}) + trace := callReadJSON(t, tools.getWorkerTrace(), string(traceInput)).(map[string]any) + if trace["inherited"] != true || int64(trace["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited trace: %#v", trace) + } + steps := trace["steps"].([]any) + if len(steps) != 1 || int64(steps[0].(map[string]any)["step_id"].(float64)) != stepID { + t.Fatalf("inherited trace steps: %#v", steps) + } + output := callReadJSON(t, tools.getWorkerOutput(), string(traceInput)).(map[string]any) + if output["inherited"] != true || output["final_text"] != "shared trace full detail" { + t.Fatalf("inherited worker output: %#v", output) + } + + search := callReadJSON(t, tools.searchAllWorkerTraces(), `{"q":"shared trace marker"}`).(map[string]any) + hits := search["hits"].([]any) + if len(hits) != 1 || hits[0].(map[string]any)["inherited"] != true { + t.Fatalf("inherited trace search: %#v", hits) + } +} diff --git a/agent/cancelcause.go b/agent/cancelcause.go index cfda073..f4705bc 100644 --- a/agent/cancelcause.go +++ b/agent/cancelcause.go @@ -6,25 +6,12 @@ import ( "fmt" ) -// AbortCause names WHY a planner/worker run's context was cancelled. -// -// Background: when a run's ctx is cancelled mid-flight the harness finishes with -// ReasonAbortedStreaming / ReasonAbortedTools and Terminal.Err = context.Canceled. -// ARTEX cancels that ctx from a dozen different places (pause, delete, kill_work, -// goal-met, settle drain backstop, chat stop, process shutdown, hard run timeout, -// …) but by the time it reaches captureRun they all collapse into one bare -// context.Canceled — so the trace could only print a guess-list, not a diagnosis. -// -// Fix: every cancel site uses context.WithCancelCause / WithTimeoutCause and -// attaches one of the causes below; captureRun recovers it via context.Cause and -// prints exactly who cancelled the run, why, and what happens to the intent. -// -// When you add a new cancel site, add its cause HERE — a site that cancels -// without one shows up in the trace as "canceled_no_cause", by design. +// AbortCause names why an agent run's context was cancelled. Every cancellation +// site should attach one so the activity trace can report the real initiator. type AbortCause struct { - Code string // stable machine-readable id (logs / future filtering) - Short string // 一句话,进折叠摘要行 - Text string // 完整说明:谁取消的、为什么、这条意图接下来会怎样 + Code string + Short string + Text string } func (c *AbortCause) Error() string { return c.Text } @@ -33,84 +20,71 @@ func cause(code, short, text string) *AbortCause { return &AbortCause{Code: code, Short: short, Text: text} } -// Causef builds an ad-hoc cause with runtime detail (e.g. which intent, which -// deadline). Use the predeclared vars below when there is nothing dynamic to add. +// Causef builds a cause that includes runtime-specific detail. func Causef(code, short, format string, args ...any) *AbortCause { return &AbortCause{Code: code, Short: short, Text: fmt.Sprintf(format, args...)} } -// The full enumeration of cancel sites, grouped by which context gets cancelled. var ( - // ---- 任务级 exec ctx(Engine.cancelExec:取消该任务在跑的 planner + 全部 worker)---- - + // Task-level execution context. AbortPausedByUser = cause("paused_by_user", "用户暂停了任务", - "用户暂停了任务(前端「暂停」按钮 / POST /api/tasks/{id}/state action=pause)。本次运行被主动取消,该意图已退回 frontier(open),恢复任务后会被重新领取、从头重跑") - - AbortPausedByOrchestrator = cause("paused_by_orchestrator", "编排 agent 调用 pause_task 暂停了任务", - "编排 agent 调用了 pause_task 工具暂停本任务。本次运行被主动取消,该意图已退回 frontier(open),恢复后重跑") - + "用户通过任务控制接口(POST /api/tasks/{id}/control,action=pause)暂停了任务。本次 Planner/Worker 运行被主动取消;运行中的意图会退回 frontier(open),恢复任务后重新领取并从头执行") + AbortPausedByOrchestrator = cause("paused_by_orchestrator", "编排 Agent 暂停了任务", + "编排 Agent 调用了 pause_task 工具暂停本任务。本次 Planner/Worker 运行被主动取消;运行中的意图会退回 frontier(open),恢复后重新执行") AbortTaskDeleted = cause("task_deleted", "任务被删除", - "任务被删除(DELETE /api/tasks/{id})——删除前先取消在跑的 planner/worker。本次运行的结果不会再被使用") + "任务正在删除(DELETE /api/tasks/{id}),删除屏障已取消该任务正在运行的 Planner、Worker 和主 Agent;本次运行结果不会再被使用") + AbortPausedOnReload = cause("paused_on_reload", "后端恢复了任务的暂停状态", + "后端启动时根据数据库中持久化的状态恢复了任务暂停。本次运行被取消;正常情况下恢复阶段没有正在运行的 Agent") + AbortGoalMet = cause("goal_met", "规划者判定任务目标已达成", + "规划者判定任务目标已达成并将任务置为 done,随后取消仍在运行的 Worker;这些意图会标记为 stopped,而不是失败") + AbortSettleDrainTimeout = cause("settle_drain_timeout", "任务超时收尾的等待时间已用尽", + "任务到达 timeout 后等待正在运行的 Worker 优雅收尾,但 90 秒 drain 宽限仍不足,因此执行硬取消;意图会标记为 exhausted,收尾阶段已经写入的事实和资产会保留") - AbortPausedOnReload = cause("paused_on_reload", "后端启动时恢复了「已暂停」状态", - "后端启动时按库里持久化的「已暂停」状态恢复暂停。本次运行被取消(正常情况下此时并没有在跑的 run)") + // Per-work context. + AbortKilledByPlanner = cause("killed_by_planner", "规划者终止了这条意图", + "规划者调用 kill_work 主动终止了这条意图,通常表示方向跑偏或已无继续价值;意图会标记为 stopped,不会自动重新领取") + AbortWorkPausedByUser = cause("work_paused_by_user", "用户暂停了这条 Worker 意图", + "用户暂停了正在运行的 Worker。本次调用被取消,意图转为 paused;已经登记的意图、事实、漏洞和活动记录全部保留,恢复后从头重新执行") + AbortWorkCancelledByUser = cause("work_cancelled_by_user", "用户取消了这条 Worker 意图", + "用户取消了正在运行的 Worker。本次调用被取消;Worker 退出写入区后,服务端会事务性删除该意图及其直接产生的事实、漏洞和执行记录") + AbortWorkFinished = cause("work_finished", "Worker 已正常结束并释放 context", + "Worker 已正常结束,引擎在 detachWork 中释放其 context 资源。这不是运行中断;若它出现在中断消息中,说明取消与收场事件发生了竞态") + AbortPausedRaceGuard = cause("paused_race_guard", "任务暂停期间拒绝启动新运行", + "任务处于暂停状态时,引擎拒绝发出新的执行 context,用于防止 claim 与暂停之间的竞态导致 Worker 继续启动;已领取的意图会退回 frontier") - AbortGoalMet = cause("goal_met", "规划者判定目标已达成,任务落终态 done", - "规划者判定任务目标已达成、任务落终态 done,随即取消所有在跑的 worker——它们手头的意图跑完也没意义了。该意图标记为 stopped,不是失败") + // Main Agent and standalone conversation contexts. + AbortChatStoppedByUser = cause("chat_stopped_by_user", "用户停止了本轮对话", + "用户点击了停止,主动中止本轮主 Agent 或会话 Agent 运行。已经产生的活动记录会保留,可以继续发送下一条消息") + AbortChatPausedWithTask = cause("chat_paused_with_task", "任务暂停并中止了主 Agent 对话", + "用户暂停任务时,正在运行的主 Agent 对话也被同步取消。已经产生的活动记录会保留;恢复任务后不会自动重放本轮消息") + AbortChatTurnFinished = cause("chat_turn_finished", "本轮对话已正常结束并释放 context", + "本轮对话已正常结束,服务端正在释放该轮 context 资源。这不是运行中断;若它出现在中断消息中,说明取消与收场事件发生了竞态") - AbortSettleDrainTimeout = cause("settle_drain_timeout", "任务级超时收尾:drain 宽限用尽后强制取消", - "任务级超时收尾的硬兜底:到达任务 timeout 后等在跑 worker 优雅收尾,超过 90s drain 宽限仍未结束 → 强制取消。该意图标记为 exhausted,收尾阶段通常已把事实/资产写回") - - // ---- 单个 work 的 ctx(Engine.workCancel[intentID])---- - - AbortKilledByPlanner = cause("killed_by_planner", "规划者用 kill_work 终止了这条意图", - "规划者调用 kill_work 主动终止了这条意图(判定方向跑偏/无意义,止损)。该意图标记为 stopped,不会被自动重领") - - AbortWorkFinished = cause("work_finished", "work 已正常结束,引擎释放其 context", - "work 已正常结束,引擎在 unregisterWork 里释放其 context 资源。这不是中断——若它出现在中断消息里,说明取消与收场事件存在竞态") - - AbortPausedRaceGuard = cause("paused_race_guard", "任务处于暂停态,引擎拒绝启动新 run", - "任务处于暂停态时请求执行上下文,引擎直接返回一个已取消的 context(守卫 claim→Execute 的竞态,避免暂停后仍启动新 run)。该意图会退回 frontier") - - // ---- 对话 / 主 Agent 的单轮 ctx(Server.chatCancel[busyKey])---- - - AbortChatStoppedByUser = cause("chat_stopped_by_user", "用户点了「停止」中止本轮对话", - "用户点了「停止」,主动中止本轮对话运行(POST .../chat/stop 或 .../conversations/{id}/stop)。已产生的步骤都保留,可以直接发下一条消息") - - AbortChatTurnFinished = cause("chat_turn_finished", "本轮对话已正常结束,服务端释放其 context", - "本轮对话已正常结束,服务端在 defer 里释放其 context 资源。这不是中断——若它出现在中断消息里,说明取消与收场事件存在竞态") - - // ---- 进程级 ---- - - AbortShutdown = cause("shutdown", "后端进程正在关停(重启 / 更新 / Ctrl-C)", - "后端进程收到 SIGINT/SIGTERM 正在关停(重启 / 更新 / 手动 Ctrl-C)。所有在跑的 planner/worker 一并取消;重启后残留的 running 意图会被重置为 open 重跑") - - // ---- 单次 run 级(worker.Execute / planner.Plan 的硬兜底)---- - - AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底触发:某一步卡死了", - "单次运行的硬超时兜底被触发:正常情况下软墙钟预算(MaxDuration)会在回合边界优雅收尾,只有当某一步(模型请求或某个工具)自身卡死、久久不返回、导致回合边界检查根本跑不到时,才会走到「软预算 + 90s 宽限」的硬 cancel。请重点看中断前最后一个 tool_use——多半是它挂住了") + // Process-level and per-run hard backstop. + AbortShutdown = cause("shutdown", "后端进程正在关闭", + "后端进程收到 SIGINT 或 SIGTERM,正在重启、更新或关闭。所有运行中的 Agent 会被取消;重启后残留的 running 意图会重置为 open 并重新执行") + AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底已触发", + "单次运行超过软墙钟预算及额外宽限,说明模型请求或某个工具长时间没有返回,导致正常的回合边界收尾无法执行。请重点检查中断前最后一个未返回的工具调用") ) -// AbortReason resolves the cause attached to a cancelled run context. ok=false -// means the ctx carries no cause at all (nothing was cancelled). +// AbortReason resolves the named cause attached to a cancelled run context. func AbortReason(ctx context.Context) (code, short, text string, ok bool) { c := context.Cause(ctx) if c == nil { return "", "", "", false } - if ac, ok := errors.AsType[*AbortCause](c); ok { + var ac *AbortCause + if errors.As(c, &ac) { return ac.Code, ac.Short, ac.Text, true } switch { case errors.Is(c, context.DeadlineExceeded): - return "deadline_exceeded", "上游 context 到达 deadline(未挂具名原因)", - "上游 context 到达 deadline,但设置方没有用 WithTimeoutCause 挂具名原因: " + c.Error(), true + return "deadline_exceeded", "上游 context 到达 deadline", + "上游 context 到达 deadline,但设置方没有通过 WithTimeoutCause 附加具名原因: " + c.Error(), true case errors.Is(c, context.Canceled): - // a plain Canceled means some cancel site still lacks a cause — say so - // plainly instead of inventing one. - return "canceled_no_cause", "取消方未挂具名原因(漏网的取消点)", - "上游 context 被取消,但取消方没有用 context.WithCancelCause 挂上具名原因。" + - "这属于漏网的取消点——请在 agent/cancelcause.go 补一条枚举并接到该取消处", true + return "canceled_no_cause", "取消方未附加具名原因", + "上游 context 被取消,但取消方没有通过 context.WithCancelCause 附加具名原因;请在 agent/cancelcause.go 登记原因并接入该取消点", true + default: + return "other", firstLine(c.Error(), 80), c.Error(), true } - return "other", firstLine(c.Error(), 80), c.Error(), true } diff --git a/agent/capture.go b/agent/capture.go index bbede4f..ca36acc 100644 --- a/agent/capture.go +++ b/agent/capture.go @@ -6,10 +6,10 @@ import ( "time" "unicode/utf8" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/agentcore" "github.com/Autumn-27/norma/harness" "github.com/Autumn-27/norma/llm" - "github.com/Autumn-27/artex/db" ) // captureRun drives one agent turn-to-completion over Session.Prompt and emits a @@ -65,9 +65,6 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tkind = kind tbuf.WriteString(text) } - - // breadcrumbs for the terminal message: when a run is aborted or times out, - // "which tool was still running, for how long" is the whole diagnosis. lastTool := &runTrace{startedAt: time.Now()} var finalText string @@ -75,7 +72,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, for ev, err := range s.Prompt(ctx, input) { if err != nil { flush() - if ctx.Err() != nil { // ctx cancelled = stopped by the engine, not a failure + if ctx.Err() != nil { // engine/user cancellation, not a provider failure sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool) rec(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}) return finalText, reason, ctx.Err() @@ -130,7 +127,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, } flush() // flush any trailing thinking / non-final text sum, detail := ev.Terminal.Text, ev.Terminal.Text - if sum == "" { + if sum == "" || ev.Terminal.Reason == harness.ReasonAbortedTools || ev.Terminal.Reason == harness.ReasonAbortedStreaming { sum, detail = terminalText(ctx, ev.Terminal, lastTool) } u := ev.Terminal.Usage // cumulative token usage for this session @@ -162,9 +159,7 @@ func blocksText(blocks []llm.ContentBlock) string { return b.String() } -// firstLine returns a single-line, length-capped preview for the summary column. -// max counts RUNES, not bytes: these previews are mostly 中文, and byte-slicing -// used to cut a multi-byte rune in half and render as "恢�…". +// firstLine returns a single-line, rune-capped preview for the summary column. func firstLine(s string, max int) string { s = strings.TrimSpace(s) if before, _, found := strings.Cut(s, "\n"); found { diff --git a/agent/goals.go b/agent/goals.go index dbbae9f..3501157 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -84,6 +84,16 @@ type GoalSpec struct { // runtime). The returned specs are read back from the store so callers can emit // per-goal activity and detect the "LLM produced nothing" case for their fallback. func DecomposeGoals(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { + if prov == nil { + return nil + } + return DecomposeGoalsWithProvider(ctx, prov, dataDir, goalText, desc, as, ts, taskID, emit) +} + +// DecomposeGoalsWithProvider is the task-runtime variant used when a task has an +// ordered provider chain. It preserves the same tools and write behavior while +// letting the caller own provider selection/failover. +func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { if prov == nil { return nil } diff --git a/agent/mainagent.go b/agent/mainagent.go index 6736ded..86b198b 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -21,6 +21,7 @@ type MainAgent struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) proxyAddr string // recording proxy for WebFetch (empty = direct) proxyCACert string // recording proxy's CA cert path (HTTPS verify) @@ -32,6 +33,15 @@ func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns} } +func (m *MainAgent) SetCompactionWindowResolver(fn func() int) { m.windowFn = fn } + +func (m *MainAgent) compactionWindow() int { + if m.windowFn != nil { + return m.windowFn() + } + return m.window +} + // SetProxy points the main agent's WebFetch at the recording proxy plus the CA // cert it trusts to verify HTTPS through it (empty addr = direct). func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert } @@ -70,9 +80,9 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t tsx.SetAssetStore(as, as.Companies()) } tsx.SetTaskID(taskID) - tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) - tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running - tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标:…」触发 + tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) + tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running + tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标:…」触发 // 领域工具 + 基础默认工具集(Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash) base := append(tsx.MainAgentTools(), actool.DefaultTools()...) tools, def, cleanup := AugmentTools(ctx, "mainagent", base) @@ -101,9 +111,9 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA WorkingDir: mainDir, // 本任务工作目录 /tasks/ ToolOutputDir: cmdOutDir(mainDir), - MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) - Compaction: compactionConfig(m.window), // long chats stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢 + MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) + Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢 // 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。 Settlement: wrapupSettlement("mainagent", nil), } diff --git a/agent/planner.go b/agent/planner.go index b567d30..b2757fe 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -24,6 +24,7 @@ type Planner struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) killWork func(intentID int64) error // engine callback to terminate a running work (nil = off) steerWork func(intentID int64, msg string) error // engine callback to steer a running work mid-run (nil = off) @@ -44,6 +45,15 @@ func NewPlanner(prov llm.Provider, model, workDir string, tx *transcript.Store, return &Planner{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, todos: map[int64]*actool.TodoStore{}} } +func (p *Planner) SetCompactionWindowResolver(fn func() int) { p.windowFn = fn } + +func (p *Planner) compactionWindow() int { + if p.windowFn != nil { + return p.windowFn() + } + return p.window +} + // SetProxy points the planner's WebFetch at the recording proxy plus the CA cert // it trusts to verify HTTPS through it (empty addr = direct). func (p *Planner) SetProxy(addr, caCert string) { p.proxyAddr, p.proxyCACert = addr, caCert } @@ -94,10 +104,11 @@ func renderPlannerTodos(items []actool.Todo) string { // TriggerEvent describes what concretely caused this planning round to fire, so // the planner looks first at the actual change instead of re-scanning the whole // overview. Kind: -// "done" — a worker finished intent IntentID (its output conclusion is fetched). -// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). -// "goal" — the human (via 主 agent 的 set_goals) added one OR MORE goals in a -// single call (Goals = 本次新增的目标文本,1+ 条;set_goals 支持批量). +// +// "done" — a worker finished intent IntentID (its output conclusion is fetched). +// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). +// "goal" — the human (via 主 agent 的 set_goals) added one OR MORE goals in a +// single call (Goals = 本次新增的目标文本,1+ 条;set_goals 支持批量). type TriggerEvent struct { Kind string IntentID int64 @@ -308,7 +319,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts ToolOutputDir: cmdOutDir(taskDir), MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management) MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余 - Compaction: compactionConfig(p.window), + Compaction: compactionConfig(p.compactionWindow()), // 跨唤醒共享的规划待办:让串行链在多轮之间保留(session 是新的,store 不是)。 Todos: p.todoFor(ts.ID()), // 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、 diff --git a/agent/provider.go b/agent/provider.go index 7a0a78e..465812b 100644 --- a/agent/provider.go +++ b/agent/provider.go @@ -8,8 +8,14 @@ package agent import ( + "bytes" "context" + "fmt" + "io" + "net/http" + "net/url" "os" + "regexp" "strings" "time" @@ -167,12 +173,16 @@ func (c Config) Provider() string { // limiter lives on the single provider instance — so planner + all workers + // main agent (which share this provider) are bounded by one shared rate limit. func (c Config) NewProvider() (llm.Provider, error) { + client, err := quotaAwareHTTPClient(c.Proxy) + if err != nil { + return nil, err + } lc := llm.Config{ - Format: c.Format, - BaseURL: c.BaseURL, - APIKey: c.APIKey, - Model: c.Model, - Proxy: c.Proxy, + Format: c.Format, + BaseURL: c.BaseURL, + APIKey: c.APIKey, + Model: c.Model, + HTTPClient: client, } // Derive the provider thinking params from the single UI selector: // "" → omit both (provider default); "off" → disable; else → enable + effort. @@ -191,6 +201,93 @@ func (c Config) NewProvider() (llm.Provider, error) { return llm.NewProvider(lc) } +// IsQuotaExhaustedMessage deliberately recognizes only explicit balance, +// billing, credit, or quota-exhaustion signals. Generic 429/rate-limit text, +// authentication failures, network errors, and server failures are excluded. +var nonFailoverHTTPStatus = regexp.MustCompile(`(?:status(?:\s+code)?|http(?:\s+status)?)\s*[=:]?\s*(?:401|403|5\d\d)\b`) +var transientQuotaLimit = regexp.MustCompile(`(?i)(?:\b(?:rpm|tpm|rpd|qps)\b|quota[_\s-]*metric|rate[_\s-]*limit|too many requests|(?:requests?|tokens?)\s+(?:per|/)\s*(?:second|minute)|(?:per|/)\s*(?:second|minute)\s+(?:requests?|tokens?)|generate[_\s-]*requests[_\s-]*per[_\s-]*(?:minute|second)|tokens?[_\s-]*per[_\s-]*(?:minute|second))`) + +func IsQuotaExhaustedMessage(message string) bool { + message = strings.ToLower(message) + // Authentication/authorization and provider-side 5xx failures never rotate, + // even when a gateway happens to echo a quota-looking phrase in the body. + if nonFailoverHTTPStatus.MatchString(message) { + return false + } + // Provider APIs frequently describe an ordinary rate limit as "quota + // exceeded", especially Google-style responses containing a quota metric. + // These limits recover with time and must stay on the current provider. + if transientQuotaLimit.MatchString(message) { + return false + } + markers := []string{ + "insufficient_quota", "quota_exceeded", "quota exceeded", "quota exhausted", + "exceeded your current quota", "billing_hard_limit_reached", + "billing hard limit", "billing_not_active", "credit balance", "insufficient credit", + "insufficient balance", "balance is too low", "payment required", "status 402", + "余额不足", "额度不足", "额度已用尽", "欠费", + } + for _, marker := range markers { + if strings.Contains(message, marker) { + return true + } + } + // gRPC RESOURCE_EXHAUSTED is overloaded for both account quota and ordinary + // request-rate limiting. Preserve it as an explicit exhaustion signal only + // when the same error does not identify a transient rate limit. + return strings.Contains(message, "resource_exhausted") && + !strings.Contains(message, "rate limit") && + !strings.Contains(message, "too many requests") +} + +// quotaAwareTransport preserves Norma's normal retry behavior except for a 429 +// whose body explicitly says the account quota/balance is exhausted. Norma's +// retry loop treats every 429 as transient; normalizing only that response to +// 402 lets a task router fail over immediately while retaining the original +// response body for provider-specific classification and audit logs. +type quotaAwareTransport struct{ base http.RoundTripper } + +func (t quotaAwareTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil || resp == nil || resp.StatusCode != http.StatusTooManyRequests { + return resp, err + } + body, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(body)) + resp.ContentLength = int64(len(body)) + if readErr != nil { + return resp, nil + } + if IsQuotaExhaustedMessage(string(body)) { + resp.StatusCode = http.StatusPaymentRequired + resp.Status = "402 Payment Required" + } + return resp, nil +} + +func quotaAwareHTTPClient(proxy string) (*http.Client, error) { + transport := http.DefaultTransport.(*http.Transport).Clone() + proxy = strings.TrimSpace(proxy) + if proxy == "" { + transport.Proxy = http.ProxyFromEnvironment + } else { + proxyURL, err := url.Parse(proxy) + if err != nil { + return nil, fmt.Errorf("llm: invalid proxy %q: %w", proxy, err) + } + switch proxyURL.Scheme { + case "http", "https", "socks5": + case "": + return nil, fmt.Errorf("llm: proxy %q missing scheme (use http://, https:// or socks5://)", proxy) + default: + return nil, fmt.Errorf("llm: unsupported proxy scheme %q (use http, https or socks5)", proxyURL.Scheme) + } + transport.Proxy = http.ProxyURL(proxyURL) + } + return &http.Client{Transport: quotaAwareTransport{base: transport}}, nil +} + // TestConnection makes a minimal real completion to verify the provider/model/ // endpoint/key actually work. Returns the round-trip latency. func TestConnection(ctx context.Context, c Config) (time.Duration, error) { diff --git a/agent/provider_quota_test.go b/agent/provider_quota_test.go new file mode 100644 index 0000000..d905ad3 --- /dev/null +++ b/agent/provider_quota_test.go @@ -0,0 +1,131 @@ +package agent + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/Autumn-27/norma/llm" +) + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestIsQuotaExhaustedMessage(t *testing.T) { + t.Parallel() + positive := []string{ + `{"error":{"code":"insufficient_quota"}}`, + `RESOURCE_EXHAUSTED`, + `You exceeded your current quota, please check your plan and billing details.`, + `billing_not_active`, + `credit balance is too low`, + `账户余额不足,请充值`, + } + for _, message := range positive { + if !IsQuotaExhaustedMessage(message) { + t.Errorf("expected quota classification for %q", message) + } + } + + negative := []string{ + `status 429: rate limit exceeded`, + `RESOURCE_EXHAUSTED: rate limit exceeded`, + `too many requests per minute`, + `quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`, + `RESOURCE_EXHAUSTED: TPM quota exceeded`, + `tokens per minute quota exceeded`, + `rate_limit_exceeded: requests per second`, + `status 401: invalid api key`, + `status 401: insufficient_quota`, + `HTTP 403: billing_hard_limit_reached`, + `status 500: internal server error`, + `status=503: insufficient_quota`, + `context length exceeded`, + } + for _, message := range negative { + if IsQuotaExhaustedMessage(message) { + t.Errorf("unexpected quota classification for %q", message) + } + } +} + +func TestQuotaAwareTransportOnlyNormalizesExplicitQuota429(t *testing.T) { + t.Parallel() + tests := []struct { + name string + status int + body string + wantStatus int + }{ + {name: "quota", status: http.StatusTooManyRequests, body: `{"code":"insufficient_quota"}`, wantStatus: http.StatusPaymentRequired}, + {name: "ordinary rate limit", status: http.StatusTooManyRequests, body: `{"message":"rate limit exceeded"}`, wantStatus: http.StatusTooManyRequests}, + {name: "server error", status: http.StatusInternalServerError, body: `insufficient_quota`, wantStatus: http.StatusInternalServerError}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + transport := quotaAwareTransport{base: roundTripperFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: tt.status, + Status: http.StatusText(tt.status), + Body: io.NopCloser(strings.NewReader(tt.body)), + Header: make(http.Header), + }, nil + })} + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.invalid", nil) + if err != nil { + t.Fatal(err) + } + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != tt.wantStatus { + t.Fatalf("status=%d, want %d", resp.StatusCode, tt.wantStatus) + } + gotBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if string(gotBody) != tt.body { + t.Fatalf("body=%q, want %q", gotBody, tt.body) + } + }) + } +} + +func TestProviderDoesNotRetryExplicitQuota429(t *testing.T) { + t.Parallel() + var requests atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + w.Header().Set("content-type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"code":"insufficient_quota","message":"You exceeded your current quota"}}`) + })) + defer upstream.Close() + + provider, err := ConfigFrom("openai", "test-model", upstream.URL, "test-key", "").NewProvider() + if err != nil { + t.Fatal(err) + } + var streamErr error + for _, err := range provider.Stream(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{llm.UserText("ping")}, + }) { + if err != nil { + streamErr = err + } + } + if streamErr == nil || !IsQuotaExhaustedMessage(streamErr.Error()) { + t.Fatalf("expected explicit quota error, got %v", streamErr) + } + if got := requests.Load(); got != 1 { + t.Fatalf("explicit quota request retried %d times, want exactly one request", got) + } +} diff --git a/agent/terminalreason.go b/agent/terminalreason.go index 90228b2..838cbd1 100644 --- a/agent/terminalreason.go +++ b/agent/terminalreason.go @@ -9,26 +9,15 @@ import ( "github.com/Autumn-27/norma/harness" ) -// This file turns a run's terminal state into something a human can act on. -// -// Before: a cancelled run produced one fixed sentence — "(运行被中断:任务暂停 / -// 终止 / 重启所致;未完成)" — a guess-list, not a diagnosis. The information to do -// better was already there (the cancel cause, the terminal reason, the turn -// count, the tool that never returned); it was just being thrown away. -// -// terminalText produces a self-contained ONE-LINE summary (the transcript's -// collapsed row, cut at the first newline by firstLine) plus a multi-line detail -// block rendered in full when the row is expanded. - -// runTrace records which tool call is in flight, so an aborted / timed-out run can -// name the step it died on — the single most useful fact when a run hangs. +// runTrace retains the latest tool call so an interrupted run can identify the +// operation that was still in flight. type runTrace struct { startedAt time.Time - id string // tool_use id of the most recent call - name string // its tool name - input string // its raw input JSON + id string + name string + input string at time.Time - pending bool // true = no tool_result arrived for it yet + pending bool } func (t *runTrace) start(id, name, input string) { @@ -41,73 +30,73 @@ func (t *runTrace) done(id string) { } } -// reasonHint explains each harness terminal reason in plain 中文: what it means, -// whether results were still written back, and what the engine does next. Covers -// every TerminalReason the harness can emit. var reasonHint = map[harness.TerminalReason]string{ - harness.ReasonCompleted: "模型主动结束了本轮(正常收场),但没有留下文字总结。事实/资产以工具调用为准,看上面的 tool_use 记录", - harness.ReasonMaxTurns: "撞到步数上限(MaxTurns):SDK 已自动跑了一轮收尾(settlement)把事实/资产写回,只是缺一段文字总结。" + - "意图标记为 exhausted——表示「这个方向试过但没做完」,规划者会据此换角度,不是失败", - harness.ReasonTimeout: "撞到单次运行的软墙钟预算(MaxDuration):SDK 在回合边界优雅收尾并写回事实/资产,只是缺文字总结。" + - "意图标记为 exhausted,不是失败", - harness.ReasonModelError: "模型/API 调用失败(网络、鉴权、限流、供应商 5xx 等)。引擎会退避后自动重跑几次;重试用尽仍失败则意图标记为 blocked", - harness.ReasonBlockingLimit: "上下文长度撞硬上限,请求在发出前就被拦下。通常是单次运行塞进了过多工具输出——考虑收窄意图粒度或压缩工具返回", - harness.ReasonPromptTooLong: "提示词过长且压缩(compaction)重试已用尽,无法继续。同样指向单次运行累积的上下文过大", - harness.ReasonImageError: "多模态内容不被当前模型支持(截图/图片块)。换支持视觉的模型,或让工具别回图", - harness.ReasonStopHookPrevented: "Stop 钩子阻止了本轮结束(守卫判定还不能收场),且随后没能继续下去。" + - "检查任务的 Guard 规则是否过严", - harness.ReasonHookStopped: "某个工具/钩子主动叫停了继续执行(守卫拦截,如越界目标、禁用命令)。看上面最后一条 tool_result 里的拦截说明", - harness.ReasonAbortedStreaming: "运行在【模型输出流式生成阶段】被取消", - harness.ReasonAbortedTools: "运行在【工具执行阶段】被取消", + harness.ReasonCompleted: "模型正常结束了本轮,但没有留下文字总结;事实和资产以本轮工具调用记录为准", + harness.ReasonMaxTurns: "达到步数上限(MaxTurns):SDK 已执行收尾并写回事实和资产,意图会标记为 exhausted,供规划者换方向继续,而不是作为失败处理", + harness.ReasonTimeout: "达到单次运行的软墙钟预算(MaxDuration):SDK 已在回合边界收尾并写回事实和资产,意图会标记为 exhausted", + harness.ReasonModelError: "模型或 API 调用失败(网络、鉴权、限流、供应商 5xx 等);引擎会按策略退避重试,重试用尽后意图标记为 blocked", + harness.ReasonBlockingLimit: "上下文长度达到硬上限,请求在发出前被拦截;应收窄意图粒度或压缩工具返回", + harness.ReasonPromptTooLong: "提示词过长且上下文压缩重试已经用尽,无法继续执行", + harness.ReasonImageError: "当前模型不支持本轮多模态内容;请切换支持视觉的模型或避免工具返回图片", + harness.ReasonStopHookPrevented: "Stop 钩子阻止本轮结束,随后未能继续;请检查任务 Guard 规则是否过严", + harness.ReasonHookStopped: "工具或钩子主动停止继续执行,例如越界目标或禁用命令;请检查最后一条 tool_result 的拦截说明", + harness.ReasonAbortedStreaming: "运行在模型输出流式生成阶段被取消", + harness.ReasonAbortedTools: "运行在工具执行阶段被取消", } -// terminalText renders a terminal event with no final text into (summary, detail). -// summary is one self-contained line; detail is the expandable breakdown. +// terminalText renders a terminal event with no final text into a compact summary +// and a Markdown detail block. func terminalText(ctx context.Context, term *harness.Terminal, tr *runTrace) (string, string) { reason := term.Reason aborted := reason == harness.ReasonAbortedStreaming || reason == harness.ReasonAbortedTools + // Prompt may return ctx.Err directly without a terminal event. Preserve the + // cancellation cause instead of falling back to an empty/unknown terminal reason. + if reason == "" && ctx.Err() != nil { + aborted = true + } - // ---- one-line summary: cause first (that's what the reader is after), then - // just enough run shape to judge it without expanding. ---- var sum string - switch { - case aborted: + if aborted { _, short, _, ok := AbortReason(ctx) if !ok { - short = "取不到取消原因(run 的 context 上没有挂 cause)" + short = "未能取得取消原因" } - stage := "工具执行阶段" - if reason == harness.ReasonAbortedStreaming { + stage := "执行过程中" + switch reason { + case harness.ReasonAbortedStreaming: stage = "模型输出阶段" + case harness.ReasonAbortedTools: + stage = "工具执行阶段" } sum = "(运行被中断:" + short + ";停在" + stage + progressSuffix(term, tr) + ",未完成)" - case reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout: - sum = "(达运行预算上限(" + string(reason) + "),已收尾写回事实" + progressSuffix(term, tr) + ";本次无文字总结)" - default: - hint := reasonHint[reason] - if hint == "" { - hint = "未知终态(harness 新增了终态但这里还没登记,请补 reasonHint)" - } - sum = "(无文字总结,终态 " + string(reason) + ":" + firstLine(hint, 80) + ")" + } else if reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout { + sum = "(达到运行预算上限(" + string(reason) + "),已收尾写回事实" + progressSuffix(term, tr) + ";本次无文字总结)" + } else { + hint := terminalReasonHint(reason) + sum = "(无文字总结,终态 " + terminalReasonLabel(reason) + ":" + firstLine(hint, 80) + ")" } - // ---- expandable detail ---- var b strings.Builder b.WriteString(sum) b.WriteString("\n\n") - fmt.Fprintf(&b, "- **终态**: `%s` — %s\n", reason, reasonHint[reason]) + displayReason := terminalReasonLabel(reason) + fmt.Fprintf(&b, "- **终态**: `%s` - %s\n", displayReason, terminalReasonHint(reason)) if aborted { code, _, why, ok := AbortReason(ctx) if ok { fmt.Fprintf(&b, "- **中断原因** (`%s`): %s\n", code, why) } else { - b.WriteString("- **中断原因**: 取不到——取消方没有用 context.WithCancelCause 挂上具名原因。" + - "这属于漏网的取消点,请补 agent/cancelcause.go 里的枚举\n") + b.WriteString("- **中断原因**: 无法取得;取消方可能没有通过 context.WithCancelCause 附加具名原因\n") } } if term.Err != nil { fmt.Fprintf(&b, "- **底层错误**: `%v`\n", term.Err) } + if aborted && strings.TrimSpace(term.Text) != "" { + b.WriteString("- **取消前已生成的部分输出**:\n\n") + b.WriteString(term.Text) + b.WriteString("\n\n") + } if term.Turns > 0 { fmt.Fprintf(&b, "- **已执行**: %d 轮模型回合\n", term.Turns) } @@ -118,21 +107,34 @@ func terminalText(ctx context.Context, term *harness.Terminal, tr *runTrace) (st fmt.Fprintf(&b, "- **累计 token**: 输入 %d / 输出 %d / 缓存读 %d / 缓存写 %d\n", u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheWriteTokens) } - if tr.name != "" { - if tr.pending { - fmt.Fprintf(&b, "- **中断时正在执行的工具**: `%s`(已跑 %s,**未返回结果**) — 若是硬超时/长时间无进展,问题多半就在这一步\n ```\n %s\n ```\n", - tr.name, roundDur(time.Since(tr.at)), firstLine(tr.input, 300)) - } else { - fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`(已正常返回)\n", tr.name) - } + if tr.name == "" { + b.WriteString("- **工具调用**: 本次运行还没有发出工具调用就结束了\n") + } else if tr.pending { + fmt.Fprintf(&b, "- **中断时正在执行的工具**: `%s`(已运行 %s,**未返回结果**)\n\n ```json\n %s\n ```\n", + tr.name, roundDur(time.Since(tr.at)), firstLine(tr.input, 300)) } else { - b.WriteString("- **工具调用**: 本次运行还没发出任何工具调用就结束了\n") + fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`(已正常返回)\n", tr.name) } return sum, b.String() } -// progressSuffix adds "已跑 N 轮 / 12m30s" to a summary line — how far the run got -// before it died, which decides whether an abort cost anything worth re-running. +func terminalReasonLabel(reason harness.TerminalReason) string { + if reason == "" { + return "context_canceled" + } + return string(reason) +} + +func terminalReasonHint(reason harness.TerminalReason) string { + if hint := reasonHint[reason]; hint != "" { + return hint + } + if reason == "" { + return "运行的 context 已取消,但底层没有产生 Terminal 事件" + } + return "未知终态;harness 可能新增了 TerminalReason,请补充 reasonHint" +} + func progressSuffix(term *harness.Terminal, tr *runTrace) string { var parts []string if term.Turns > 0 { @@ -144,10 +146,9 @@ func progressSuffix(term *harness.Terminal, tr *runTrace) string { if len(parts) == 0 { return "" } - return ",已跑 " + strings.Join(parts, " / ") + return ",已运行 " + strings.Join(parts, " / ") } -// roundDur formats a duration at a readable granularity (no 3.400213s noise). func roundDur(d time.Duration) string { switch { case d < time.Minute: diff --git a/agent/terminalreason_test.go b/agent/terminalreason_test.go index a3ecd20..ff53fe4 100644 --- a/agent/terminalreason_test.go +++ b/agent/terminalreason_test.go @@ -12,8 +12,6 @@ import ( "github.com/Autumn-27/norma/llm" ) -// every TerminalReason the harness can emit must have an explanation — otherwise -// a run silently falls back to "未知终态", which is the bug this file exists to fix. func TestReasonHintCoversEveryTerminalReason(t *testing.T) { all := []harness.TerminalReason{ harness.ReasonCompleted, harness.ReasonBlockingLimit, harness.ReasonImageError, @@ -21,152 +19,151 @@ func TestReasonHintCoversEveryTerminalReason(t *testing.T) { harness.ReasonPromptTooLong, harness.ReasonStopHookPrevented, harness.ReasonHookStopped, harness.ReasonMaxTurns, harness.ReasonTimeout, } - for _, r := range all { - if strings.TrimSpace(reasonHint[r]) == "" { - t.Errorf("terminal reason %q 没有中文说明", r) + for _, reason := range all { + if strings.TrimSpace(reasonHint[reason]) == "" { + t.Errorf("terminal reason %q has no explanation", reason) } } } -// the cause attached at cancel time must survive the ctx chain the engine builds: -// task exec ctx -> per-work child -> value ctx -> the run's hard-timeout backstop. func TestAbortCausePropagatesThroughRunContextChain(t *testing.T) { type key struct{} - exec, cancelExec := context.WithCancelCause(context.Background()) - work, cancelWork := context.WithCancelCause(exec) + execCtx, cancelExec := context.WithCancelCause(context.Background()) + workCtx, cancelWork := context.WithCancelCause(execCtx) defer cancelWork(nil) - valued := context.WithValue(work, key{}, "task-1") - run, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) + valued := context.WithValue(workCtx, key{}, "task-1") + runCtx, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) defer cancelRun() cancelExec(AbortPausedByUser) - - code, _, text, ok := AbortReason(run) + code, _, text, ok := AbortReason(runCtx) if !ok || code != "paused_by_user" { - t.Fatalf("code=%q ok=%v,期望 paused_by_user", code, ok) + t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok) } - if !strings.Contains(text, "退回 frontier") { - t.Errorf("原因文案没传下来: %q", text) + if !strings.Contains(text, "frontier") { + t.Fatalf("cause detail did not propagate: %q", text) } } func TestAbortReasonFallbacks(t *testing.T) { if _, _, _, ok := AbortReason(context.Background()); ok { - t.Error("未取消的 ctx 不应报出中断原因") + t.Fatal("live context must not report an abort reason") } - // a cancel site that forgot to attach a cause must be called out, not guessed at plain, cancel := context.WithCancel(context.Background()) cancel() - code, _, _, ok := AbortReason(plain) - if !ok || code != "canceled_no_cause" { - t.Errorf("code=%q ok=%v,期望 canceled_no_cause", code, ok) + if code, _, _, ok := AbortReason(plain); !ok || code != "canceled_no_cause" { + t.Fatalf("code=%q ok=%v, want canceled_no_cause", code, ok) } - - timed, cancelT := context.WithTimeout(context.Background(), time.Nanosecond) - defer cancelT() + timed, cancelTimed := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancelTimed() <-timed.Done() if code, _, _, _ := AbortReason(timed); code != "deadline_exceeded" { - t.Errorf("code=%q,期望 deadline_exceeded", code) + t.Fatalf("code=%q, want deadline_exceeded", code) } } -// an aborted run's one-line summary must name the actual cause, and the detail -// must pin down the tool that never came back. func TestTerminalTextAbortedNamesCauseAndHangingTool(t *testing.T) { ctx, cancel := context.WithCancelCause(context.Background()) cancel(AbortKilledByPlanner) - - tr := &runTrace{startedAt: time.Now().Add(-90 * time.Second)} - tr.start("tu_1", "Bash", `{"command":"nmap -p- 10.0.0.1"}`) - + trace := &runTrace{startedAt: time.Now().Add(-90 * time.Second)} + trace.start("tu_1", "Bash", `{"command":"nmap -p- 10.0.0.1"}`) term := &harness.Terminal{ Reason: harness.ReasonAbortedTools, Err: context.Canceled, Turns: 7, Usage: llm.Usage{InputTokens: 1200, OutputTokens: 340}, } - sum, detail := terminalText(ctx, term, tr) - - if !strings.Contains(sum, "kill_work") { - t.Errorf("摘要没点出取消方: %q", sum) - } - if strings.Contains(sum, "\n") { - t.Errorf("摘要必须是单行(前端按首行截断): %q", sum) + summary, detail := terminalText(ctx, term, trace) + if !strings.Contains(summary, "规划者") || strings.Contains(summary, "\n") { + t.Fatalf("unexpected summary: %q", summary) } for _, want := range []string{"killed_by_planner", "aborted_tools", "7 轮", "Bash", "未返回结果", "1200"} { if !strings.Contains(detail, want) { - t.Errorf("详情缺少 %q:\n%s", want, detail) + t.Errorf("detail missing %q:\n%s", want, detail) } } } -// a tool that DID return must not be reported as the hang site. +func TestTerminalTextDirectContextCancellationKeepsCause(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortChatStoppedByUser) + summary, detail := terminalText(ctx, &harness.Terminal{Err: context.Canceled}, &runTrace{startedAt: time.Now()}) + if !strings.Contains(summary, "用户停止") || !strings.Contains(detail, "chat_stopped_by_user") { + t.Fatalf("direct context cancellation lost cause: %q\n%s", summary, detail) + } + if !strings.Contains(detail, "context_canceled") { + t.Fatalf("missing synthesized terminal label: %s", detail) + } +} + +func TestTerminalTextAbortedPreservesPartialOutput(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortPausedByUser) + summary, detail := terminalText(ctx, &harness.Terminal{ + Reason: harness.ReasonAbortedStreaming, + Text: "已经生成的半段回答", + }, &runTrace{startedAt: time.Now()}) + if !strings.Contains(summary, "用户暂停") { + t.Fatalf("abort summary lost cause: %q", summary) + } + if !strings.Contains(detail, "已经生成的半段回答") { + t.Fatalf("abort detail lost partial output: %s", detail) + } +} + func TestTerminalTextCompletedToolNotBlamed(t *testing.T) { ctx, cancel := context.WithCancelCause(context.Background()) cancel(AbortShutdown) - tr := &runTrace{startedAt: time.Now()} - tr.start("tu_1", "Read", `{"path":"/etc/hosts"}`) - tr.done("tu_1") - - _, detail := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonAbortedStreaming}, tr) - if strings.Contains(detail, "未返回结果") { - t.Errorf("已返回的工具被误报为卡住:\n%s", detail) - } - if !strings.Contains(detail, "已正常返回") { - t.Errorf("详情没记录最后一个工具:\n%s", detail) + trace := &runTrace{startedAt: time.Now()} + trace.start("tu_1", "Read", `{"path":"/etc/hosts"}`) + trace.done("tu_1") + _, detail := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonAbortedStreaming}, trace) + if strings.Contains(detail, "未返回结果") || !strings.Contains(detail, "已正常返回") { + t.Fatalf("completed tool was blamed:\n%s", detail) } } -// non-abort terminals keep their own explanation instead of the interruption text. func TestTerminalTextNonAbortReasons(t *testing.T) { - ctx := context.Background() - tr := &runTrace{startedAt: time.Now()} - - sum, _ := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonMaxTurns}, tr) - if !strings.Contains(sum, "达运行预算上限") || strings.Contains(sum, "中断") { - t.Errorf("max_turns 摘要不对: %q", sum) + trace := &runTrace{startedAt: time.Now()} + summary, _ := terminalText(context.Background(), &harness.Terminal{Reason: harness.ReasonMaxTurns}, trace) + if !strings.Contains(summary, "运行预算上限") || strings.Contains(summary, "中断") { + t.Fatalf("unexpected max_turns summary: %q", summary) } - - sum, detail := terminalText(ctx, &harness.Terminal{ - Reason: harness.ReasonModelError, Err: errors.New("429 rate limited")}, tr) - if !strings.Contains(sum, "model_error") { - t.Errorf("model_error 摘要不对: %q", sum) - } - if !strings.Contains(detail, "429 rate limited") { - t.Errorf("详情没带上底层错误:\n%s", detail) + summary, detail := terminalText(context.Background(), &harness.Terminal{ + Reason: harness.ReasonModelError, Err: errors.New("429 rate limited"), + }, trace) + if !strings.Contains(summary, "model_error") || !strings.Contains(detail, "429 rate limited") { + t.Fatalf("model_error detail incomplete: %q\n%s", summary, detail) } } -// every enumerated cause must carry both a compact headline (for the collapsed -// row) and a full explanation (for the expanded detail). func TestAbortCausesAreWellFormed(t *testing.T) { all := []*AbortCause{ AbortPausedByUser, AbortPausedByOrchestrator, AbortTaskDeleted, AbortPausedOnReload, - AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkFinished, - AbortPausedRaceGuard, AbortChatStoppedByUser, AbortChatTurnFinished, + AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkPausedByUser, + AbortWorkCancelledByUser, AbortWorkFinished, AbortPausedRaceGuard, + AbortChatStoppedByUser, AbortChatPausedWithTask, AbortChatTurnFinished, AbortShutdown, AbortRunHardTimeout, } seen := map[string]bool{} - for _, c := range all { + for _, abort := range all { switch { - case c.Code == "" || seen[c.Code]: - t.Errorf("code 缺失或重复: %q", c.Code) - case c.Short == "" || utf8.RuneCountInString(c.Short) > 40: - t.Errorf("%s: Short 必须存在且 ≤40 字(进折叠摘要行): %q", c.Code, c.Short) - case utf8.RuneCountInString(c.Text) <= utf8.RuneCountInString(c.Short): - t.Errorf("%s: Text 应比 Short 更完整", c.Code) + case abort.Code == "" || seen[abort.Code]: + t.Errorf("missing or duplicate code: %q", abort.Code) + case abort.Short == "" || utf8.RuneCountInString(abort.Short) > 40: + t.Errorf("%s has invalid short text: %q", abort.Code, abort.Short) + case utf8.RuneCountInString(abort.Text) <= utf8.RuneCountInString(abort.Short): + t.Errorf("%s detail must be longer than short text", abort.Code) } - seen[c.Code] = true + seen[abort.Code] = true } } -// summaries are mostly 中文; capping by bytes used to slice a rune in half. func TestFirstLineCapsByRuneNotByte(t *testing.T) { - got := firstLine(strings.Repeat("恢复", 10), 5) - if got != "恢复恢复恢…" { - t.Errorf("firstLine 按字节截断了: %q", got) + if got := firstLine(strings.Repeat("恢复", 10), 5); got != "恢复恢复恢…" { + t.Fatalf("firstLine split a Unicode character: %q", got) } if got := firstLine("头\n尾", 100); got != "头" { - t.Errorf("多行没截到首行: %q", got) + t.Fatalf("firstLine did not stop at newline: %q", got) } } diff --git a/agent/tools.go b/agent/tools.go index 93b71b7..d0958a2 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -23,6 +23,10 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[ var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "summary": p["summary"], "state": n.State} + if n.Inherited { + m["source_task_id"] = n.SourceTaskID + m["inherited"] = true + } // asset_ids is the structured "which assets this direction covers" signal for // dedup; fall back to legacy payload keys (target_ids plural, then target_id // single) so intents stored before the rename still surface their anchors. @@ -265,132 +269,438 @@ func (t *ToolSet) graphOverview() actool.CoreTool { // the model needn't spend a turn calling the tool — every plan round starts with // an empty context and always needs this first). func (t *ToolSet) graphOverviewData() map[string]any { - { - out := map[string]any{} - // goals summary folded in so the planner needn't call list_goals each round. - goals, _ := t.ts.ListByKind(db.KindGoal, 100) - gsum := make([]map[string]any, 0, len(goals)) - for _, g := range goals { - var p map[string]any - _ = json.Unmarshal(g.Payload, &p) - gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + out := map[string]any{} + // goals summary folded in so the planner needn't call list_goals each round. + goals, _ := t.ts.ListByKind(db.KindGoal, 100) + gsum := make([]map[string]any, 0, len(goals)) + for _, g := range goals { + var p map[string]any + _ = json.Unmarshal(g.Payload, &p) + gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + } + out["goals"] = gsum + // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示;folded in so the + // planner reads them every round when generating intents (否则只写不读). + hints, _ := t.ts.ListByKind(db.KindHint, 50) + hsum := make([]map[string]any, 0, len(hints)) + for _, h := range hints { + var p map[string]any + _ = json.Unmarshal(h.Payload, &p) + hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + out["hints"] = hsum + // lineage from the exploration edges: an intent's parents (what it + // derived_from — possibly several facts combined) and its yields (the + // facts/findings it produced). factFrom maps a fact → the intent that + // produced it. This is the relationship layer the flat lists lacked. + edges, _ := t.ts.Edges(5000) + parentsOf := map[int64][]int64{} + yieldsOf := map[int64][]int64{} + factFrom := map[int64]int64{} + for _, e := range edges { + switch e.Rel { + case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent) + parentsOf[e.To] = append(parentsOf[e.To], e.From) + case db.RelYields: // intent --yields--> fact/finding + yieldsOf[e.From] = append(yieldsOf[e.From], e.To) + factFrom[e.To] = e.From } - out["goals"] = gsum - // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示;folded in so the - // planner reads them every round when generating intents (否则只写不读). - hints, _ := t.ts.ListByKind(db.KindHint, 50) - hsum := make([]map[string]any, 0, len(hints)) - for _, h := range hints { - var p map[string]any - _ = json.Unmarshal(h.Payload, &p) - hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + fr, _ := t.ts.Frontier(100) + out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) + all, _ := t.ts.ListByKind(db.KindIntent, 300) + var running, recentDone []*db.Node + for _, n := range all { + switch n.State { + case "running": + running = append(running, n) + case "done", "blocked", "exhausted": + if len(recentDone) < 15 { + recentDone = append(recentDone, n) + } } - out["hints"] = hsum - // lineage from the exploration edges: an intent's parents (what it - // derived_from — possibly several facts combined) and its yields (the - // facts/findings it produced). factFrom maps a fact → the intent that - // produced it. This is the relationship layer the flat lists lacked. - edges, _ := t.ts.Edges(5000) + } + out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) + out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) + out["frontier_open"] = len(fr) + // findings (confirmed vulns) and facts (worker exploration results) are + // now distinct node kinds. recent_facts surfaces fact summaries (esp. + // negative results) so the planner sees them in one call; full content + // via node_detail(id). + vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000) + factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first + out["findings"] = len(vulnNodes) // 确认漏洞数(目标判定看它) + out["facts"] = len(factNodes) // 探索事实/结论数(含否定结论) + recentFacts := make([]map[string]any, 0, 20) + for _, n := range factNodes { + if len(recentFacts) >= 20 { + break + } + m := compactNode(n) + if from := factFrom[n.ID]; from > 0 { + m["from_intent"] = from // 本事实由哪个意图产生 + } + // confidence 带进概览:让规划者一眼看出哪条结论只是 inferred(尤其否定结论 + // 别当铁案);evidence 较长,留给 node_detail(id)。 + var fp map[string]any + if json.Unmarshal(n.Payload, &fp) == nil { + if c, ok := fp["confidence"].(string); ok && c != "" { + m["confidence"] = c + } + } + recentFacts = append(recentFacts, m) + } + out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?},详情/证据用 node_detail(id) + // the original task (root) so the planner always has it, not just the + // decomposed goals. + if description, goal, err := t.ts.Root(); err == nil { + out["task"] = map[string]any{"description": description, "goal": goal} + } + // Direct source tasks are a live, read-only blackboard view. Keep their + // summaries in a separate field so their intents never enter this task's + // frontier or get mistaken for locally claimable work. + out["related_tasks"] = t.relatedTaskOverviews() + // coverage:粗略的资产测试覆盖度参考——范围(task_scope)内的资产里,被 fact 碰过的 + // 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。 + if t.as != nil && t.ts != nil && t.taskID > 0 { + if cov, err := t.as.TaskCoverageWithSources(t.taskID); err == nil { + m := map[string]any{ + "denominator": cov.Denominator, + "tested": cov.Tested, + "by_type": cov.ByType, + "note": "coverage资产测试覆盖度(包括接口等各种相关资产),粗略估计、仅供参考:包含当前任务与直接关联任务的 scope、事实锚点;关联 scope 只读。容器型资产/大量枚举会让它偏低,勿据此认为已测完;可用 add_task_scope 增补本任务范围、list_untested_assets 看未测资产【通常不调用list_untested_assets,按照任务推进即可】;", + } + if cov.Denominator == 0 { + m["pct"] = nil + m["status"] = "范围未锚定" + } else { + m["pct"] = cov.Pct + } + // scope:当前测试范围的根资产(task_scope 原始行),让 agent 知道这个任务到底 + // 圈定了哪些目标(不是全部测试资产,而是范围边界本身)。 + if rows, err := t.as.ListTaskScopeWithSources(t.taskID); err == nil && len(rows) > 0 { + scope := make([]map[string]any, 0, len(rows)) + for _, r := range rows { + e := map[string]any{"kind": r.Kind, "source": r.Source, "task_id": r.TaskID} + if r.TaskID != t.taskID { + inheritedMap(e, r.TaskID) + } + switch { + case r.Domain != "": + e["value"] = r.Domain + case r.Net != "": + e["value"] = r.Net + case r.CompanyID != nil: + e["company_id"] = *r.CompanyID + } + scope = append(scope, e) + } + m["scope"] = scope + } + if hosts, err := t.as.HostsByTaskWithSources(t.taskID); err == nil { + const hostContextLimit = 500 + visible := hosts + if len(visible) > hostContextLimit { + visible = visible[:hostContextLimit] + } + m["hosts"] = visible + m["host_count"] = len(hosts) + if len(visible) < len(hosts) { + m["hosts_truncated"] = true + } + } + out["coverage"] = m + } + } + return out +} + +func inheritedMap(m map[string]any, sourceTaskID int64) map[string]any { + m["source_task_id"] = sourceTaskID + m["inherited"] = true + return m +} + +const ( + relatedOverviewTotalTextRunes = 48_000 + relatedOverviewMaxTextPerSource = 8_000 + relatedOverviewMaxGoalsPerSource = 8 + relatedOverviewMaxHintsPerSource = 6 + relatedOverviewMaxFactsPerSource = 12 + relatedOverviewMaxFindingsPerTask = 6 + relatedOverviewMaxIntentsPerTask = 8 + relatedOverviewMaxScopePerSource = 12 +) + +// overviewTextBudget bounds inherited prompt text while preserving a fair slice +// for every direct source. Full evidence remains available through the on-demand +// read tools, so truncation here does not discard persisted blackboard data. +type overviewTextBudget struct { + remaining int + truncated bool +} + +func relatedOverviewBudgetForSources(sourceCount int) int { + if sourceCount <= 0 { + return 0 + } + if sourceCount > db.MaxTaskSourceCount { + sourceCount = db.MaxTaskSourceCount + } + perSource := relatedOverviewTotalTextRunes / sourceCount + if perSource > relatedOverviewMaxTextPerSource { + perSource = relatedOverviewMaxTextPerSource + } + return perSource +} + +func (b *overviewTextBudget) take(value any, fieldLimit int) string { + var text string + switch value := value.(type) { + case string: + text = strings.TrimSpace(value) + case nil: + return "" + default: + text = strings.TrimSpace(fmt.Sprint(value)) + } + if text == "" { + return "" + } + if b.remaining <= 0 || fieldLimit <= 0 { + b.truncated = true + return "" + } + runes := []rune(text) + limit := fieldLimit + if limit > b.remaining { + limit = b.remaining + } + if len(runes) > limit { + b.truncated = true + if limit == 1 { + text = "…" + } else { + text = string(runes[:limit-1]) + "…" + } + runes = []rune(text) + } + b.remaining -= len(runes) + return text +} + +func recentTerminalIntents(store *db.ExplorationStore, limit int) []*db.Node { + if limit <= 0 { + return []*db.Node{} + } + const batch = 300 + cursor := int64(0) + out := make([]*db.Node, 0, limit) + for len(out) < limit { + page, more, err := store.ListByKindPage(db.KindIntent, cursor, batch) + if err != nil || len(page) == 0 { + break + } + for _, intent := range page { + switch intent.State { + case "done", "blocked", "exhausted", "stopped": + out = append(out, intent) + } + if len(out) >= limit { + break + } + } + if !more { + break + } + cursor = page[len(page)-1].ID + } + return out +} + +// relatedTaskOverviews distills persistent blackboard state from direct source +// tasks. It intentionally reads each source's local store methods, never its own +// related sources, so inheritance is one level only. +func (t *ToolSet) relatedTaskOverviews() []map[string]any { + sources, err := t.ts.DirectSourceStores() + if err != nil { + return []map[string]any{} + } + if len(sources) > db.MaxTaskSourceCount { + sources = sources[:db.MaxTaskSourceCount] + } + perSourceTextBudget := relatedOverviewBudgetForSources(len(sources)) + out := make([]map[string]any, 0, len(sources)) + for _, source := range sources { + ts := source.Store + budget := overviewTextBudget{remaining: perSourceTextBudget} + item := map[string]any{ + "source_task_id": source.Task.TaskID, + "inherited": true, + "task": map[string]any{ + "description": budget.take(source.Task.Description, 800), + "goal": budget.take(source.Task.Goal, 800), + "status": source.Task.Status, + }, + } + stats, statsErr := ts.Stats() + + edges, _ := ts.Edges(5000) parentsOf := map[int64][]int64{} yieldsOf := map[int64][]int64{} factFrom := map[int64]int64{} - for _, e := range edges { - switch e.Rel { - case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent) - parentsOf[e.To] = append(parentsOf[e.To], e.From) - case db.RelYields: // intent --yields--> fact/finding - yieldsOf[e.From] = append(yieldsOf[e.From], e.To) - factFrom[e.To] = e.From + for _, edge := range edges { + switch edge.Rel { + case db.RelDerivedFrom, db.RelSpawns: + parentsOf[edge.To] = append(parentsOf[edge.To], edge.From) + case db.RelYields: + yieldsOf[edge.From] = append(yieldsOf[edge.From], edge.To) + factFrom[edge.To] = edge.From } } - fr, _ := t.ts.Frontier(100) - out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) - all, _ := t.ts.ListByKind(db.KindIntent, 300) - var running, recentDone []*db.Node - for _, n := range all { - switch n.State { - case "running": - running = append(running, n) - case "done", "blocked", "exhausted": - if len(recentDone) < 15 { - recentDone = append(recentDone, n) - } + + goals, _ := ts.ListByKind(db.KindGoal, relatedOverviewMaxGoalsPerSource) + goalSummary := make([]map[string]any, 0, len(goals)) + for _, goal := range goals { + var payload map[string]any + _ = json.Unmarshal(goal.Payload, &payload) + goalSummary = append(goalSummary, inheritedMap(map[string]any{ + "id": goal.ID, "state": goal.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["goals"] = goalSummary + + hints, _ := ts.ListByKind(db.KindHint, relatedOverviewMaxHintsPerSource) + hintSummary := make([]map[string]any, 0, len(hints)) + for _, hint := range hints { + var payload map[string]any + _ = json.Unmarshal(hint.Payload, &payload) + hintSummary = append(hintSummary, inheritedMap(map[string]any{ + "id": hint.ID, "state": hint.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["hints"] = hintSummary + + facts, _ := ts.ListByKind(db.KindFact, relatedOverviewMaxFactsPerSource) + findings, _ := ts.ListByKind(db.KindFinding, relatedOverviewMaxFindingsPerTask) + intentNodes, _ := ts.ListByKind(db.KindIntent, 300) + terminalIntent := make(map[int64]bool, len(intentNodes)) + for _, intent := range intentNodes { + terminalIntent[intent.ID] = inheritedIntentSummaryState(intent.State) + } + item["facts"] = len(facts) + item["findings"] = len(findings) + if statsErr == nil { + item["facts"] = stats[db.KindFact] + item["findings"] = stats[db.KindFinding] + if stats[db.KindGoal] > len(goals) || stats[db.KindHint] > len(hints) || + stats[db.KindFact] > len(facts) || stats[db.KindFinding] > len(findings) { + budget.truncated = true } } - out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) - out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) - out["frontier_open"] = len(fr) - // findings (confirmed vulns) and facts (worker exploration results) are - // now distinct node kinds. recent_facts surfaces fact summaries (esp. - // negative results) so the planner sees them in one call; full content - // via node_detail(id). - vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000) - factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first - out["findings"] = len(vulnNodes) // 确认漏洞数(目标判定看它) - out["facts"] = len(factNodes) // 探索事实/结论数(含否定结论) - recentFacts := make([]map[string]any, 0, 20) - for _, n := range factNodes { - if len(recentFacts) >= 20 { - break + recentFindings := make([]map[string]any, 0, len(findings)) + for _, finding := range findings { + entry := inheritedMap(compactFinding(finding), source.Task.TaskID) + entry["summary"] = budget.take(entry["summary"], 400) + recentFindings = append(recentFindings, entry) + } + item["recent_findings"] = recentFindings + recentFacts := make([]map[string]any, 0, len(facts)) + for _, fact := range facts { + m := inheritedMap(compactNode(fact), source.Task.TaskID) + m["summary"] = budget.take(m["summary"], 400) + if from := factFrom[fact.ID]; from > 0 && terminalIntent[from] { + m["from_intent"] = from } - m := compactNode(n) - if from := factFrom[n.ID]; from > 0 { - m["from_intent"] = from // 本事实由哪个意图产生 - } - // confidence 带进概览:让规划者一眼看出哪条结论只是 inferred(尤其否定结论 - // 别当铁案);evidence 较长,留给 node_detail(id)。 - var fp map[string]any - if json.Unmarshal(n.Payload, &fp) == nil { - if c, ok := fp["confidence"].(string); ok && c != "" { - m["confidence"] = c + var payload map[string]any + if json.Unmarshal(fact.Payload, &payload) == nil { + if confidence, ok := payload["confidence"].(string); ok && confidence != "" { + m["confidence"] = confidence } } recentFacts = append(recentFacts, m) } - out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?},详情/证据用 node_detail(id) - // the original task (root) so the planner always has it, not just the - // decomposed goals. - if description, goal, err := t.ts.Root(); err == nil { - out["task"] = map[string]any{"description": description, "goal": goal} + item["recent_facts"] = recentFacts + + recentDone := recentTerminalIntents(ts, relatedOverviewMaxIntentsPerTask) + for _, intent := range recentDone { + intent.Inherited = true + intent.SourceTaskID = source.Task.TaskID } - // coverage:粗略的资产测试覆盖度参考——范围(task_scope)内的资产里,被 fact 碰过的 - // 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。 - if t.as != nil && t.ts != nil && t.taskID > 0 { - if cov, err := t.as.TaskCoverage(t.taskID, t.ts.ID()); err == nil { - m := map[string]any{ - "denominator": cov.Denominator, - "tested": cov.Tested, - "by_type": cov.ByType, - "note": "coverage资产测试覆盖度(包括接口等各种相关资产),粗略估计、仅供参考:容器型资产/大量枚举会让它偏低,勿据此认为已测完;scope 是当前任务测试范围的根资产(root_domain/subdomain/ip/cidr/company 原始行),即覆盖度分母的边界、本任务授权/圈定的目标本身;可用 add_task_scope 增补范围、list_untested_assets 看未测的具体资产【通常不调用list_untested_assets,按照任务推进即可】;", + intentResults := compactIntents(recentDone, parentsOf, yieldsOf) + for i, intent := range recentDone { + intentResults[i]["summary"] = budget.take(intentResults[i]["summary"], 400) + acts, _, err := ts.ActivityPageForTerminalIntent(intent.ID, 0, 20) + if err != nil { + continue + } + var resultSummary, textFallback string + for _, activity := range acts { + switch activity.Kind { + case "result": + resultSummary = activity.Summary + case "text": + textFallback = activity.Summary } - if cov.ScopeRows == 0 { - m["pct"] = nil - m["status"] = "范围未锚定" - } else { - m["pct"] = cov.Pct - } - // scope:当前测试范围的根资产(task_scope 原始行),让 agent 知道这个任务到底 - // 圈定了哪些目标(不是全部测试资产,而是范围边界本身)。 - if rows, err := t.as.ListTaskScope(t.taskID); err == nil && len(rows) > 0 { - scope := make([]map[string]any, 0, len(rows)) - for _, r := range rows { - e := map[string]any{"kind": r.Kind, "source": r.Source} - switch { - case r.Domain != "": - e["value"] = r.Domain - case r.Net != "": - e["value"] = r.Net - case r.CompanyID != nil: - e["company_id"] = *r.CompanyID - } - scope = append(scope, e) - } - m["scope"] = scope - } - out["coverage"] = m + } + if resultSummary == "" { + resultSummary = textFallback + } + if resultSummary != "" { + intentResults[i]["result_summary"] = budget.take(resultSummary, 800) } } - return out + item["recent_intent_results"] = intentResults + if statsErr == nil { + item["node_stats"] = stats + } + + if t.as != nil { + if scopeRows, err := t.as.ListTaskScope(source.Task.TaskID); err == nil && len(scopeRows) > 0 { + scopeCount := len(scopeRows) + if len(scopeRows) > relatedOverviewMaxScopePerSource { + scopeRows = scopeRows[:relatedOverviewMaxScopePerSource] + budget.truncated = true + } + scope := make([]map[string]any, 0, len(scopeRows)) + for _, row := range scopeRows { + entry := map[string]any{"kind": row.Kind, "source": budget.take(row.Source, 300)} + switch { + case row.Domain != "": + entry["value"] = budget.take(row.Domain, 400) + case row.Net != "": + entry["value"] = budget.take(row.Net, 400) + case row.CompanyID != nil: + entry["company_id"] = *row.CompanyID + } + scope = append(scope, entry) + } + item["asset_scope"] = scope + item["asset_scope_count"] = scopeCount + } + if coverage, err := t.as.TaskCoverage(source.Task.TaskID, source.Task.ExplorationID); err == nil { + item["asset_coverage"] = map[string]any{ + "denominator": coverage.Denominator, + "tested": coverage.Tested, + "pct": coverage.Pct, + "by_type": coverage.ByType, + } + } + } + if budget.truncated { + item["summary_truncated"] = true + } + out = append(out, item) + } + return out +} + +func inheritedIntentSummaryState(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false } } @@ -399,7 +709,11 @@ func (t *ToolSet) graphOverviewData() map[string]any { func compactNode(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) - return map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } + return m } // compactFinding is compactNode plus the vuln-specific vulnclass/severity. @@ -407,6 +721,9 @@ func compactFinding(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } if vc, ok := p["vulnclass"]; ok && vc != nil && vc != "" { m["vulnclass"] = vc } @@ -417,16 +734,23 @@ func compactFinding(n *db.Node) map[string]any { } func (t *ToolSet) listFindings() actool.CoreTool { - return readTool("list_findings", "列【确认漏洞】(紧凑:id+task_id+intent_id+vulnclass+severity+摘要+状态)。task_id=该漏洞所属任务, intent_id=产生它的意图。这里只有漏洞,不含普通探索事实(那是 list_facts)。详情用 node_detail(id)。", + return readTool("list_findings", "列本任务及直接关联任务的【确认漏洞】(紧凑:id+task_id+intent_id+vulnclass+severity+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。这里只含漏洞;普通探索事实用 list_facts,详情用 node_detail(id)。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFinding, 500) - intentOf, _ := t.ts.FindingIntents() // finding id -> 产生它的 intent id - taskID := t.ts.ID() // exploration id = 任务 id(本 store 内所有 finding 同属) + f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500) + intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id + taskID := t.taskID + if taskID <= 0 { + taskID, _ = t.ts.TaskID() + } out := make([]map[string]any, 0, len(f)) for _, n := range f { m := compactFinding(n) - m["task_id"] = taskID + if n.Inherited { + m["task_id"] = n.SourceTaskID + } else { + m["task_id"] = taskID + } if iid, ok := intentOf[n.ID]; ok { m["intent_id"] = iid } @@ -437,10 +761,10 @@ func (t *ToolSet) listFindings() actool.CoreTool { } func (t *ToolSet) listFacts() actool.CoreTool { - return readTool("list_facts", "列【探索事实/结论】(紧凑:id+摘要+状态),即 worker 按意图探出的结论(含'端口关闭/不可注入'等否定结论)。详情用 node_detail(id)。漏洞看 list_findings。", + return readTool("list_facts", "列本任务及直接关联任务的【探索事实/结论】(紧凑:id+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。详情用 node_detail(id),漏洞看 list_findings。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFact, 500) + f, _ := t.ts.ListByKindWithSources(db.KindFact, 500) out := make([]map[string]any, 0, len(f)) for _, n := range f { out = append(out, compactNode(n)) @@ -450,7 +774,7 @@ func (t *ToolSet) listFacts() actool.CoreTool { } func (t *ToolSet) nodeDetail() actool.CoreTool { - return readTool("node_detail", "按 id 取某个【探索图节点】的完整内容(事实/发现/意图/目标:摘要 + 详情/证据/PoC)。仅限探索图节点 id(来自 list_facts / list_findings / graph_overview 的 recent_facts 返回的 id)。查【资产】请用 list_assets / asset_neighbors——资产 id 不是探索节点 id,不要传进来。", + return readTool("node_detail", "按 id 取本任务或直接关联任务的【探索图节点】完整内容。继承节点带 source_task_id/inherited=true 且只读。仅限 list_facts/list_findings/graph_overview 返回的探索节点 id;资产请用 list_assets/asset_neighbors。", obj(map[string]any{"id": idp("探索图节点 id(非资产 id)")}, "id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -461,7 +785,7 @@ func (t *ToolSet) nodeDetail() actool.CoreTool { if id <= 0 { return actool.Errorf("id 必填"), nil } - n, err := t.ts.GetNode(id) + n, err := t.ts.GetNodeWithSources(id) if err != nil { return actool.Errorf(err.Error()), nil } @@ -493,9 +817,9 @@ func (t *ToolSet) addOneIntent(it intentItem) (int64, error) { // 先校验锚点(建节点前,避免坏锚点留下孤儿意图)。 parents := pidList(it.ParentIDs) for _, pidv := range parents { - n, err := t.ts.GetNode(pidv) + n, err := t.ts.GetNodeWithSources(pidv) if err != nil || n == nil { - return 0, fmt.Errorf("parent_id %d 不存在:parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id;顶层全新方向请留空 parent_ids", pidv) + return 0, fmt.Errorf("parent_id %d 不存在于本任务或直接关联任务:parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id;顶层全新方向请留空 parent_ids", pidv) } if n.Kind != db.KindFact && n.Kind != db.KindFinding { return 0, fmt.Errorf("parent_id %d 是 %q 节点,不能作为意图锚点:意图只能锚在已确认的【事实(fact)/发现(finding)】上,不能挂在意图/目标/提示上;顶层全新方向请留空 parent_ids", pidv, n.Kind) @@ -604,6 +928,14 @@ func (t *ToolSet) proveGoal() actool.CoreTool { if goal == 0 || ev == 0 { return actool.Errorf("goal_id 和 evidence_id 必填"), nil } + goalNode, err := t.ts.GetNode(goal) + if err != nil || goalNode == nil || goalNode.Kind != db.KindGoal { + return actool.Errorf("goal_id 必须是本任务的目标节点(关联任务目标只读)"), nil + } + evidenceNode, err := t.ts.GetNodeWithSources(ev) + if err != nil || evidenceNode == nil || (evidenceNode.Kind != db.KindFact && evidenceNode.Kind != db.KindFinding) { + return actool.Errorf("evidence_id 必须是本任务或直接关联任务的事实/漏洞节点"), nil + } _ = t.ts.Link(ev, db.RelProves, goal) _ = t.ts.SetNodeState(goal, "met") // 每标记一个目标 met,就检查本任务是否【所有目标】都已 met;若是,自动判定 @@ -668,12 +1000,19 @@ func (t *ToolSet) addFinding() actool.CoreTool { } var id int64 if t.ts != nil { + intent := pid(a.IntentID) + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图(关联任务意图只读)"), nil + } + } var err error id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors) if err != nil { return actool.Errorf(err.Error()), nil } - if intent := pid(a.IntentID); intent > 0 { + if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding } _, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors) @@ -728,15 +1067,21 @@ func (t *ToolSet) recordOneFact(it factItem, defaultIntent int64) (int64, error) if c := strings.TrimSpace(it.Confidence); c != "" { payload["confidence"] = c } + intent := pid(it.IntentID) + if intent <= 0 { + intent = defaultIntent + } + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return 0, fmt.Errorf("intent_id 必须是本任务的意图(关联任务意图只读)") + } + } // a fact is its OWN node kind (distinct from a vuln finding). id, err := t.ts.AddNode(db.KindFact, payload, 5, "confirmed", t.worker, pidList(it.AssetIDs)) if err != nil { return 0, err } - intent := pid(it.IntentID) - if intent <= 0 { - intent = defaultIntent - } if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> fact } @@ -984,6 +1329,10 @@ func (t *ToolSet) killWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图(关联任务意图只读)"), nil + } if err := t.killWork(id); err != nil { return actool.Errorf(err.Error()), nil } @@ -1014,6 +1363,10 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图(关联任务意图只读)"), nil + } if strings.TrimSpace(a.Message) == "" { return actool.Errorf("message 必填"), nil } @@ -1026,7 +1379,7 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { // getWorkerOutput returns a work's final (or截至中止时的) conclusion text by intent id. func (t *ToolSet) getWorkerOutput() actool.CoreTool { - return readTool("get_worker_output", "取某条意图(work)的最终输出结论。正常结束返回其总结;被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", + return readTool("get_worker_output", "取本任务或直接关联任务某条意图(work)的最终输出结论。关联任务结果带 source_task_id/inherited=true 且只读。正常结束返回其总结;被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", obj(map[string]any{"intent_id": idp("意图 id(= work 句柄)")}, "intent_id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -1037,7 +1390,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } - acts, _, err := t.ts.ActivityList(&id, 0, 1000) + intentNode, err := t.ts.GetNodeWithSources(id) + if err != nil { + return actool.Errorf(err.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } + acts, _, err := t.ts.ActivityListWithSources(id, 0, 1000) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1056,16 +1416,25 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { pick, terminated = fallback, true } if pick == nil { + if intentNode.Inherited { + return jsonResult(inheritedMap(map[string]any{ + "intent_id": id, "final_text": "(该 work 尚无任何输出)", "terminated": true, + }, intentNode.SourceTaskID)) + } return actool.Text("(该 work 尚无任何输出)"), nil } - detail, _ := t.ts.ActivityDetail(pick.ID) + detail, _ := t.ts.ActivityDetailWithSources(pick.ID) if detail == "" { detail = pick.Summary } - return jsonResult(map[string]any{ + result := map[string]any{ "intent_id": id, "worker_name": pick.Worker, "final_text": detail, "summary": pick.Summary, "is_error": pick.IsError, "terminated": terminated, - }) + } + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1075,10 +1444,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { func traceSteps(acts []db.Activity) []map[string]any { steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } return steps } @@ -1092,7 +1465,7 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { "① 只传 intent_id → 返回该 work 每一步的摘要流(summary≤100字,含 step_id;只是动作轮廓,不含完整输出);\n"+ "② intent_id + q → 只返回命中关键字的步骤摘要(在摘要和完整输出里都搜;仍只给 summary,要看内容用③);\n"+ "③ intent_id + step_ids(最多5个)→ 返回这些步骤的完整内容(detail)。\n"+ - "典型流程:先①/②定位可疑步骤的 step_id,再用③取其完整输出。不含思考(thinking)步骤。", + "典型流程:先①/②定位可疑步骤的 step_id,再用③取其完整输出。不含思考(thinking)步骤。支持直接关联任务的历史 trace;其结果带 source_task_id/inherited=true 且只读。", obj(map[string]any{ "intent_id": idp("意图 id(= work 句柄)"), "q": str("关键字:只返回摘要/完整输出命中它的步骤(可选;与 step_ids 互斥)"), @@ -1111,6 +1484,13 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + intentNode, nodeErr := t.ts.GetNodeWithSources(id) + if nodeErr != nil { + return actool.Errorf(nodeErr.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } // ③ detail drill-down by step ids (max 5), thinking excluded by the store. if len(a.StepIDs) > 0 { var ids []int64 @@ -1122,31 +1502,47 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if len(ids) > 5 { return actool.Errorf("step_ids 一次最多 5 个(先用不带 step_ids 的调用定位,再分批取)"), nil } - acts, err := t.ts.ActivityByIDs(ids) + acts, err := t.ts.ActivityByIDsWithSources(ids) if err != nil { return actool.Errorf(err.Error()), nil } steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + if acts[i].NodeID == nil || *acts[i].NodeID != id || acts[i].Inherited != intentNode.Inherited || + (acts[i].Inherited && acts[i].SourceTaskID != intentNode.SourceTaskID) { + continue + } + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "detail": acts[i].Detail, - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } - return jsonResult(map[string]any{"intent_id": id, "steps": steps}) + result := map[string]any{"intent_id": id, "steps": steps} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) } // ①/② summary stream, optionally keyword-filtered; 100-char summaries. var acts []db.Activity var err error if strings.TrimSpace(a.Q) != "" { - acts, err = t.ts.ActivityTraceSearch(&id, a.Q, a.Limit) + acts, err = t.ts.ActivityTraceSearchWithSources(id, a.Q, a.Limit) } else { - acts, err = t.ts.ActivityTrace(id, a.Limit) + acts, err = t.ts.ActivityTraceWithSources(id, a.Limit) } if err != nil { return actool.Errorf(err.Error()), nil } - return jsonResult(map[string]any{"intent_id": id, "steps": traceSteps(acts)}) + result := map[string]any{"intent_id": id, "steps": traceSteps(acts)} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1172,7 +1568,7 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { return actool.Errorf("q 必填"), nil } // 排除调用者自身这条意图的步骤(worker 的自有 trace 已在其上下文里)。 - acts, err := t.ts.ActivityTraceSearchExcluding(t.ownerNode, a.Q, a.Limit) + acts, err := t.ts.ActivityTraceSearchAllWithSources(t.ownerNode, a.Q, a.Limit) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1182,11 +1578,15 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { if acts[i].NodeID != nil { intent = *acts[i].NodeID } - hits = append(hits, map[string]any{ + hit := map[string]any{ "intent_id": intent, "step_id": acts[i].ID, "worker": acts[i].Worker, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(hit, acts[i].SourceTaskID) + } + hits = append(hits, hit) } return jsonResult(map[string]any{"query": a.Q, "hits": hits}) }) @@ -1217,13 +1617,16 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if limit <= 0 { limit = 50 } - all, err := t.ts.ListByKind(db.KindIntent, 500) + all, err := t.ts.ListByKindWithSources(db.KindIntent, 500) if err != nil { return actool.Errorf(err.Error()), nil } q := strings.ToLower(strings.TrimSpace(a.Q)) out := make([]map[string]any, 0, limit) for _, n := range all { + if n.Inherited && n.State == "running" { + continue + } switch n.State { case "running", "done", "exhausted", "blocked", "stopped": // has run → has a process default: @@ -1235,7 +1638,11 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if q != "" && !strings.Contains(strings.ToLower(summary), q) { continue } - out = append(out, map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State}) + item := map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State} + if n.Inherited { + inheritedMap(item, n.SourceTaskID) + } + out = append(out, item) if len(out) >= limit { break } diff --git a/agent/tools_insert.go b/agent/tools_insert.go index c808988..65f7cf2 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -398,13 +398,12 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { ) } -// listUntestedAssets lets the plan agent pull the CURRENT TASK's in-scope, not-yet- -// tested assets on demand (filter by type, paginated) — so it can decide what to -// test next itself, instead of coverage pushing a backlog list into every prompt. +// listUntestedAssets lets the plan agent pull the current + directly inherited +// scope's not-yet-tested assets on demand (filter by type, paginated). func (t *ToolSet) listUntestedAssets() actool.CoreTool { return readTool( "list_untested_assets", - "查询【本任务】范围内、还没被测过的资产(供你自己判断要不要补测,不代替你决策)。\n"+ + "查询【本任务及直接关联任务】范围内、还没被事实锚点覆盖的资产(关联范围只读,供你自己判断要不要补测,不代替你决策)。\n"+ "可选按资产类型过滤:root_domain/subdomain/service/app/endpoint/ip。\n"+ "分页:page 从 1 起、page_size 默认 10。返回 {assets:[{id,type,label}], total, page, page_size}。仅任务上下文可用。", obj(map[string]any{ @@ -432,7 +431,7 @@ func (t *ToolSet) listUntestedAssets() actool.CoreTool { a.PageSize = 10 } offset := (a.Page - 1) * a.PageSize - assets, total, err := t.as.ListUntestedAssets(t.taskID, t.ts.ID(), strings.TrimSpace(a.Type), a.PageSize, offset) + assets, total, err := t.as.ListUntestedAssetsWithSources(t.taskID, strings.TrimSpace(a.Type), a.PageSize, offset) if err != nil { return actool.Errorf(err.Error()), nil } diff --git a/agent/tools_overview_test.go b/agent/tools_overview_test.go new file mode 100644 index 0000000..8fa6a80 --- /dev/null +++ b/agent/tools_overview_test.go @@ -0,0 +1,31 @@ +package agent + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/Autumn-27/artex/db" +) + +func TestOverviewTextBudgetIsFairAndUTF8Safe(t *testing.T) { + if got := relatedOverviewBudgetForSources(1); got != relatedOverviewMaxTextPerSource { + t.Fatalf("single source budget=%d want=%d", got, relatedOverviewMaxTextPerSource) + } + perSource := relatedOverviewBudgetForSources(db.MaxTaskSourceCount) + if perSource*db.MaxTaskSourceCount > relatedOverviewTotalTextRunes { + t.Fatalf("aggregate budget exceeded: per_source=%d", perSource) + } + + budget := overviewTextBudget{remaining: 5} + got := budget.take(strings.Repeat("中", 10), 20) + if !utf8.ValidString(got) { + t.Fatalf("budget truncation produced invalid UTF-8: %q", got) + } + if utf8.RuneCountInString(got) != 5 || !budget.truncated || budget.remaining != 0 { + t.Fatalf("unexpected truncation: got=%q budget=%+v", got, budget) + } + if tail := budget.take("more", 20); tail != "" { + t.Fatalf("exhausted budget returned more text: %q", tail) + } +} diff --git a/agent/worker.go b/agent/worker.go index 71e1eb9..47e25e9 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -49,6 +49,7 @@ type Worker struct { mem *memory.Store // cross-engagement tradecraft memory (G4) tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) // runTimeout is the wall-clock budget for the main exploration of one intent // (0 = unlimited). When it fires, the run is cut and a settlement round is @@ -84,6 +85,15 @@ func NewWorker(prov llm.Provider, model, workDir string, tx *transcript.Store, w return &Worker{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, extraTools: extra} } +func (w *Worker) SetCompactionWindowResolver(fn func() int) { w.windowFn = fn } + +func (w *Worker) compactionWindow() int { + if w.windowFn != nil { + return w.windowFn() + } + return w.window +} + // SetProxy configures the recording proxy address that workers route target // traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through // that MITM proxy. Empty addr disables the hint. @@ -306,8 +316,8 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. // large tool output spills to cmd-output/ with a head + pointer (SDK tool.Capture); // full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。 ToolOutputDir: cmdOutDir(runDir), - Compaction: compactionConfig(w.window), // long tool-heavy runs stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢 + Compaction: compactionConfig(w.compactionWindow()), // long tool-heavy runs stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢 } if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic) opts.Hooks = hooks diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..7d046cc --- /dev/null +++ b/build.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# ARTEX 单二进制构建脚本:静态导出前端并嵌入 Go 二进制。 +# +# 可选环境变量: +# ARTEX_TARGET_OS=linux 目标系统,默认当前系统 +# ARTEX_TARGET_ARCH=amd64 目标架构,默认当前架构 +# ARTEX_BUILD_VERSION=v0.3.0 构建版本,默认取 git describe +# ARTEX_OUTPUT=/path/to/artex 输出路径 +# ARTEX_SKIP_NPM_CI=1 跳过 npm ci,复用现有依赖 +# ARTEX_GOSUMDB=sum.golang.org Go checksum database +set -euo pipefail + +cd "$(cd "$(dirname "$0")" && pwd)" + +info() { printf '\033[36m[*]\033[0m %s\n' "$*"; } +ok() { printf '\033[32m[+]\033[0m %s\n' "$*"; } +die() { printf '\033[31m[x]\033[0m %s\n' "$*" >&2; exit 1; } + +command -v go >/dev/null 2>&1 || die "未检测到 Go(项目需要 Go 1.26 或更高版本)" +command -v npm >/dev/null 2>&1 || die "未检测到 npm(前端静态构建需要 Node.js/npm)" +command -v rsync >/dev/null 2>&1 || die "未检测到 rsync" + +ARTEX_GOSUMDB="${ARTEX_GOSUMDB:-sum.golang.org}" +ARTEX_TARGET_OS="${ARTEX_TARGET_OS:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOOS)}" +ARTEX_TARGET_ARCH="${ARTEX_TARGET_ARCH:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOARCH)}" + +if [ -z "${ARTEX_BUILD_VERSION:-}" ]; then + if command -v git >/dev/null 2>&1 && git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + ARTEX_BUILD_VERSION="$(git describe --tags --always --dirty)" + else + ARTEX_BUILD_VERSION="dev" + fi +fi + +ARTEX_BINARY_NAME="artex" +[ "$ARTEX_TARGET_OS" = "windows" ] && ARTEX_BINARY_NAME="artex.exe" +ARTEX_OUTPUT="${ARTEX_OUTPUT:-dist/artex-${ARTEX_TARGET_OS}-${ARTEX_TARGET_ARCH}/${ARTEX_BINARY_NAME}}" + +info "构建前端静态资源" +if [ "${ARTEX_SKIP_NPM_CI:-0}" != "1" ]; then + (cd web && npm ci) +fi +(cd web && npm run build:static) + +info "同步前端资源到 server/webui/dist" +mkdir -p server/webui/dist +rsync -a --delete web/out/ server/webui/dist/ + +info "编译 ${ARTEX_TARGET_OS}/${ARTEX_TARGET_ARCH},版本 ${ARTEX_BUILD_VERSION}" +mkdir -p "$(dirname "$ARTEX_OUTPUT")" +GOSUMDB="$ARTEX_GOSUMDB" \ +CGO_ENABLED=0 \ +GOOS="$ARTEX_TARGET_OS" \ +GOARCH="$ARTEX_TARGET_ARCH" \ +go build \ + -tags embedui \ + -trimpath \ + -ldflags "-s -w -X main.version=${ARTEX_BUILD_VERSION}" \ + -o "$ARTEX_OUTPUT" \ + ./cmd/artex + +ok "编译完成:$ARTEX_OUTPUT" +if command -v file >/dev/null 2>&1; then + file "$ARTEX_OUTPUT" +fi diff --git a/cmd/artex/main.go b/cmd/artex/main.go index b5579dc..cbc5b42 100644 --- a/cmd/artex/main.go +++ b/cmd/artex/main.go @@ -73,14 +73,8 @@ func main() { sigCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - // re-wrap with a named cause so a run cancelled by shutdown says "后端进程正在 - // 关停" in its trace instead of a bare context.Canceled (see agent/cancelcause.go). - ctx, shutdown := context.WithCancelCause(sigCtx) + ctx, shutdown := shutdownContext(sigCtx) defer shutdown(agent.AbortShutdown) - go func() { - <-sigCtx.Done() - shutdown(agent.AbortShutdown) - }() mgr, err := server.NewManager(*dataDir, *proxy) if err != nil { @@ -113,3 +107,18 @@ func main() { defer cancel() _ = httpSrv.Shutdown(shutdownCtx) } + +// shutdownContext deliberately does not derive from signalCtx. If it did, the +// parent's plain context.Canceled could win the race before AbortShutdown was +// attached to the child, losing the diagnostic cause in every running Agent. +func shutdownContext(signalCtx context.Context) (context.Context, context.CancelCauseFunc) { + ctx, shutdown := context.WithCancelCause(context.Background()) + go func() { + select { + case <-signalCtx.Done(): + shutdown(agent.AbortShutdown) + case <-ctx.Done(): + } + }() + return ctx, shutdown +} diff --git a/cmd/artex/main_test.go b/cmd/artex/main_test.go new file mode 100644 index 0000000..149c7c6 --- /dev/null +++ b/cmd/artex/main_test.go @@ -0,0 +1,19 @@ +package main + +import ( + "context" + "testing" + + "github.com/Autumn-27/artex/agent" +) + +func TestShutdownContextPreservesNamedCause(t *testing.T) { + signalCtx, signalCancel := context.WithCancel(context.Background()) + ctx, shutdown := shutdownContext(signalCtx) + defer shutdown(nil) + signalCancel() + <-ctx.Done() + if code, _, _, ok := agent.AbortReason(ctx); !ok || code != "shutdown" { + t.Fatalf("code=%q ok=%v, want shutdown", code, ok) + } +} diff --git a/db/assets.go b/db/assets.go index 583cc87..c1bcb47 100644 --- a/db/assets.go +++ b/db/assets.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "net/url" + "sort" "strings" ) @@ -913,7 +914,7 @@ func (s *AssetStore) CountByType(typ string) (int, error) { } // QueryByCompany returns assets for a company, optionally filtered by type. -// limit <= 0 means no limit — pair it with CountByCompany for server-side paging. +// limit <= 0 means no limit. func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), @@ -940,8 +941,6 @@ FROM assets WHERE company_id = $1` return scanAssets(rows) } -// CountByCompany returns the full number of assets attributed to a company -// (optionally filtered by type), ignoring limit/offset. func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) { q := `SELECT count(*) FROM assets WHERE company_id = $1` args := []any{companyID} @@ -954,8 +953,6 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) { return n, err } -// pageClause appends limit/offset placeholders to args and returns the trailing -// ORDER BY … LIMIT … OFFSET fragment. limit <= 0 means no limit. func pageClause(args *[]any, limit, offset int) string { q := ` ORDER BY last_seen DESC` if limit > 0 { @@ -970,7 +967,6 @@ func pageClause(args *[]any, limit, offset int) string { } // QueryByTask returns assets rows that have a given task_id in task_ids. -// limit <= 0 means no limit — pair it with CountByTask for server-side paging. func (s *AssetStore) QueryByTask(taskID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), @@ -997,8 +993,6 @@ FROM assets WHERE $1 = ANY(task_ids)` return scanAssets(rows) } -// CountByTask returns the full number of assets attached to a task (optionally -// filtered by type), ignoring limit/offset — lets the UI paginate on the server. func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { q := `SELECT count(*) FROM assets WHERE $1 = ANY(task_ids)` args := []any{taskID} @@ -1011,7 +1005,6 @@ func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { return n, err } -// CountsByTypeForTask returns asset counts per type for a single task. func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) { rows, err := s.db.Query(`SELECT type, COUNT(*) FROM assets WHERE $1 = ANY(task_ids) GROUP BY type`, taskID) if err != nil { @@ -1021,6 +1014,151 @@ func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) { return scanTypeCounts(rows) } +// DeleteByTaskID removes assets owned only by taskID and detaches taskID from +// assets shared with other tasks. Full task deletion uses the coordinated +// transaction in DeleteTaskCascadePrepared; this method remains for callers +// that explicitly manage only asset associations. +func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) { + res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID) + if err != nil { + return 0, err + } + deleted, _ := res.RowsAffected() + if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil { + return deleted, err + } + return deleted, nil +} + +// HostsByTask returns the exact HTTP host candidates attached to a task's +// assets. Domain/IP columns cover root domains, subdomains and non-HTTP +// services; URL covers HTTP services and endpoints. +func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) { + rows, err := s.db.Query(` +SELECT COALESCE(domain,''), COALESCE(ip,''), COALESCE(url,'') +FROM assets WHERE $1 = ANY(task_ids)`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + + hosts := make(map[string]struct{}) + add := func(host string) { + host = strings.TrimSpace(strings.ToLower(host)) + if host != "" { + hosts[host] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + if err := rows.Scan(&domain, &ip, &rawURL); err != nil { + return nil, err + } + add(domain) + add(ip) + if u, err := url.Parse(rawURL); err == nil { + add(u.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(hosts)) + for host := range hosts { + out = append(out, host) + } + sort.Strings(out) + return out, nil +} + +// HostsForTaskDeletion returns hosts that belong to the task being deleted and +// are not referenced by any other live task. Current-task candidates include +// both task_ids ownership and exploration anchors so legacy seeded assets (which +// were anchor-only) are covered. Protection is host-wide: if another live task +// references any asset for a candidate host, that host's global traffic remains. +func (s *AssetStore) HostsForTaskDeletion(taskID, explorationID int64) ([]string, error) { + return hostsForTaskDeletion(s.db, taskID, explorationID) +} + +type rowsQuerier interface { + Query(query string, args ...any) (*sql.Rows, error) +} + +// hostsForTaskDeletion is shared by the read-only AssetStore API and the task +// deletion transaction. Coordinated traffic deletion must call it through the +// transaction path so asset/anchor writes stay locked until the task delete is +// committed. +func hostsForTaskDeletion(q rowsQuerier, taskID, explorationID int64) ([]string, error) { + rows, err := q.Query(` +WITH current_assets AS ( + SELECT id FROM assets WHERE $1 = ANY(task_ids) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + WHERE n.exploration_id=$2 +), +other_assets AS ( + SELECT DISTINCT a.id + FROM assets a + WHERE EXISTS ( + SELECT 1 FROM tasks t + WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids) + ) OR EXISTS ( + SELECT 1 + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + JOIN tasks t ON t.exploration_id=n.exploration_id + WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL + ) +) +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), true +FROM assets a JOIN current_assets c ON c.id=a.id +UNION ALL +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), false +FROM assets a JOIN other_assets o ON o.id=a.id`, taskID, explorationID) + if err != nil { + return nil, err + } + defer rows.Close() + + candidates := make(map[string]struct{}) + protected := make(map[string]struct{}) + add := func(dst map[string]struct{}, raw string) { + raw = strings.TrimSpace(strings.ToLower(raw)) + if raw != "" { + dst[raw] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + var candidate bool + if err := rows.Scan(&domain, &ip, &rawURL, &candidate); err != nil { + return nil, err + } + dst := protected + if candidate { + dst = candidates + } + add(dst, domain) + add(dst, ip) + if u, err := url.Parse(rawURL); err == nil { + add(dst, u.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(candidates)) + for host := range candidates { + if _, shared := protected[host]; !shared { + out = append(out, host) + } + } + sort.Strings(out) + return out, nil +} + const assetSelectCols = `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), COALESCE(c_segment::text,''), port, @@ -1053,45 +1191,6 @@ func (s *AssetStore) GetByIDs(ids []int64) ([]*Asset, error) { return scanAssets(rows) } -// DeleteByTaskID removes a task's assets: hard-deletes assets owned solely by this -// task (task_ids = {taskID}) and de-associates the rest (drops this task from their -// task_ids, keeping assets shared with other tasks intact). Returns rows deleted. -func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) { - res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID) - if err != nil { - return 0, err - } - n, _ := res.RowsAffected() - if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil { - return n, err - } - return n, nil -} - -// HostsByTask returns the distinct hostnames/IPs the task's assets carry (domain, -// root_domain, ip) — used to purge host-keyed traffic recorded for that task. -func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) { - rows, err := s.db.Query(` -SELECT DISTINCT h FROM ( - SELECT NULLIF(domain,'') AS h FROM assets WHERE $1 = ANY(task_ids) - UNION SELECT NULLIF(root_domain,'') FROM assets WHERE $1 = ANY(task_ids) - UNION SELECT NULLIF(ip,'') FROM assets WHERE $1 = ANY(task_ids) -) q WHERE h IS NOT NULL AND h <> ''`, taskID) - if err != nil { - return nil, err - } - defer rows.Close() - var out []string - for rows.Next() { - var h string - if err := rows.Scan(&h); err != nil { - return nil, err - } - out = append(out, h) - } - return out, rows.Err() -} - // DeleteByCompanyID hard-deletes all assets belonging to a company. Returns rows deleted. func (s *AssetStore) DeleteByCompanyID(companyID int64) (int64, error) { res, err := s.db.Exec(`DELETE FROM assets WHERE company_id = $1`, companyID) @@ -1161,7 +1260,6 @@ func (s *AssetStore) CountsByType() (map[string]int, error) { return scanTypeCounts(rows) } -// scanTypeCounts scans a (type, count) result set into a map. func scanTypeCounts(rows *sql.Rows) (map[string]int, error) { out := map[string]int{} for rows.Next() { diff --git a/db/assets_test.go b/db/assets_test.go index 367f7db..6e92030 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -290,12 +290,12 @@ func TestUpsertHTTPService(t *testing.T) { sc := 200 cl := int64(1024) id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ - URL: "https://www.httptest.example.com/", - Technologies: []string{"nginx", "vue"}, - StatusCode: &sc, + URL: "https://www.httptest.example.com/", + Technologies: []string{"nginx", "vue"}, + StatusCode: &sc, ContentLength: &cl, - PageTitle: "Test Site", - TaskID: 1, + PageTitle: "Test Site", + TaskID: 1, }) if err != nil { t.Fatal(err) diff --git a/db/config.go b/db/config.go index 67d9edc..8d91487 100644 --- a/db/config.go +++ b/db/config.go @@ -1,8 +1,11 @@ package db import ( + "context" "database/sql" "encoding/json" + "errors" + "fmt" "time" ) @@ -133,9 +136,236 @@ VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$ return p.ID, err } +var ( + ErrActiveLLMProfileDelete = errors.New("cannot delete the active LLM profile; activate another profile first") + ErrLLMProfileReferencesChanged = errors.New("LLM profile references changed while deleting; retry the request") + ErrLLMProfileNotFound = errors.New("LLM profile not found") +) + func (d *DB) DeleteProfile(id int64) error { - _, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) - return err + return d.DeleteProfileContext(context.Background(), id) +} + +// DeleteProfileContext removes a non-default profile while preserving the task +// failover cursor. Reference changes are retried because a task, agent, or +// conversation may start pointing at the profile between the initial scan and +// the profile row lock. A bound prevents a continuously changing workload from +// keeping an HTTP request alive forever. +func (d *DB) DeleteProfileContext(ctx context.Context, id int64) error { + const ( + maxAttempts = 8 + maxDuration = 15 * time.Second + ) + ctx, cancel := context.WithTimeout(ctx, maxDuration) + defer cancel() + for attempt := 0; attempt < maxAttempts; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + retry, err := d.deleteProfile(ctx, id) + if err != nil || !retry { + return err + } + } + return fmt.Errorf("%w: profile %d", ErrLLMProfileReferencesChanged, id) +} + +func (d *DB) deleteProfile(ctx context.Context, id int64) (bool, error) { + tx, err := d.BeginTx(ctx, nil) + if err != nil { + return false, err + } + defer tx.Rollback() + + // Reference setters lock their task/agent/conversation row before an + // llm_profiles row can be locked by a foreign-key check. Keep deletion in the + // same order: the old profile -> child order deadlocked with a setter's child + // -> profile order. Sorting also gives concurrent profile deletions a stable + // order when their chains overlap multiple tasks. + lockedTasks, err := lockProfileReferenceRows(ctx, tx, `SELECT t.id +FROM tasks t +WHERE t.llm_profile_id=$1 + OR t.active_llm_profile_id=$1 + OR EXISTS ( + SELECT 1 FROM task_llm_profiles x + WHERE x.task_id=t.id AND x.profile_id=$1 + ) +ORDER BY t.id +FOR UPDATE OF t`, id) + if err != nil { + return false, err + } + lockedAgents, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM agents +WHERE llm_profile_id=$1 +ORDER BY id +FOR UPDATE`, id) + if err != nil { + return false, err + } + lockedConversations, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM conversations +WHERE llm_profile_id=$1 +ORDER BY id +FOR UPDATE`, id) + if err != nil { + return false, err + } + + var isDefault bool + if err := tx.QueryRowContext(ctx, `SELECT is_default FROM llm_profiles WHERE id=$1 FOR UPDATE`, id).Scan(&isDefault); err != nil { + if err == sql.ErrNoRows { + return false, ErrLLMProfileNotFound + } + return false, err + } + if isDefault { + return false, ErrActiveLLMProfileDelete + } + + type affectedTask struct { + id int64 + position int + wasActive bool + } + // A task may have committed a new reference after the first statement took + // its snapshot but before this transaction acquired the profile lock. The + // profile lock now prevents further references; retry if that committed task + // was not part of the task-first lock set. Never acquire a new task lock while + // holding the profile lock, because that would recreate the inversion. + rows, err := tx.QueryContext(ctx, `SELECT ref_kind, ref_id FROM ( + SELECT 'task'::text AS ref_kind, t.id AS ref_id +FROM tasks t +WHERE t.llm_profile_id=$1 + OR t.active_llm_profile_id=$1 + OR EXISTS ( + SELECT 1 FROM task_llm_profiles x + WHERE x.task_id=t.id AND x.profile_id=$1 + ) + UNION ALL + SELECT 'agent', a.id FROM agents a WHERE a.llm_profile_id=$1 + UNION ALL + SELECT 'conversation', c.id FROM conversations c WHERE c.llm_profile_id=$1 +) refs +ORDER BY ref_kind, ref_id`, id) + if err != nil { + return false, err + } + for rows.Next() { + var ( + kind string + rowID int64 + ) + if err := rows.Scan(&kind, &rowID); err != nil { + rows.Close() + return false, err + } + locked := false + switch kind { + case "task": + _, locked = lockedTasks[rowID] + case "agent": + _, locked = lockedAgents[rowID] + case "conversation": + _, locked = lockedConversations[rowID] + } + if !locked { + rows.Close() + return true, nil + } + } + if err := rows.Err(); err != nil { + rows.Close() + return false, err + } + if err := rows.Close(); err != nil { + return false, err + } + + rows, err = tx.QueryContext(ctx, `SELECT x.task_id, x.position, COALESCE(t.active_llm_profile_id=$1, false) +FROM task_llm_profiles x +JOIN tasks t ON t.id=x.task_id +WHERE x.profile_id=$1 +ORDER BY x.task_id`, id) + if err != nil { + return false, err + } + var affected []affectedTask + for rows.Next() { + var task affectedTask + if err := rows.Scan(&task.id, &task.position, &task.wasActive); err != nil { + rows.Close() + return false, err + } + affected = append(affected, task) + } + if err := rows.Err(); err != nil { + rows.Close() + return false, err + } + if err := rows.Close(); err != nil { + return false, err + } + if _, err := tx.ExecContext(ctx, `DELETE FROM llm_profiles WHERE id=$1`, id); err != nil { + return false, err + } + // Only a successor after the deleted cursor is eligible. If there is none, + // clear the explicit chain so the task falls back to its Agent/global provider; + // profiles before a manually selected cursor must never be revived. + for _, task := range affected { + if !task.wasActive { + if _, err := tx.ExecContext(ctx, `UPDATE tasks SET llm_chain_revision=llm_chain_revision+1 WHERE id=$1`, task.id); err != nil { + return false, err + } + continue + } + var next int64 + err := tx.QueryRowContext(ctx, `SELECT profile_id FROM task_llm_profiles +WHERE task_id=$1 AND position>$2 AND status='ready' +ORDER BY position +LIMIT 1`, task.id, task.position).Scan(&next) + if err != nil && err != sql.ErrNoRows { + return false, err + } + if err == sql.ErrNoRows { + if _, err := tx.ExecContext(ctx, `DELETE FROM task_llm_profiles WHERE task_id=$1`, task.id); err != nil { + return false, err + } + if _, err := tx.ExecContext(ctx, `UPDATE tasks +SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, task.id); err != nil { + return false, err + } + continue + } + if _, err := tx.ExecContext(ctx, `UPDATE tasks +SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, task.id, next); err != nil { + return false, err + } + } + if err := tx.Commit(); err != nil { + return false, err + } + return false, nil +} + +func lockProfileReferenceRows(ctx context.Context, tx *sql.Tx, query string, profileID int64) (map[int64]struct{}, error) { + rows, err := tx.QueryContext(ctx, query, profileID) + if err != nil { + return nil, err + } + defer rows.Close() + locked := make(map[int64]struct{}) + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + locked[id] = struct{}{} + } + if err := rows.Err(); err != nil { + return nil, err + } + return locked, nil } // SetActiveProfile makes one profile the global default (single-default invariant). @@ -148,9 +378,13 @@ func (d *DB) SetActiveProfile(id int64) error { if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil { return err } - if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil { + res, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id) + if err != nil { return err } + if n, _ := res.RowsAffected(); n != 1 { + return ErrLLMProfileNotFound + } return tx.Commit() } @@ -313,8 +547,45 @@ func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error { // the binding (id == nil) so the agent follows the task/conversation pin, else the // global active profile. Precedence at runtime: agent binding → task/conv pin → active. func (d *DB) SetAgentLLMProfile(key string, id *int64) error { - _, err := d.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE key=$2`, id, key) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + // DeleteProfile locks reference rows before the profile row. Keep the same + // order here so a concurrent rebind cannot form a child/profile deadlock. + var agentID int64 + if err := tx.QueryRow(`SELECT id FROM agents WHERE key=$1 FOR UPDATE`, key).Scan(&agentID); err != nil { + if err == sql.ErrNoRows { + // Preserve the previous UPDATE semantics: an unknown key is a no-op. + return tx.Commit() + } + return err + } + if err := lockLLMProfileForReference(tx, id); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE id=$2`, id, agentID); err != nil { + return err + } + return tx.Commit() +} + +// lockLLMProfileForReference makes the profile side of the shared lock-order +// protocol explicit. Callers must already own the referencing child row. +func lockLLMProfileForReference(tx *sql.Tx, profileID *int64) error { + if profileID == nil { + return nil + } + var lockedID int64 + if err := tx.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR KEY SHARE`, *profileID).Scan(&lockedID); err != nil { + if err == sql.ErrNoRows { + return ErrLLMProfileNotFound + } + return err + } + return nil } // SetAgentWebSearch toggles whether an agent uses network search (still gated by diff --git a/db/config_test.go b/db/config_test.go index 6f1743e..16f5d8c 100644 --- a/db/config_test.go +++ b/db/config_test.go @@ -1,7 +1,9 @@ package db import ( + "context" "encoding/json" + "errors" "testing" ) @@ -53,6 +55,15 @@ func TestPoolProfilesOrder(t *testing.T) { } } +func TestDeleteProfileContextHonorsCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var d DB + if err := d.DeleteProfileContext(ctx, 1); !errors.Is(err, context.Canceled) { + t.Fatalf("DeleteProfileContext error=%v, want context cancellation", err) + } +} + func TestConfigStores(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { @@ -69,6 +80,9 @@ func TestConfigStores(t *testing.T) { if err := d.SetActiveProfile(pid); err != nil { t.Fatal(err) } + if err := d.DeleteProfile(pid); !errors.Is(err, ErrActiveLLMProfileDelete) { + t.Fatalf("deleting active profile error=%v, want %v", err, ErrActiveLLMProfileDelete) + } act, err := d.ActiveProfile() if err != nil || act == nil || act.APIKey != "secret123" { t.Fatalf("active profile/key: %+v err=%v", act, err) diff --git a/db/conversation.go b/db/conversation.go index 9e7e97a..2b2ad41 100644 --- a/db/conversation.go +++ b/db/conversation.go @@ -69,19 +69,60 @@ func scanConv(row interface{ Scan(...any) error }) (Conversation, error) { // CreateConversation opens a new chat thread for agentKey with an initial title. // llmProfileID may be nil to use the globally active profile. func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) { - c, err := scanConv(d.QueryRow(` -INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, $3) -RETURNING `+convCols, agentKey, title, llmProfileID)) + tx, err := d.Begin() if err != nil { return nil, err } + defer tx.Rollback() + + // Insert the child without a profile first. The new row is exclusively owned + // by this transaction before it takes a profile lock, matching DeleteProfile's + // child-row -> profile-row protocol. + c, err := scanConv(tx.QueryRow(` +INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, NULL) +RETURNING `+convCols, agentKey, title)) + if err != nil { + return nil, err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return nil, err + } + if llmProfileID != nil { + c, err = scanConv(tx.QueryRow(`UPDATE conversations SET llm_profile_id=$2 +WHERE id=$1 RETURNING `+convCols, c.ID, llmProfileID)) + if err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } return &c, nil } // UpdateConversationProfile sets (or clears) the LLM profile override for a conversation. func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error { - _, err := d.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + var lockedID int64 + if err := tx.QueryRow(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, id).Scan(&lockedID); err != nil { + if err == sql.ErrNoRows { + // Preserve the previous UPDATE semantics: an unknown id is a no-op. + return tx.Commit() + } + return err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID); err != nil { + return err + } + return tx.Commit() } // ListConversations returns all threads, most-recently-updated first. diff --git a/db/exploration.go b/db/exploration.go index 94f6086..be01753 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -25,34 +25,40 @@ func utf8Clean(s string) string { // Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint. type Node struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - Payload json.RawMessage `json:"payload"` - Priority int `json:"priority"` - State string `json:"state"` - Origin string `json:"origin,omitempty"` - Owner string `json:"owner,omitempty"` - Anchors []int64 `json:"anchors,omitempty"` - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + Kind string `json:"kind"` + Payload json.RawMessage `json:"payload"` + Priority int `json:"priority"` + State string `json:"state"` + Origin string `json:"origin,omitempty"` + Owner string `json:"owner,omitempty"` + BlockedReason string `json:"blocked_reason,omitempty"` + Anchors []int64 `json:"anchors,omitempty"` + CreatedAt time.Time `json:"created_at"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // Activity is one worker execution step (= old task_activity). type Activity struct { - ID int64 `json:"id"` - NodeID *int64 `json:"node_id,omitempty"` - Worker string `json:"worker,omitempty"` - Kind string `json:"kind,omitempty"` - Tool string `json:"tool,omitempty"` - ToolUseID string `json:"tool_use_id,omitempty"` - IsError bool `json:"is_error"` - Summary string `json:"summary,omitempty"` - Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + NodeID *int64 `json:"node_id,omitempty"` + Worker string `json:"worker,omitempty"` + Kind string `json:"kind,omitempty"` + Tool string `json:"tool,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + IsError bool `json:"is_error"` + Summary string `json:"summary,omitempty"` + Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads + Metadata json.RawMessage `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` // Token usage — set only on the terminal kind='result' record; nil elsewhere. - InputTokens *int `json:"input_tokens,omitempty"` - OutputTokens *int `json:"output_tokens,omitempty"` - CacheReadTokens *int `json:"cache_read_tokens,omitempty"` - CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + InputTokens *int `json:"input_tokens,omitempty"` + OutputTokens *int `json:"output_tokens,omitempty"` + CacheReadTokens *int `json:"cache_read_tokens,omitempty"` + CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // TokenUsage is a per-worker token aggregate (TokenStatsByWorker). @@ -201,17 +207,17 @@ ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, re // SetNodeState updates any node's state (never deletes). func (s *ExplorationStore) SetNodeState(id int64, state string) error { - _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) + _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1, blocked_reason=NULL WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) return err } -// SetIntentState updates an intent node's state (done/blocked/exhausted/open). +// SetIntentState updates an intent node's state. // ResetRunningIntents returns any intent left in 'running' back to 'open' so it // is re-claimed. Called on startup: a 'running' intent with no live worker (a // backend restart or crashed worker goroutine left it stuck) would otherwise spin // forever in the UI. Workers re-run an intent from scratch, so reopening is safe. func (s *ExplorationStore) ResetRunningIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) if err != nil { return 0, err @@ -226,7 +232,7 @@ WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) // changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button. func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) { res, err := s.db.Exec(`UPDATE exploration_nodes -SET state='open', completed_at=NULL +SET state='open', completed_at=NULL, blocked_reason=NULL WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state IN ('blocked','exhausted','stopped')`, id, s.expID) if err != nil { @@ -240,7 +246,7 @@ WHERE id=$1 AND exploration_id=$2 AND kind='intent' // (batch rerun after e.g. an LLM/network outage that blocked many at once). Returns the // number reopened. Workers re-run each from scratch; kept graph writes remain. func (s *ExplorationStore) ReopenBlockedIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID) if err != nil { return 0, err @@ -253,17 +259,92 @@ func (s *ExplorationStore) SetIntentState(id int64, state string) error { // terminal states stamp completed_at; reopening (back to open/running) clears it. terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped" _, err := s.db.Exec(`UPDATE exploration_nodes -SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END +SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $4 THEN now() ELSE NULL END WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal) return err } -const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at` +// IntentCleanup summarizes the blackboard records removed when a user cancels +// one worker intent. Global assets and traffic are deliberately not part of this +// operation; exploration anchors disappear only because their owning nodes do. +type IntentCleanup struct { + Intents int64 `json:"intents"` + Facts int64 `json:"facts"` + Findings int64 `json:"findings"` + Activities int64 `json:"activities"` +} + +// CancelIntent removes one local intent and the fact/finding nodes it directly +// yielded. The standalone finding row must be deleted before its node (whose FK +// otherwise uses ON DELETE SET NULL), so the entire cleanup is kept in one DB +// transaction. The caller must first stop a running worker to prevent late writes. +func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) { + var out IntentCleanup + tx, err := s.db.Begin() + if err != nil { + return out, err + } + defer tx.Rollback() + + var state string + if err := tx.QueryRow(`SELECT state FROM exploration_nodes + WHERE id=$1 AND exploration_id=$2 AND kind='intent' FOR UPDATE`, id, s.expID).Scan(&state); err != nil { + if err == sql.ErrNoRows { + return out, fmt.Errorf("intent not found") + } + return out, err + } + if state != "running" && state != "paused" { + return out, fmt.Errorf("intent state %s cannot be cancelled", state) + } + + if err := tx.QueryRow(`SELECT + COUNT(*) FILTER (WHERE n.kind='fact'), + COUNT(*) FILTER (WHERE n.kind='finding') + FROM exploration_edges e + JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.kind IN ('fact','finding')`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil { + return out, err + } + + if _, err := tx.Exec(`DELETE FROM findings f USING exploration_edges e, exploration_nodes n + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.id=e.dst_id AND n.exploration_id=e.exploration_id + AND n.kind='finding' AND f.node_id=n.id`, s.expID, id, RelYields); err != nil { + return out, err + } + res, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id) + if err != nil { + return out, err + } + out.Activities, _ = res.RowsAffected() + + res, err = tx.Exec(`DELETE FROM exploration_nodes + WHERE exploration_id=$1 AND ( + id=$2 OR id IN ( + SELECT e.dst_id FROM exploration_edges e + JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.kind IN ('fact','finding') + ) + )`, s.expID, id, RelYields) + if err != nil { + return out, err + } + if removed, _ := res.RowsAffected(); removed != 1+out.Facts+out.Findings { + return out, fmt.Errorf("intent cleanup changed concurrently") + } + out.Intents = 1 + return out, tx.Commit() +} + +const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at` func scanNode(sc interface{ Scan(...any) error }) (*Node, error) { var n Node var payload []byte - if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil { + if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.BlockedReason, &n.CreatedAt); err != nil { return nil, err } n.Payload = json.RawMessage(payload) @@ -455,6 +536,31 @@ WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields return out, rows.Err() } +// FindingIntentsTerminal is the inherited-history variant: it returns finding +// lineage only when the producing intent has reached an immutable terminal +// state. This keeps a source task's live worker topology private. +func (s *ExplorationStore) FindingIntentsTerminal() (map[int64]int64, error) { + rows, err := s.db.Query(`SELECT e.dst_id, e.src_id + FROM exploration_edges e + JOIN exploration_nodes finding ON finding.id=e.dst_id AND finding.exploration_id=e.exploration_id + JOIN exploration_nodes intent ON intent.id=e.src_id AND intent.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.rel=$2 AND finding.kind='finding' + AND intent.kind='intent' AND intent.state IN ('done','blocked','exhausted','stopped')`, s.expID, RelYields) + if err != nil { + return nil, err + } + defer rows.Close() + out := map[int64]int64{} + for rows.Next() { + var findingID, intentID int64 + if err := rows.Scan(&findingID, &intentID); err != nil { + return nil, err + } + out[findingID] = intentID + } + return out, rows.Err() +} + // Frontier returns open intents ordered by priority desc then id (FIFO tiebreak). func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) { if limit <= 0 { @@ -516,12 +622,16 @@ func (s *ExplorationStore) Stats() (map[string]int, error) { // AppendActivity records one worker step and returns its id. func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) { + metadata := a.Metadata + if len(metadata) == 0 { + metadata = json.RawMessage(`{}`) + } var id int64 err := s.db.QueryRow(` -INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13) +INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14) RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail), - a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) + metadata, a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) return id, err } @@ -659,7 +769,7 @@ func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) } var rows *sql.Rows var err error - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` if nodeID != nil { rows, err = s.db.Query(`SELECT `+cols+` FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit) @@ -673,6 +783,43 @@ FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID, defer rows.Close() out := []Activity{} cursor := sinceID + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, sinceID, err + } + if a.ID > cursor { + cursor = a.ID + } + out = append(out, a) + } + return out, cursor, rows.Err() +} + +// ActivityListForTerminalIntent is the inherited, incremental-read boundary. +// The intent state check and activity read happen in one statement so a source +// intent reopened between API-level checks cannot expose its new in-flight trace. +// Model reasoning and accounting rows are never inherited. +func (s *ExplorationStore) ActivityListForTerminalIntent(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + if limit <= 0 { + limit = 300 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND a.id>$3 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $4`, s.expID, nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + defer rows.Close() + out := []Activity{} + cursor := sinceID for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -719,7 +866,7 @@ func (s *ExplorationStore) ActivityPage(f ActivitySessionFilter, before int64, l if limit <= 0 { limit = 200 } - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` args := []any{s.expID} cond, cargs := f.cond(len(args) + 1) args = append(args, cargs...) @@ -736,6 +883,51 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) } defer rows.Close() desc := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, false, err + } + desc = append(desc, a) + } + if err := rows.Err(); err != nil { + return nil, false, err + } + hasMore := len(desc) > limit + if hasMore { + desc = desc[:limit] + } + // reverse the newest-first window into ascending id order for display. + out := make([]Activity, len(desc)) + for i, a := range desc { + out[len(desc)-1-i] = a + } + return out, hasMore, nil +} + +// ActivityPageForTerminalIntent is the inherited session-history boundary. It +// combines ownership, terminal-state and row-kind checks in one query, closing +// the race where a previously completed source intent is reopened before paging. +func (s *ExplorationStore) ActivityPageForTerminalIntent(nodeID, before int64, limit int) ([]Activity, bool, error) { + if limit <= 0 { + limit = 200 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND ($3 <= 0 OR a.id < $3) + ORDER BY a.id DESC LIMIT $4`, s.expID, nodeID, before, limit+1) + if err != nil { + return nil, false, err + } + defer rows.Close() + desc := []Activity{} for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -751,7 +943,6 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) if hasMore { desc = desc[:limit] } - // reverse the newest-first window into ascending id order for display. out := make([]Activity, len(desc)) for i, a := range desc { out[len(desc)-1-i] = a @@ -817,6 +1008,26 @@ ORDER BY id LIMIT $3`, s.expID, nodeID, limit) return scanTrace(rows) } +// ActivityTraceForTerminalIntent returns one inherited work trace only while its +// source intent is still terminal at query time. +func (s *ExplorationStore) ActivityTraceForTerminalIntent(nodeID int64, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 500 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $3`, s.expID, nodeID, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // ActivityTraceSearch finds steps whose summary OR detail matches q // (case-insensitive), returning summaries only. nodeID != nil scopes to one work; // nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no @@ -844,13 +1055,59 @@ AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like, return scanTrace(rows) } +// ActivityTraceSearchForTerminalIntent is the scoped inherited search variant; +// terminal eligibility is evaluated by the same statement that reads the rows. +func (s *ExplorationStore) ActivityTraceSearchForTerminalIntent(nodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $3 OR a.detail ILIKE $3) + ORDER BY a.id LIMIT $4`, s.expID, nodeID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + +// ActivityTraceSearchTerminalIntents searches inherited worker history without +// exposing planner/main rows or work that is still open/running in the source. +func (s *ExplorationStore) ActivityTraceSearchTerminalIntents(q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $2 OR a.detail ILIKE $2) + ORDER BY a.id LIMIT $3`, s.expID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // AssetRef is a compact exploration node (intent/fact/finding) that references an // asset via an anchor — for the coverage graph's node drawer. type AssetRef struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - State string `json:"state"` - Summary string `json:"summary"` + ID int64 `json:"id"` + Kind string `json:"kind"` + State string `json:"state"` + Summary string `json:"summary"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // AssetRefs returns the intents / facts / findings in THIS exploration anchored to @@ -943,6 +1200,42 @@ FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings. return out, rows.Err() } +// ActivityByIDsForTerminalIntents is the inherited-detail read boundary. It +// deliberately excludes node-less planner/main rows, non-intent activities, live +// source work, and thinking. Local callers that need legacy behavior use +// ActivityByIDs instead. +func (s *ExplorationStore) ActivityByIDsForTerminalIntents(ids []int64) ([]Activity, error) { + if len(ids) == 0 { + return nil, nil + } + ph := make([]string, len(ids)) + args := make([]any, len(ids)+1) + args[0] = s.expID + for i, id := range ids { + ph[i] = fmt.Sprintf("$%d", i+2) + args[i+1] = id + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.detail,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.kind NOT IN ('thinking','usage') AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.id IN (`+strings.Join(ph, ",")+`) ORDER BY a.id`, args...) + if err != nil { + return nil, err + } + defer rows.Close() + out := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil { + return nil, err + } + out = append(out, a) + } + return out, rows.Err() +} + // AddStandaloneFinding writes a finding to the standalone findings table, which // persists across task deletion (task_id / node_id become NULL when the task or // exploration node is deleted). taskID and nodeID may be 0 (stored as NULL). diff --git a/db/exploration_sources.go b/db/exploration_sources.go new file mode 100644 index 0000000..3ccbb96 --- /dev/null +++ b/db/exploration_sources.go @@ -0,0 +1,350 @@ +package db + +import ( + "database/sql" + "sort" +) + +// DirectSourceStore binds one directly related task to its exploration store. +// It is intentionally a read-side helper: callers keep using the receiver store +// for every graph mutation, frontier lookup, and intent claim. +type DirectSourceStore struct { + Task TaskSource + Store *ExplorationStore +} + +// DirectSourceStores resolves the live, direct task relations for this +// exploration. It deliberately queries on every call: relations disappear when +// a source task is deleted, and inherited context must not retain stale rows. +// Sources of a source are never expanded. +func (s *ExplorationStore) DirectSourceStores() ([]DirectSourceStore, error) { + rows, err := s.db.Query(` +SELECT source.id, source.exploration_id, source.description, source.goal, source.status +FROM tasks owner +JOIN task_relations relation ON relation.task_id=owner.id +JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL +WHERE owner.exploration_id=$1 AND owner.deleted_at IS NULL +ORDER BY relation.created_at, source.id`, s.expID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []DirectSourceStore{} + for rows.Next() { + var source TaskSource + if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil { + return nil, err + } + out = append(out, DirectSourceStore{Task: source, Store: s.db.Exploration(source.ExplorationID)}) + } + return out, rows.Err() +} + +// TaskID returns the live task bound to this exploration. Explorations created +// directly in tests or maintenance code have no task and return zero. +func (s *ExplorationStore) TaskID() (int64, error) { + var id int64 + err := s.db.QueryRow(`SELECT id FROM tasks WHERE exploration_id=$1 AND deleted_at IS NULL`, s.expID).Scan(&id) + if err != nil { + if err == sql.ErrNoRows { + return 0, nil + } + return 0, err + } + return id, nil +} + +func markInheritedNode(n *Node, taskID int64) *Node { + if n == nil { + return nil + } + n.SourceTaskID = taskID + n.Inherited = true + return n +} + +func markInheritedActivities(in []Activity, taskID int64) []Activity { + for i := range in { + in[i].SourceTaskID = taskID + in[i].Inherited = true + } + return in +} + +func inheritedIntentTerminal(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false + } +} + +// ListByKindWithSources returns local nodes followed by nodes from each direct +// source in relation order. The limit remains per exploration, matching the +// existing ListByKind contract while ensuring one large task cannot hide all +// inherited context from another source. +func (s *ExplorationStore) ListByKindWithSources(kind string, limit int) ([]*Node, error) { + out, err := s.ListByKind(kind, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + nodes, err := source.Store.ListByKind(kind, limit) + if err != nil { + return nil, err + } + for _, node := range nodes { + if kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + out = append(out, markInheritedNode(node, source.Task.TaskID)) + } + } + return out, nil +} + +// GetNodeWithSources reads a node only when it belongs to this exploration or +// one of its direct sources. Inherited nodes are tagged so tool callers can keep +// them read-only and show their provenance. +func (s *ExplorationStore) GetNodeWithSources(id int64) (*Node, error) { + node, err := s.GetNode(id) + if err != nil || node != nil { + return node, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, err = source.Store.GetNode(id) + if err != nil { + return nil, err + } + if node != nil { + if node.Kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + return markInheritedNode(node, source.Task.TaskID), nil + } + } + return nil, nil +} + +// FindingIntentsWithSources combines finding lineage for the current +// exploration and each direct source. Node ids are globally unique. +func (s *ExplorationStore) FindingIntentsWithSources() (map[int64]int64, error) { + out, err := s.FindingIntents() + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + items, err := source.Store.FindingIntentsTerminal() + if err != nil { + return nil, err + } + for findingID, intentID := range items { + out[findingID] = intentID + } + } + return out, nil +} + +// ActivityTraceWithSources returns a work trace when its intent belongs to the +// current exploration or a direct source. It never searches indirect sources. +func (s *ExplorationStore) ActivityTraceWithSources(nodeID int64, limit int) ([]Activity, error) { + acts, err := s.ActivityTrace(nodeID, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceForTerminalIntent(nodeID, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityListWithSources is the source-aware equivalent used by +// get_worker_output. Node ids are global, so the first owning exploration is +// unambiguous even when the work has not emitted any activity yet. +func (s *ExplorationStore) ActivityListWithSources(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + node, err := s.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node != nil { + return s.ActivityList(&nodeID, sinceID, limit) + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, sinceID, err + } + for _, source := range sources { + node, err = source.Store.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, cursor, err := source.Store.ActivityListForTerminalIntent(nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + return markInheritedActivities(acts, source.Task.TaskID), cursor, nil + } + return []Activity{}, sinceID, nil +} + +// ActivityDetailWithSources keeps the legacy local-task lookup (including local +// thinking rows), while inherited details are restricted to terminal worker +// intents. Source planner/main rows have no node id and must never become part of +// inherited context. +func (s *ExplorationStore) ActivityDetailWithSources(id int64) (string, error) { + detail, err := s.ActivityDetail(id) + if err != nil || detail != "" { + return detail, err + } + acts, err := s.ActivityByIDsWithSources([]int64{id}) + if err != nil || len(acts) == 0 { + return "", err + } + return acts[0].Detail, nil +} + +// ActivityTraceSearchWithSources performs the scoped worker-trace search used +// by get_worker_trace. A node id resolves to at most one exploration because +// exploration node ids are global. +func (s *ExplorationStore) ActivityTraceSearchWithSources(nodeID int64, q string, limit int) ([]Activity, error) { + acts, err := s.ActivityTraceSearch(&nodeID, q, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceSearchForTerminalIntent(nodeID, q, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityTraceSearchAllWithSources searches local worker traces plus every +// direct source. The local owner is excluded only from the current exploration; +// inherited traces are immutable historical context. +func (s *ExplorationStore) ActivityTraceSearchAllWithSources(excludeNodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + out, err := s.ActivityTraceSearchExcluding(excludeNodeID, q, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityTraceSearchTerminalIntents(q, limit) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +// ActivityByIDsWithSources loads local step details with the legacy behavior. For +// direct sources it only returns rows attached to terminal intents, preventing an +// arbitrary global activity id from exposing source planner/main transcripts. +func (s *ExplorationStore) ActivityByIDsWithSources(ids []int64) ([]Activity, error) { + out, err := s.ActivityByIDs(ids) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityByIDsForTerminalIntents(ids) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + return out, nil +} + +// AssetRefsWithSources returns anchored nodes from this exploration and each +// direct source. Inherited entries retain their owning task id so API/UI callers +// can present them as immutable context. +func (s *ExplorationStore) AssetRefsWithSources(assetID int64) ([]AssetRef, error) { + out, err := s.AssetRefs(assetID) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + refs, err := source.Store.AssetRefs(assetID) + if err != nil { + return nil, err + } + for i := range refs { + if refs[i].Kind == KindIntent && !inheritedIntentTerminal(refs[i].State) { + continue + } + refs[i].SourceTaskID = source.Task.TaskID + refs[i].Inherited = true + out = append(out, refs[i]) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].ID > out[j].ID }) + return out, nil +} diff --git a/db/exploration_sources_test.go b/db/exploration_sources_test.go new file mode 100644 index 0000000..7c66fae --- /dev/null +++ b/db/exploration_sources_test.go @@ -0,0 +1,431 @@ +package db + +import ( + "fmt" + "testing" +) + +func TestInheritedActivityReadsRequireTerminalIntent(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("source activity boundary", "source activity boundary") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) }) + store := d.Exploration(expID) + intentID, err := store.AddIntent(map[string]any{"summary": "source work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + appendStep := func(kind, summary string) int64 { + t.Helper() + id, appendErr := store.AppendActivity(Activity{ + NodeID: &intentID, Worker: "worker", Kind: kind, Summary: summary, Detail: summary + " detail", + }) + if appendErr != nil { + t.Fatal(appendErr) + } + return id + } + textID := appendStep("text", "shared text") + thinkingID := appendStep("thinking", "private reasoning") + usageID := appendStep("usage", "private accounting") + resultID := appendStep("result", "shared result") + if err := store.SetIntentState(intentID, "done"); err != nil { + t.Fatal(err) + } + + page, more, err := store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || more || len(page) != 2 || page[0].ID != textID || page[1].ID != resultID { + t.Fatalf("terminal page boundary: page=%+v more=%v err=%v", page, more, err) + } + list, cursor, err := store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 2 || cursor != resultID { + t.Fatalf("terminal list boundary: list=%+v cursor=%d err=%v", list, cursor, err) + } + trace, err := store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 2 { + t.Fatalf("terminal trace boundary: trace=%+v err=%v", trace, err) + } + hits, err := store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 2 { + t.Fatalf("terminal search boundary: hits=%+v err=%v", hits, err) + } + details, err := store.ActivityByIDsForTerminalIntents([]int64{textID, thinkingID, usageID, resultID}) + if err != nil || len(details) != 2 || details[0].ID != textID || details[1].ID != resultID { + t.Fatalf("terminal detail boundary: details=%+v err=%v", details, err) + } + + // Reopening a source intent must close every inherited activity read even if + // callers still hold a stale terminal-state snapshot. + if err := store.SetIntentState(intentID, "running"); err != nil { + t.Fatal(err) + } + page, _, err = store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || len(page) != 0 { + t.Fatalf("reopened intent leaked through page: page=%+v err=%v", page, err) + } + list, _, err = store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 0 { + t.Fatalf("reopened intent leaked through list: list=%+v err=%v", list, err) + } + trace, err = store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 0 { + t.Fatalf("reopened intent leaked through trace: trace=%+v err=%v", trace, err) + } + hits, err = store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 0 { + t.Fatalf("reopened intent leaked through search: hits=%+v err=%v", hits, err) + } + details, err = store.ActivityByIDsForTerminalIntents([]int64{textID, resultID}) + if err != nil || len(details) != 0 { + t.Fatalf("reopened intent leaked through detail: details=%+v err=%v", details, err) + } +} + +func TestExplorationDirectSourceReadView(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand source", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("direct source", "source goal", TaskCreateOptions{ + SourceTaskIDs: []int64{grand.ID}, + }) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current", "current goal", TaskCreateOptions{ + SourceTaskIDs: []int64{source.ID}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + }) + + grandStore := d.Exploration(grand.ExplorationID) + sourceStore := d.Exploration(source.ExplorationID) + currentStore := d.Exploration(current.ExplorationID) + + grandFact, err := grandStore.AddNode(KindFact, map[string]any{"summary": "grand-only fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFact, err := sourceStore.AddNode(KindFact, map[string]any{"summary": "direct source fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFinding, err := sourceStore.AddNode(KindFinding, map[string]any{"summary": "direct source finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work"}, 9, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.Link(sourceIntent, RelYields, sourceFinding); err != nil { + t.Fatal(err) + } + currentIntent, err := currentStore.AddIntent(map[string]any{"summary": "current work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + currentFact, err := currentStore.AddNode(KindFact, map[string]any{"summary": "current fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + + sources, err := currentStore.DirectSourceStores() + if err != nil { + t.Fatal(err) + } + if len(sources) != 1 || sources[0].Task.TaskID != source.ID { + t.Fatalf("want only direct source %d, got %+v", source.ID, sources) + } + + frontier, err := currentStore.Frontier(20) + if err != nil { + t.Fatal(err) + } + if len(frontier) != 1 || frontier[0].ID != currentIntent { + t.Fatalf("source intent leaked into local frontier: %+v", frontier) + } + if claimed, err := currentStore.ClaimIntent(sourceIntent, "wrong-task-worker"); err != nil || claimed { + t.Fatalf("source intent must not be claimable: claimed=%v err=%v", claimed, err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got != nil { + t.Fatalf("live source intent must not be readable as inherited history: node=%+v err=%v", got, err) + } + if got, err := currentStore.ListByKindWithSources(KindIntent, 100); err != nil || len(got) != 1 || got[0].ID != currentIntent { + t.Fatalf("live source intent leaked through inherited list: nodes=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != 0 { + t.Fatalf("live source intent leaked through finding lineage: lineage=%+v err=%v", got, err) + } + + facts, err := currentStore.ListByKindWithSources(KindFact, 100) + if err != nil { + t.Fatal(err) + } + seenCurrent, seenSource, seenGrand := false, false, false + for _, fact := range facts { + switch fact.ID { + case currentFact: + seenCurrent = !fact.Inherited && fact.SourceTaskID == 0 + case sourceFact: + seenSource = fact.Inherited && fact.SourceTaskID == source.ID + case grandFact: + seenGrand = true + } + } + if !seenCurrent || !seenSource || seenGrand { + t.Fatalf("unexpected direct-source fact view: current=%v source=%v grand=%v", seenCurrent, seenSource, seenGrand) + } + + gotSource, err := currentStore.GetNodeWithSources(sourceFact) + if err != nil || gotSource == nil || !gotSource.Inherited || gotSource.SourceTaskID != source.ID { + t.Fatalf("source node lookup: node=%+v err=%v", gotSource, err) + } + gotGrand, err := currentStore.GetNodeWithSources(grandFact) + if err != nil || gotGrand != nil { + t.Fatalf("indirect source must be invisible: node=%+v err=%v", gotGrand, err) + } + + // Existing write methods remain bound to currentStore.expID. Passing an + // inherited id is a no-op and cannot mutate the source blackboard. + if err := currentStore.SetNodeState(sourceFact, "dismissed"); err != nil { + t.Fatal(err) + } + unchanged, err := sourceStore.GetNode(sourceFact) + if err != nil || unchanged == nil || unchanged.State != "confirmed" { + t.Fatalf("inherited fact was mutated: node=%+v err=%v", unchanged, err) + } + + if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil { + t.Fatal(err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got == nil || !got.Inherited { + t.Fatalf("terminal source intent should become readable history: node=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != sourceIntent { + t.Fatalf("terminal source finding lineage missing: lineage=%+v err=%v", got, err) + } + sourcePlannerStep, err := sourceStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "private source planner", Detail: "private source planner transcript", + }) + if err != nil { + t.Fatal(err) + } + openSourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work still running"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.SetIntentState(openSourceIntent, "running"); err != nil { + t.Fatal(err) + } + openSourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &openSourceIntent, Worker: "source-live-worker", Kind: "text", + Summary: "private live source work", Detail: "private live source transcript", + }) + if err != nil { + t.Fatal(err) + } + sourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &sourceIntent, Worker: "source-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "source-secret response", Detail: "source-secret full detail", + }) + if err != nil { + t.Fatal(err) + } + currentPlannerStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "current planner", Detail: "current planner detail", + }) + if err != nil { + t.Fatal(err) + } + currentThinkingStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "thinking", Summary: "current thinking", Detail: "current thinking detail", + }) + if err != nil { + t.Fatal(err) + } + grandStep, err := grandStore.AppendActivity(Activity{ + NodeID: &grandFact, Worker: "grand-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "grand-only response", Detail: "grand-only detail", + }) + if err != nil { + t.Fatal(err) + } + + trace, err := currentStore.ActivityTraceWithSources(sourceIntent, 20) + if err != nil || len(trace) != 1 || trace[0].ID != sourceStep || !trace[0].Inherited || trace[0].SourceTaskID != source.ID { + t.Fatalf("source trace lookup: trace=%+v err=%v", trace, err) + } + hits, err := currentStore.ActivityTraceSearchAllWithSources(0, "source-secret", 20) + if err != nil || len(hits) != 1 || hits[0].ID != sourceStep || !hits[0].Inherited { + t.Fatalf("source trace search: hits=%+v err=%v", hits, err) + } + grandHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "grand-only", 20) + if err != nil || len(grandHits) != 0 { + t.Fatalf("indirect source trace leaked: hits=%+v err=%v", grandHits, err) + } + liveHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "private live source", 20) + if err != nil || len(liveHits) != 0 { + t.Fatalf("live source trace leaked: hits=%+v err=%v", liveHits, err) + } + if liveTrace, err := currentStore.ActivityTraceWithSources(openSourceIntent, 20); err != nil || len(liveTrace) != 0 { + t.Fatalf("live inherited intent trace leaked: trace=%+v err=%v", liveTrace, err) + } + details, err := currentStore.ActivityByIDsWithSources([]int64{ + sourcePlannerStep, openSourceStep, sourceStep, currentPlannerStep, grandStep, + }) + if err != nil || len(details) != 2 || details[0].ID != sourceStep || !details[0].Inherited || details[1].ID != currentPlannerStep || details[1].Inherited { + t.Fatalf("source trace detail allow-list: details=%+v err=%v", details, err) + } + if detail, err := currentStore.ActivityDetailWithSources(sourcePlannerStep); err != nil || detail != "" { + t.Fatalf("source planner transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(openSourceStep); err != nil || detail != "" { + t.Fatalf("live source intent transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentPlannerStep); err != nil || detail != "current planner detail" { + t.Fatalf("local activity detail behavior changed: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentThinkingStep); err != nil || detail != "current thinking detail" { + t.Fatalf("local thinking detail behavior changed: detail=%q err=%v", detail, err) + } +} + +func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand assets", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("source assets", "source goal", TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}}) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current assets", "current goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + sourceTestedDomain := fmt.Sprintf("context-tested-%d.invalid", source.ID) + sourceUntestedDomain := fmt.Sprintf("context-untested-%d.invalid", source.ID) + sourceAnchoredOnlyDomain := fmt.Sprintf("context-anchor-only-%d.invalid", source.ID) + grandDomain := fmt.Sprintf("context-grand-%d.invalid", grand.ID) + sourceTestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceTestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceUntestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceUntestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceAnchoredOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceAnchoredOnlyDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + grandID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: grandDomain, TaskID: grand.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + _, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID}) + }) + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "tested source asset"}, 0, "confirmed", "worker", []int64{sourceTestedID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "anchored-only source asset"}, 0, "confirmed", "worker", []int64{sourceAnchoredOnlyID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(grand.ExplorationID).AddNode(KindFact, map[string]any{"summary": "indirect tested asset"}, 0, "confirmed", "worker", []int64{grandID}); err != nil { + t.Fatal(err) + } + + scopes, err := assets.ListTaskScopeWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if len(scopes) != 2 || scopes[0].TaskID != source.ID || scopes[1].TaskID != source.ID { + t.Fatalf("direct source scopes only: %+v", scopes) + } + cov, err := assets.TaskCoverageWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if cov.Denominator != 3 || cov.Tested != 2 { + t.Fatalf("combined coverage: %+v", cov) + } + untested, total, err := assets.ListUntestedAssetsWithSources(current.ID, "", 10, 0) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(untested) != 1 || untested[0].ID != sourceUntestedID { + t.Fatalf("combined untested assets: total=%d assets=%+v", total, untested) + } + hosts, err := assets.HostsByTaskWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + hostSet := map[string]bool{} + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[sourceTestedDomain] || !hostSet[sourceUntestedDomain] || !hostSet[sourceAnchoredOnlyDomain] || hostSet[grandDomain] { + t.Fatalf("direct source hosts only: %v", hosts) + } + refs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(sourceAnchoredOnlyID) + if err != nil || len(refs) != 1 || !refs[0].Inherited || refs[0].SourceTaskID != source.ID || refs[0].Kind != KindFact { + t.Fatalf("source asset refs provenance: refs=%+v err=%v", refs, err) + } + grandRefs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(grandID) + if err != nil || len(grandRefs) != 0 { + t.Fatalf("indirect source asset refs leaked: refs=%+v err=%v", grandRefs, err) + } + graph, err := assets.BuildCoverageGraph(current.ID, current.ExplorationID) + if err != nil { + t.Fatal(err) + } + seenAnchoredOnly := false + for _, node := range graph.Nodes { + if node.AssetID == sourceAnchoredOnlyID { + seenAnchoredOnly = node.Tested && node.InScope + } + } + if !seenAnchoredOnly { + t.Fatalf("anchored-only source asset missing from coverage graph: %+v", graph.Nodes) + } +} diff --git a/db/exploration_test.go b/db/exploration_test.go index 1c31d9f..97d9f62 100644 --- a/db/exploration_test.go +++ b/db/exploration_test.go @@ -1,6 +1,9 @@ package db -import "testing" +import ( + "fmt" + "testing" +) func TestExplorationFlow(t *testing.T) { d, err := Open(testDSN(t)) @@ -112,3 +115,134 @@ func TestExplorationFlow(t *testing.T) { t.Fatalf("stats: %+v", st) } } + +func TestIntentPauseResumeAndCancelCleanup(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("test", "worker control cleanup") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + es := d.Exploration(expID) + + assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("cancel-intent-%d.invalid", expID), + }) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, assetID) + + otherIntent, err := es.AddIntent(map[string]any{"summary": "keep intent"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + intentID, err := es.AddIntent(map[string]any{"summary": "cancel intent"}, 10, nil, "planner") + if err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-control-test"); err != nil || !claimed { + t.Fatalf("initial claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + frontier, err := es.Frontier(10) + if err != nil { + t.Fatal(err) + } + for _, node := range frontier { + if node.ID == intentID { + t.Fatalf("paused intent %d must not enter frontier", intentID) + } + } + if claimed, err := es.ClaimIntent(intentID, "worker-while-paused"); err != nil || claimed { + t.Fatalf("paused intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "open"); err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-after-resume"); err != nil || !claimed { + t.Fatalf("resumed intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + directFact, err := es.AddNode("fact", map[string]any{"summary": "remove fact"}, 0, "confirmed", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + directFinding, err := es.AddNode("finding", map[string]any{"summary": "remove finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + keptFact, err := es.AddNode("fact", map[string]any{"summary": "keep fact"}, 0, "confirmed", "other-worker", nil) + if err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFact); err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFinding); err != nil { + t.Fatal(err) + } + findingRowID, err := es.AddStandaloneFinding(0, directFinding, "test", "cancelled finding", SeverityHigh, "summary", "evidence", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + activityID, err := es.AppendActivity(Activity{NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "remove activity"}) + if err != nil { + t.Fatal(err) + } + keptActivityID, err := es.AppendActivity(Activity{NodeID: &otherIntent, Worker: "other-worker", Kind: "result", Summary: "keep activity"}) + if err != nil { + t.Fatal(err) + } + + cleanup, err := es.CancelIntent(intentID) + if err != nil { + t.Fatal(err) + } + if cleanup.Intents != 1 || cleanup.Facts != 1 || cleanup.Findings != 1 || cleanup.Activities != 1 { + t.Fatalf("unexpected cleanup counts: %+v", cleanup) + } + for _, nodeID := range []int64{intentID, directFact, directFinding} { + node, err := es.GetNode(nodeID) + if err != nil { + t.Fatal(err) + } + if node != nil { + t.Fatalf("node %d survived intent cancellation", nodeID) + } + } + for _, nodeID := range []int64{otherIntent, keptFact} { + node, err := es.GetNode(nodeID) + if err != nil || node == nil { + t.Fatalf("unrelated node %d removed: node=%v err=%v", nodeID, node, err) + } + } + + assertCount := func(query string, want int, args ...any) { + t.Helper() + var got int + if err := d.QueryRow(query, args...).Scan(&got); err != nil { + t.Fatal(err) + } + if got != want { + t.Fatalf("query count=%d, want %d: %s", got, want, query) + } + } + assertCount(`SELECT COUNT(*) FROM findings WHERE id=$1`, 0, findingRowID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 0, activityID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 1, keptActivityID) + assertCount(`SELECT COUNT(*) FROM exploration_edges WHERE exploration_id=$1 AND (src_id=$2 OR dst_id=$2)`, 0, expID, intentID) + assertCount(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1`, 0, directFact) + assertCount(`SELECT COUNT(*) FROM assets WHERE id=$1`, 1, assetID) +} diff --git a/db/findings.go b/db/findings.go index 6fc64db..9dfe6e2 100644 --- a/db/findings.go +++ b/db/findings.go @@ -9,7 +9,8 @@ import ( "time" ) -// DBFinding is a row in the standalone findings table (persists across task deletion). +// DBFinding is a row in the standalone findings table. It persists across task +// deletion unless the caller explicitly requests related finding cleanup. type DBFinding struct { ID int64 TaskID *int64 diff --git a/db/schema.sql b/db/schema.sql index 32ba6f9..90779f7 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -129,19 +129,40 @@ CREATE TABLE IF NOT EXISTS exploration_nodes ( state TEXT NOT NULL DEFAULT 'open', origin TEXT, owner TEXT, + blocked_reason TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), completed_at TIMESTAMPTZ, CONSTRAINT ck_node_kind CHECK (kind IN ('begin','goal','intent','fact','finding','hint')), CONSTRAINT ck_node_state CHECK ( (kind='begin' AND state IN ('open')) OR - (kind='intent' AND state IN ('open','running','done','blocked','exhausted','stopped')) OR + (kind='intent' AND state IN ('open','running','paused','done','blocked','exhausted','stopped')) OR (kind='goal' AND state IN ('open','met','abandoned')) OR (kind='fact' AND state IN ('confirmed','dismissed','origin')) OR (kind='finding' AND state IN ('confirmed','dismissed')) OR (kind='hint' AND state IN ('active','consumed')) ) ); +ALTER TABLE exploration_nodes ADD COLUMN IF NOT EXISTS blocked_reason TEXT; +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conrelid='exploration_nodes'::regclass + AND conname='ck_node_state' + AND pg_get_constraintdef(oid) NOT LIKE '%paused%' + ) THEN + ALTER TABLE exploration_nodes DROP CONSTRAINT ck_node_state; + ALTER TABLE exploration_nodes ADD CONSTRAINT ck_node_state CHECK ( + (kind='begin' AND state IN ('open')) OR + (kind='intent' AND state IN ('open','running','paused','done','blocked','exhausted','stopped')) OR + (kind='goal' AND state IN ('open','met','abandoned')) OR + (kind='fact' AND state IN ('confirmed','dismissed','origin')) OR + (kind='finding' AND state IN ('confirmed','dismissed')) OR + (kind='hint' AND state IN ('active','consumed')) + ); + END IF; +END $$; CREATE INDEX IF NOT EXISTS idx_expnodes_part ON exploration_nodes(exploration_id, kind); CREATE INDEX IF NOT EXISTS idx_expnodes_frontier ON exploration_nodes(exploration_id, priority DESC) WHERE kind='intent' AND state='open'; @@ -179,12 +200,14 @@ CREATE TABLE IF NOT EXISTS activity ( is_error BOOLEAN NOT NULL DEFAULT false, summary TEXT, detail TEXT, + metadata JSONB NOT NULL DEFAULT '{}', input_tokens INTEGER, output_tokens INTEGER, cache_read_tokens INTEGER, cache_write_tokens INTEGER, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +ALTER TABLE activity ADD COLUMN IF NOT EXISTS metadata JSONB NOT NULL DEFAULT '{}'; CREATE INDEX IF NOT EXISTS idx_act_node ON activity(exploration_id, node_id, id); CREATE INDEX IF NOT EXISTS idx_act_since ON activity(exploration_id, id); -- Main/Plan history pages filter by worker (both carry NULL node_id, so idx_act_node @@ -256,6 +279,8 @@ CREATE TABLE IF NOT EXISTS tasks ( paused BOOLEAN NOT NULL DEFAULT false, queued BOOLEAN NOT NULL DEFAULT false, llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, + active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, + llm_chain_revision BIGINT NOT NULL DEFAULT 0, company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL, parent_ref TEXT, timeout_seconds INTEGER NOT NULL DEFAULT 0, @@ -276,6 +301,57 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300; -- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。 ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_chain_revision BIGINT NOT NULL DEFAULT 0; + +-- Direct, read-only task context inheritance. Relations are intentionally not +-- recursive: a task sees only the source tasks explicitly chosen at creation. +CREATE TABLE IF NOT EXISTS task_relations ( + task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + source_task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (task_id, source_task_id), + CONSTRAINT ck_task_relation_not_self CHECK (task_id <> source_task_id) +); +CREATE INDEX IF NOT EXISTS idx_task_relations_source ON task_relations(source_task_id); + +-- Ordered task-level LLM failover chain. A quota-exhausted entry is skipped +-- until the user saves/resets the chain, which clears all failure state. +CREATE TABLE IF NOT EXISTS task_llm_profiles ( + task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + profile_id BIGINT NOT NULL REFERENCES llm_profiles(id) ON DELETE CASCADE, + position INTEGER NOT NULL CHECK (position >= 0), + status TEXT NOT NULL DEFAULT 'ready' + CHECK (status IN ('ready','quota_exhausted')), + last_error TEXT, + exhausted_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (task_id, profile_id), + UNIQUE (task_id, position) +); +CREATE INDEX IF NOT EXISTS idx_task_llm_profiles_order ON task_llm_profiles(task_id, position); +CREATE INDEX IF NOT EXISTS idx_task_llm_profiles_profile ON task_llm_profiles(profile_id, task_id); +CREATE INDEX IF NOT EXISTS idx_tasks_llm_profile ON tasks(llm_profile_id) WHERE llm_profile_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_tasks_active_llm_profile ON tasks(active_llm_profile_id) WHERE active_llm_profile_id IS NOT NULL; +DROP TRIGGER IF EXISTS trg_task_llm_profiles_upd ON task_llm_profiles; +CREATE TRIGGER trg_task_llm_profiles_upd BEFORE UPDATE ON task_llm_profiles + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + +-- One-time-compatible backfill: old pinned tasks become one-entry chains. A user +-- can still clear the chain later because the update path also clears the legacy +-- llm_profile_id column, preventing this block from re-adding it on restart. +INSERT INTO task_llm_profiles(task_id, profile_id, position) +SELECT t.id, t.llm_profile_id, 0 +FROM tasks t +WHERE t.llm_profile_id IS NOT NULL + AND NOT EXISTS (SELECT 1 FROM task_llm_profiles x WHERE x.task_id=t.id) +ON CONFLICT DO NOTHING; +UPDATE tasks t +SET active_llm_profile_id = t.llm_profile_id +WHERE t.active_llm_profile_id IS NULL + AND t.llm_profile_id IS NOT NULL + AND EXISTS (SELECT 1 FROM task_llm_profiles x WHERE x.task_id=t.id AND x.profile_id=t.llm_profile_id); -- 任务测试范围(资产覆盖度的分母 + 授权边界)。 -- 自动填(source='auto'):insertAssets 顶层按 worker 显式插入的资产类型加保守范围 @@ -335,6 +411,7 @@ ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_merge_mode TEXT NOT NUL ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_max_parallel INTEGER NOT NULL DEFAULT 5; -- per-agent LLM 绑定(agent 级默认模型):列自初版即在上方 CREATE 中,此 ALTER 仅为极旧库兜底(幂等)。 ALTER TABLE agents ADD COLUMN IF NOT EXISTS llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL; +CREATE INDEX IF NOT EXISTS idx_agents_llm_profile ON agents(llm_profile_id) WHERE llm_profile_id IS NOT NULL; DROP TRIGGER IF EXISTS trg_agents_upd ON agents; CREATE TRIGGER trg_agents_upd BEFORE UPDATE ON agents FOR EACH ROW EXECUTE FUNCTION set_updated_at(); @@ -454,6 +531,7 @@ CREATE TABLE IF NOT EXISTS conversations ( created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL; DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations; CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations FOR EACH ROW EXECUTE FUNCTION set_updated_at(); diff --git a/db/task_assets_context.go b/db/task_assets_context.go new file mode 100644 index 0000000..722cb36 --- /dev/null +++ b/db/task_assets_context.go @@ -0,0 +1,215 @@ +package db + +import ( + "net/url" + "sort" + "strconv" + "strings" +) + +// directTaskContextCTE is the current task plus exactly its explicitly related +// source tasks. It is deliberately non-recursive. +const directTaskContextCTE = ` +context_tasks AS ( + SELECT t.id AS task_id, t.exploration_id + FROM tasks t + WHERE t.id=$1 AND t.deleted_at IS NULL + UNION ALL + SELECT source.id, source.exploration_id + FROM task_relations relation + JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL + WHERE relation.task_id=$1 +)` + +const contextCoverageCTE = directTaskContextCTE + `, +target AS ( + SELECT DISTINCT a.id, a.type, + COALESCE(a.url, a.domain, a.ip, a.app_name, a.root_domain, '') AS label + FROM assets a + JOIN task_scope ts ON ( + (ts.kind='company' AND a.company_id = ts.company_id) + OR (ts.kind='root_domain' AND a.root_domain = ts.domain) + OR (ts.kind='subdomain' AND a.domain = ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + ) + JOIN context_tasks ctx ON ctx.task_id=ts.task_id + UNION + SELECT a.id, a.type, + COALESCE(a.url, a.domain, a.ip, a.app_name, a.root_domain, '') AS label + FROM assets a + JOIN exploration_anchors ea ON ea.asset_id=a.id + JOIN exploration_nodes en ON en.id=ea.node_id + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id +), +tested AS ( + SELECT DISTINCT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes en ON en.id=ea.node_id AND en.kind='fact' + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id +)` + +// ListTaskScopeWithSources returns the current task's scope followed by the +// scopes of its direct source tasks. TaskScope.TaskID preserves provenance. +func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) { + rows, err := s.db.Query(`WITH `+directTaskContextCTE+` +SELECT ts.id, ts.task_id, ts.kind, COALESCE(ts.company_id,0), COALESCE(ts.domain,''), + COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'') +FROM task_scope ts +JOIN context_tasks ctx ON ctx.task_id=ts.task_id +ORDER BY CASE WHEN ts.task_id=$1 THEN 0 ELSE 1 END, ts.id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []TaskScope{} + for rows.Next() { + var scope TaskScope + var companyID int64 + if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.Domain, &scope.Net, &scope.Source, &scope.Reason); err != nil { + return nil, err + } + if companyID > 0 { + scope.CompanyID = &companyID + } + out = append(out, scope) + } + return out, rows.Err() +} + +// TaskCoverageWithSources computes one coverage view over the union of the +// current task and its direct sources: source scopes and anchored assets extend +// the denominator, while fact anchors count as tested. No row is copied. +func (s *AssetStore) TaskCoverageWithSources(taskID int64) (*Coverage, error) { + cov := &Coverage{ByType: []CoverageByType{}} + _ = s.db.QueryRow(`WITH `+directTaskContextCTE+` +SELECT count(*) FROM task_scope ts JOIN context_tasks ctx ON ctx.task_id=ts.task_id`, taskID).Scan(&cov.ScopeRows) + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` +SELECT target.type, count(*) AS total, + count(*) FILTER (WHERE target.id IN (SELECT asset_id FROM tested)) AS tested +FROM target GROUP BY target.type ORDER BY target.type`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var byType CoverageByType + if err := rows.Scan(&byType.Type, &byType.Total, &byType.Tested); err != nil { + return nil, err + } + cov.ByType = append(cov.ByType, byType) + cov.Denominator += byType.Total + cov.Tested += byType.Tested + } + if err := rows.Err(); err != nil { + return nil, err + } + if cov.Denominator > 0 { + pct := float64(cov.Tested) / float64(cov.Denominator) + cov.Pct = &pct + } + return cov, nil +} + +// ListUntestedAssetsWithSources is the direct-source-aware backlog query used +// by inherited tasks. Source fact anchors remove assets from the backlog. +func (s *AssetStore) ListUntestedAssetsWithSources(taskID int64, typ string, limit, offset int) ([]CoverageAsset, int, error) { + if limit <= 0 { + limit = 10 + } + if offset < 0 { + offset = 0 + } + typeFilter := "" + args := []any{taskID} + if typ != "" { + typeFilter = " AND target.type = $2" + args = append(args, typ) + } + var total int + if err := s.db.QueryRow(`WITH `+contextCoverageCTE+` +SELECT count(*) FROM target WHERE target.id NOT IN (SELECT asset_id FROM tested)`+typeFilter, args...).Scan(&total); err != nil { + return nil, 0, err + } + pageArgs := append(append([]any{}, args...), limit, offset) + limitPosition := strconv.Itoa(len(args) + 1) + offsetPosition := strconv.Itoa(len(args) + 2) + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` +SELECT target.id, target.type, target.label FROM target +WHERE target.id NOT IN (SELECT asset_id FROM tested)`+typeFilter+` +ORDER BY target.id LIMIT $`+limitPosition+` OFFSET $`+offsetPosition, pageArgs...) + if err != nil { + return nil, total, err + } + defer rows.Close() + out := []CoverageAsset{} + for rows.Next() { + var asset CoverageAsset + if err := rows.Scan(&asset.ID, &asset.Type, &asset.Label); err != nil { + return nil, total, err + } + out = append(out, asset) + } + return out, total, rows.Err() +} + +// HostsByTaskWithSources resolves exact HTTP host candidates from assets that +// are attached to, anchored by, or in scope for the current task or a direct +// source. Traffic remains global and is not copied. This read helper must not be +// used for destructive task cleanup; HostsByTask intentionally retains that +// narrower, task-owned behavior. +func (s *AssetStore) HostsByTaskWithSources(taskID int64) ([]string, error) { + rows, err := s.db.Query(`WITH `+directTaskContextCTE+`, +context_assets AS ( + SELECT DISTINCT a.id + FROM assets a + WHERE EXISTS (SELECT 1 FROM context_tasks ctx WHERE ctx.task_id=ANY(a.task_ids)) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes en ON en.id=ea.node_id + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id + UNION + SELECT DISTINCT a.id + FROM assets a + JOIN task_scope ts ON ( + (ts.kind='company' AND a.company_id=ts.company_id) + OR (ts.kind='root_domain' AND a.root_domain=ts.domain) + OR (ts.kind='subdomain' AND a.domain=ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + ) + JOIN context_tasks ctx ON ctx.task_id=ts.task_id +) +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,'') +FROM assets a JOIN context_assets ctx ON ctx.id=a.id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + hosts := map[string]struct{}{} + add := func(host string) { + host = strings.TrimSpace(strings.ToLower(host)) + if host != "" { + hosts[host] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + if err := rows.Scan(&domain, &ip, &rawURL); err != nil { + return nil, err + } + add(domain) + add(ip) + if parsed, err := url.Parse(rawURL); err == nil { + add(parsed.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(hosts)) + for host := range hosts { + out = append(out, host) + } + sort.Strings(out) + return out, nil +} diff --git a/db/task_context.go b/db/task_context.go new file mode 100644 index 0000000..4927b02 --- /dev/null +++ b/db/task_context.go @@ -0,0 +1,388 @@ +package db + +import ( + "database/sql" + "fmt" + "strings" + "time" + "unicode/utf8" +) + +const IntentBlockedLLMQuota = "llm_quota_exhausted" + +// TaskLLMProfile is one ordered entry in a task's explicit failover chain. +type TaskLLMProfile struct { + ProfileID int64 `json:"profile_id"` + Position int `json:"position"` + Status string `json:"status"` + LastError string `json:"last_error,omitempty"` + ExhaustedAt *time.Time `json:"exhausted_at,omitempty"` +} + +// TaskSource identifies one directly related task and its exploration. +type TaskSource struct { + TaskID int64 + ExplorationID int64 + Description string + Goal string + Status string +} + +// TaskLLMTransition reports the shared task-level result of marking one profile +// quota-exhausted. NextProfileID is nil when the explicit chain is exhausted. +type TaskLLMTransition struct { + PreviousProfileID int64 + NextProfileID *int64 + ChainExhausted bool + Advanced bool + // Stale means the chain was replaced after the failing call selected its + // provider. The caller may retry a pre-stream request against the new chain, + // but must not report or persist a transition for this result. + Stale bool +} + +func (d *DB) hydrateTaskContext(t *Task) error { + if t == nil { + return nil + } + // A nil cursor on a non-empty chain is a persisted end-of-chain marker. Do not + // repair it from an earlier ready entry: the user may have manually selected a + // profile in the middle and legitimately exhausted every candidate after it. + legacyProfileID, activeProfileID, revision, chain, err := d.taskLLMContext(t.ID) + if err != nil { + return err + } + t.LLMProfileID = legacyProfileID + t.ActiveLLMProfileID = activeProfileID + t.LLMChainRevision = revision + sources, err := d.TaskSourceIDs(t.ID) + if err != nil { + return err + } + t.SourceTaskIDs = sources + t.LLMProfileIDs = make([]int64, 0, len(chain)) + t.LLMFailoverState = "default" + var latest *TaskLLMProfile + activeReady := false + for i := range chain { + entry := chain[i] + t.LLMProfileIDs = append(t.LLMProfileIDs, entry.ProfileID) + if t.ActiveLLMProfileID != nil && entry.ProfileID == *t.ActiveLLMProfileID && entry.Status == "ready" { + activeReady = true + } + if entry.ExhaustedAt != nil && (latest == nil || entry.ExhaustedAt.After(*latest.ExhaustedAt)) { + copy := entry + latest = © + } + } + if len(chain) > 0 { + if activeReady { + t.LLMFailoverState = "ready" + } else { + t.LLMFailoverState = "chain_exhausted" + } + } + if latest != nil { + t.LLMFailoverReason = latest.LastError + } + return nil +} + +// taskLLMContext reads the compatibility profile, current cursor, and ordered +// chain in one statement. Runtime chain edits commit atomically, and one SQL +// statement gives hydration one matching snapshot instead of a transient mix of +// an old task cursor and a newly replaced chain (or vice versa). +func (d *DB) taskLLMContext(taskID int64) (*int64, *int64, int64, []TaskLLMProfile, error) { + rows, err := d.Query(` +SELECT t.llm_profile_id, t.active_llm_profile_id, t.llm_chain_revision, + p.profile_id, p.position, p.status, COALESCE(p.last_error,''), p.exhausted_at +FROM tasks t +LEFT JOIN task_llm_profiles p ON p.task_id=t.id +WHERE t.id=$1 +ORDER BY p.position NULLS LAST`, taskID) + if err != nil { + return nil, nil, 0, nil, err + } + defer rows.Close() + var ( + legacyProfileID *int64 + activeProfileID *int64 + revision int64 + chain []TaskLLMProfile + found bool + ) + for rows.Next() { + var legacy, active, profileID, position sql.NullInt64 + var rowRevision int64 + var status, lastError sql.NullString + var exhaustedAt sql.NullTime + if err := rows.Scan(&legacy, &active, &rowRevision, &profileID, &position, &status, &lastError, &exhaustedAt); err != nil { + return nil, nil, 0, nil, err + } + if !found { + found = true + revision = rowRevision + if legacy.Valid { + id := legacy.Int64 + legacyProfileID = &id + } + if active.Valid { + id := active.Int64 + activeProfileID = &id + } + } + if !profileID.Valid { + continue + } + entry := TaskLLMProfile{ + ProfileID: profileID.Int64, + Position: int(position.Int64), + Status: status.String, + LastError: lastError.String, + } + if exhaustedAt.Valid { + ts := exhaustedAt.Time + entry.ExhaustedAt = &ts + } + chain = append(chain, entry) + } + if err := rows.Err(); err != nil { + return nil, nil, 0, nil, err + } + if !found { + return nil, nil, 0, nil, sql.ErrNoRows + } + return legacyProfileID, activeProfileID, revision, chain, nil +} + +func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) { + rows, err := d.Query(`SELECT source_task_id FROM task_relations WHERE task_id=$1 ORDER BY created_at, source_task_id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + out = append(out, id) + } + return out, rows.Err() +} + +func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) { + rows, err := d.Query(` +SELECT t.id, t.exploration_id, t.description, t.goal, t.status +FROM task_relations r +JOIN tasks t ON t.id=r.source_task_id AND t.deleted_at IS NULL +WHERE r.task_id=$1 +ORDER BY r.created_at, r.source_task_id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []TaskSource + for rows.Next() { + var source TaskSource + if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil { + return nil, err + } + out = append(out, source) + } + return out, rows.Err() +} + +func (d *DB) TaskLLMProfiles(taskID int64) ([]TaskLLMProfile, error) { + rows, err := d.Query(`SELECT profile_id, position, status, COALESCE(last_error,''), exhausted_at +FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []TaskLLMProfile + for rows.Next() { + var entry TaskLLMProfile + if err := rows.Scan(&entry.ProfileID, &entry.Position, &entry.Status, &entry.LastError, &entry.ExhaustedAt); err != nil { + return nil, err + } + out = append(out, entry) + } + return out, rows.Err() +} + +// ReplaceTaskLLMProfiles atomically replaces and resets the explicit task chain. +// activeProfileID=0 selects the first entry. An empty list restores the existing +// agent-binding/global fallback behavior. +func (d *DB) ReplaceTaskLLMProfiles(taskID int64, profileIDs []int64, activeProfileID int64) error { + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + // Keep the task -> profile lock order shared by all task LLM mutations and + // DeleteProfile. Inserts below may take KEY SHARE locks on llm_profiles for + // their foreign keys, so the task row must be locked before any of them. + var status string + if err := tx.QueryRow(`SELECT status FROM tasks WHERE id=$1 AND deleted_at IS NULL FOR UPDATE`, taskID).Scan(&status); err != nil { + if err == sql.ErrNoRows { + return fmt.Errorf("task %d not found", taskID) + } + return err + } + if IsTerminal(status) { + return fmt.Errorf("terminal task cannot change LLM profiles") + } + if _, err := tx.Exec(`DELETE FROM task_llm_profiles WHERE task_id=$1`, taskID); err != nil { + return err + } + if err := insertTaskLLMProfiles(tx, taskID, profileIDs); err != nil { + return err + } + if len(profileIDs) == 0 { + activeProfileID = 0 + } else if activeProfileID == 0 { + activeProfileID = profileIDs[0] + } + if activeProfileID > 0 { + found := false + for _, id := range profileIDs { + found = found || id == activeProfileID + } + if !found { + return fmt.Errorf("active LLM profile %d is not in the task chain", activeProfileID) + } + } + var active any + if activeProfileID > 0 { + active = activeProfileID + } + if _, err := tx.Exec(`UPDATE tasks +SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, taskID, active); err != nil { + return err + } + return tx.Commit() +} + +// MarkTaskLLMProfileQuotaExhausted advances the shared task cursor once. A late +// in-flight error from an older profile records that entry as exhausted but does +// not advance past the profile another call already selected. +func (d *DB) MarkTaskLLMProfileQuotaExhausted(taskID, profileID int64, reason string) (TaskLLMTransition, error) { + return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, 0, false, reason) +} + +// MarkTaskLLMProfileQuotaExhaustedAtRevision applies a provider failure only +// when it belongs to the chain snapshot used to start that call. +func (d *DB) MarkTaskLLMProfileQuotaExhaustedAtRevision(taskID, profileID, revision int64, reason string) (TaskLLMTransition, error) { + return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, revision, true, reason) +} + +func (d *DB) markTaskLLMProfileQuotaExhausted(taskID, profileID, revision int64, checkRevision bool, reason string) (TaskLLMTransition, error) { + var out TaskLLMTransition + tx, err := d.Begin() + if err != nil { + return out, err + } + defer tx.Rollback() + var ( + active sql.NullInt64 + currentRevision int64 + ) + if err := tx.QueryRow(`SELECT active_llm_profile_id, llm_chain_revision FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&active, ¤tRevision); err != nil { + return out, err + } + out.PreviousProfileID = profileID + if checkRevision && revision != currentRevision { + out.Stale = true + return out, tx.Commit() + } + var ( + position int + entryStatus string + ) + if err := tx.QueryRow(`SELECT position, status FROM task_llm_profiles WHERE task_id=$1 AND profile_id=$2`, taskID, profileID). + Scan(&position, &entryStatus); err != nil { + if err == sql.ErrNoRows { + // The chain was edited while this request was in flight. Its provider + // error remains valid for the caller, but it must not mutate the new chain. + return out, tx.Commit() + } + return out, err + } + // Once the cursor reached the end, late failures must be idempotent. This also + // prevents an older in-flight request from reviving a ready entry before a + // manually selected starting position. + if !active.Valid { + out.ChainExhausted = true + return out, tx.Commit() + } + reason = truncateUTF8(strings.TrimSpace(reason), 1000) + if entryStatus != "quota_exhausted" { + if _, err := tx.Exec(`UPDATE task_llm_profiles +SET status='quota_exhausted', last_error=$3, exhausted_at=now() +WHERE task_id=$1 AND profile_id=$2`, taskID, profileID, reason); err != nil { + return out, err + } + } + if active.Valid && active.Int64 != profileID { + next := active.Int64 + out.NextProfileID = &next + return out, tx.Commit() + } + var next int64 + err = tx.QueryRow(`SELECT profile_id FROM task_llm_profiles +WHERE task_id=$1 AND position>$2 AND status='ready' +ORDER BY position LIMIT 1`, taskID, position).Scan(&next) + switch err { + case nil: + out.Advanced = true + out.NextProfileID = &next + if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=$2, llm_profile_id=$2 WHERE id=$1`, taskID, next); err != nil { + return out, err + } + case sql.ErrNoRows: + out.Advanced = true + out.ChainExhausted = true + if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=NULL, llm_profile_id=NULL WHERE id=$1`, taskID); err != nil { + return out, err + } + default: + return out, err + } + return out, tx.Commit() +} + +func truncateUTF8(value string, maxBytes int) string { + value = strings.ToValidUTF8(value, "\uFFFD") + if maxBytes <= 0 { + return "" + } + if len(value) <= maxBytes { + return value + } + end := maxBytes + for end > 0 && !utf8.ValidString(value[:end]) { + end-- + } + return value[:end] +} + +func (s *ExplorationStore) SetIntentBlockedReason(id int64, reason string) error { + _, err := s.db.Exec(`UPDATE exploration_nodes +SET state='blocked', blocked_reason=NULLIF($1,''), completed_at=now() +WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, reason, id, s.expID) + return err +} + +func (s *ExplorationStore) ReopenIntentsByBlockedReason(reason string) (int64, error) { + res, err := s.db.Exec(`UPDATE exploration_nodes +SET state='open', blocked_reason=NULL, completed_at=NULL +WHERE exploration_id=$1 AND kind='intent' AND state='blocked' AND blocked_reason=$2`, s.expID, reason) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + return n, nil +} diff --git a/db/task_context_lock_test.go b/db/task_context_lock_test.go new file mode 100644 index 0000000..06b455d --- /dev/null +++ b/db/task_context_lock_test.go @@ -0,0 +1,527 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" +) + +func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + first, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-first-%d", suffix), Format: "openai", Model: "first", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + second, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-second-%d", suffix), Format: "openai", Model: "second", APIKey: "test-key", + }) + if err != nil { + _ = d.DeleteProfile(first) + t.Fatal(err) + } + task, err := d.CreateTaskWithOptions("LLM lock order", "verify concurrent mutation locks", TaskCreateOptions{ + LLMProfileIDs: []int64{first, second}, + }) + if err != nil { + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + }) + + replaceDB, replacePID := openSingleConnectionTestDB(t, dsn) + replaceBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer replaceBlocker.Rollback() + if _, err := replaceBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + replaceDone := make(chan error, 1) + go func() { + replaceDone <- replaceDB.ReplaceTaskLLMProfiles(task.ID, []int64{second, first}, second) + }() + if err := waitForBackendBlock(d, replacePID, replaceDone); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, first, second) + if err := replaceBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(replaceDone); err != nil { + t.Fatalf("replace chain: %v", err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer deleteBlocker.Rollback() + if _, err := deleteBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + deleteDone := make(chan error, 1) + go func() { + deleteDone <- deleteDB.DeleteProfile(second) + }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, second) + if err := deleteBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != first { + t.Fatalf("deleting the active profile did not select its successor: %+v", got) + } +} + +func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-lock-profile-%d", suffix), Format: "openai", Model: "agent-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("lock_agent_%d", suffix), "lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + if err := d.SetAgentLLMProfile(agent.Key, &profileID); err != nil { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete agent profile: %v", err) + } + got, err := d.GetAgentByKey(agent.Key) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("agent binding was not cleared: agent=%+v err=%v", got, err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-lock-profile-%d", suffix), Format: "openai", Model: "conversation-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "lock conversation", &profileID) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete conversation profile: %v", err) + } + got, err := d.GetConversation(conversation.ID) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("conversation binding was not cleared: conversation=%+v err=%v", got, err) + } + }) +} + +func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-write-lock-profile-%d", suffix), Format: "openai", Model: "agent-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("write_lock_agent_%d", suffix), "write lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.SetAgentLLMProfile(agent.Key, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind agent profile: %v", err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-write-lock-profile-%d", suffix), Format: "openai", Model: "conversation-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "write lock conversation", nil) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.UpdateConversationProfile(conversation.ID, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind conversation profile: %v", err) + } + }) +} + +func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-create-race-profile-%d", suffix), Format: "openai", Model: "conversation-create-race", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + title := fmt.Sprintf("conversation create race %d", suffix) + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM conversations WHERE title=$1`, title) + _ = d.DeleteProfile(profileID) + }) + + profileBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer profileBlocker.Rollback() + if _, err := profileBlocker.Exec(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE`, profileID); err != nil { + t.Fatal(err) + } + + createDB, createPID := openSingleConnectionTestDB(t, dsn) + type createResult struct { + conversation *Conversation + err error + } + createDone := make(chan createResult, 1) + createStatus := make(chan error, 1) + go func() { + conversation, err := createDB.CreateConversation("planner", title, &profileID) + createDone <- createResult{conversation: conversation, err: err} + createStatus <- err + }() + if err := waitForBackendBlock(d, createPID, createStatus); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + if err := profileBlocker.Rollback(); err != nil { + t.Fatal(err) + } + + var created createResult + select { + case created = <-createDone: + case <-time.After(12 * time.Second): + t.Fatal("timed out waiting for conversation creation") + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile during conversation creation: %v", err) + } + var persisted int + if err := d.QueryRow(`SELECT count(*) FROM conversations WHERE title=$1`, title).Scan(&persisted); err != nil { + t.Fatal(err) + } + if created.err == nil { + if created.conversation == nil || persisted != 1 { + t.Fatalf("successful creation was not committed atomically: conversation=%+v count=%d", created.conversation, persisted) + } + } else if persisted != 0 { + t.Fatalf("failed creation left a partial conversation row: err=%v count=%d", created.err, persisted) + } +} + +func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("late-reference-profile-%d", suffix), Format: "openai", Model: "late-reference", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + task, err := d.CreateTask("late profile reference", "exercise delete retry", nil, 0, 0) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(profileID) + }) + + // Keep the new reference uncommitted while deletion takes its initial + // READ COMMITTED snapshot. The task is therefore absent from the first lock + // set, while its FK KEY SHARE lock makes deletion wait at the profile row. + referenceTx, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer referenceTx.Rollback() + if _, err := referenceTx.Exec(`UPDATE tasks +SET llm_profile_id=$2, active_llm_profile_id=$2 +WHERE id=$1`, task.ID, profileID); err != nil { + t.Fatal(err) + } + if _, err := referenceTx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position) +VALUES ($1,$2,0)`, task.ID, profileID); err != nil { + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + referenceTx.Rollback() + t.Fatal(err) + } + if err := referenceTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile after late reference: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMChainRevision != 1 { + t.Fatalf("late task reference was not handled by a locked retry: %+v", got) + } +} + +func openSingleConnectionTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForBackendBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("mutation returned before reaching the expected reference-row lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func assertProfilesUnlocked(t *testing.T, d *DB, profileIDs ...int64) { + t.Helper() + probe, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer probe.Rollback() + for _, profileID := range profileIDs { + var lockedID int64 + if err := probe.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE NOWAIT`, profileID).Scan(&lockedID); err != nil { + t.Fatalf("profile %d was locked before the task row: %v", profileID, err) + } + } +} + +func waitForMutationResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for task LLM mutation") + } +} diff --git a/db/task_context_unit_test.go b/db/task_context_unit_test.go new file mode 100644 index 0000000..519f304 --- /dev/null +++ b/db/task_context_unit_test.go @@ -0,0 +1,27 @@ +package db + +import ( + "strings" + "testing" + "unicode/utf8" +) + +func TestTruncateUTF8PreservesValidEncoding(t *testing.T) { + t.Parallel() + input := strings.Repeat("额度不足", 400) + got := truncateUTF8(input, 1000) + if len(got) > 1000 { + t.Fatalf("truncated value has %d bytes, want at most 1000", len(got)) + } + if !utf8.ValidString(got) { + t.Fatalf("truncated value is not valid UTF-8: %q", got[len(got)-8:]) + } +} + +func TestTruncateUTF8RepairsInvalidInput(t *testing.T) { + t.Parallel() + got := truncateUTF8("bad\xffvalue", 1000) + if !utf8.ValidString(got) { + t.Fatalf("repaired value is not valid UTF-8: %q", got) + } +} diff --git a/db/task_delete_concurrency_test.go b/db/task_delete_concurrency_test.go new file mode 100644 index 0000000..0d65191 --- /dev/null +++ b/db/task_delete_concurrency_test.go @@ -0,0 +1,210 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" + + _ "github.com/jackc/pgx/v5/stdlib" +) + +func TestDeleteTaskTrafficHostsUseOneLockedTransaction(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + t.Run("sharing committed before deletion lock is observed", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/concurrent-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + writer, _ := openTaskDeleteTestDB(t, dsn) + writerTx, err := writer.Begin() + if err != nil { + t.Fatal(err) + } + defer writerTx.Rollback() + if _, err := writerTx.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID); err != nil { + t.Fatal(err) + } + + deleter, deleterPID := openTaskDeleteTestDB(t, dsn) + var preparedHosts []string + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + preparedHosts = append([]string(nil), p.TrafficHosts...) + return nil + }) + deleteDone <- deleteErr + }() + + if err := waitForTaskDeleteBlock(d, deleterPID, deleteDone); err != nil { + t.Fatal(err) + } + if err := writerTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if containsDeleteHost(preparedHosts, host) { + t.Fatalf("newly shared host %q was selected for traffic deletion: %v", host, preparedHosts) + } + var remaining int + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE url=$1 AND $2=ANY(task_ids)`, serviceURL, second.ID).Scan(&remaining); err != nil { + t.Fatal(err) + } + if remaining != 1 { + t.Fatalf("concurrent owner's asset was not preserved: count=%d", remaining) + } + }) + + t.Run("sharing cannot enter between host resolution and commit", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/late-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + deleter, _ := openTaskDeleteTestDB(t, dsn) + prepared := make(chan TaskDeletePreparation, 1) + releasePrepare := make(chan struct{}) + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + prepared <- p + <-releasePrepare + return nil + }) + deleteDone <- deleteErr + }() + + var plan TaskDeletePreparation + select { + case plan = <-prepared: + case err := <-deleteDone: + close(releasePrepare) + t.Fatalf("delete returned before preparation: %v", err) + case <-time.After(5 * time.Second): + close(releasePrepare) + t.Fatal("timed out waiting for task delete preparation") + } + if !containsDeleteHost(plan.TrafficHosts, host) { + close(releasePrepare) + t.Fatalf("exclusive host %q missing from preparation: %v", host, plan.TrafficHosts) + } + + writer, writerPID := openTaskDeleteTestDB(t, dsn) + writerDone := make(chan error, 1) + go func() { + _, writeErr := writer.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID) + writerDone <- writeErr + }() + if err := waitForTaskDeleteBlock(d, writerPID, writerDone); err != nil { + close(releasePrepare) + t.Fatal(err) + } + + close(releasePrepare) + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(writerDone); err != nil { + t.Fatal(err) + } + }) +} + +func createTaskDeleteRaceFixture(t *testing.T, d *DB) (first, second *Task, host string, rootAssetID int64) { + t.Helper() + var err error + first, err = d.CreateTask("delete race owner", "delete safely", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err = d.CreateTask("delete race sharer", "preserve shared host", nil, 0, 0) + if err != nil { + _ = d.DeleteTask(first.ID) + t.Fatal(err) + } + host = fmt.Sprintf("task-delete-race-%d.example.test", first.ID) + rootAssetID, err = d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: host, TaskID: first.ID}) + if err != nil { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + t.Fatal(err) + } + return first, second, host, rootAssetID +} + +func openTaskDeleteTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForTaskDeleteBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("operation returned before reaching the task deletion lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func waitForTaskDeleteResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for concurrent task deletion operation") + } +} + +func containsDeleteHost(hosts []string, want string) bool { + for _, host := range hosts { + if host == want { + return true + } + } + return false +} diff --git a/db/task_scope.go b/db/task_scope.go index e3108eb..2c69dcd 100644 --- a/db/task_scope.go +++ b/db/task_scope.go @@ -347,23 +347,23 @@ func hostPortOf(n *CoverageGraphNode) (string, int) { return host, port } -// BuildCoverageGraph assembles the full coverage graph for a task: every in-scope -// asset (all types) plus the connector root domains / companies needed to link -// them, with a tested flag on each in-scope asset. The frontend does the folding -// and highlighting; this only returns nodes + derived containment edges. -func (s *AssetStore) BuildCoverageGraph(taskID, expID int64) (*CoverageGraphData, error) { +// BuildCoverageGraph assembles the full coverage graph for a task and its direct +// read-only sources: every in-scope asset plus connector root domains/companies, +// with current-or-source fact anchors reflected in Tested. The legacy expID +// argument is retained for API compatibility; the task registry is authoritative. +func (s *AssetStore) BuildCoverageGraph(taskID, _ int64) (*CoverageGraphData, error) { g := &CoverageGraphData{Nodes: []CoverageGraphNode{}, Edges: []CoverageGraphEdge{}} if taskID <= 0 { return g, nil } - rows, err := s.db.Query(`WITH `+covTargetCTE+` + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` SELECT a.id, a.type, COALESCE(a.company_id,0), COALESCE(a.domain,''), COALESCE(a.root_domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), COALESCE(a.port,0), COALESCE(a.service_type,''), COALESCE(a.app_name,''), COALESCE(a.page_title,''), COALESCE(a.status_code,0), (a.id IN (SELECT asset_id FROM tested)) AS tested FROM assets a JOIN target t ON t.id = a.id -ORDER BY a.id`, taskID, expID) +ORDER BY a.id`, taskID) if err != nil { return nil, err } diff --git a/db/task_source_limit_test.go b/db/task_source_limit_test.go new file mode 100644 index 0000000..01596c8 --- /dev/null +++ b/db/task_source_limit_test.go @@ -0,0 +1,20 @@ +package db + +import ( + "strings" + "testing" +) + +func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) { + sourceIDs := make([]int64, MaxTaskSourceCount+1) + for i := range sourceIDs { + sourceIDs[i] = int64(i + 1) + } + + // No database handle is needed: validation must run before Begin so an + // oversized request cannot consume a connection or create partial rows. + _, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{SourceTaskIDs: sourceIDs}) + if err == nil || !strings.Contains(err.Error(), "too many source tasks") { + t.Fatalf("expected source-count validation error, got %v", err) + } +} diff --git a/db/tasks.go b/db/tasks.go index 86dfac6..e5c4dc9 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -1,23 +1,34 @@ package db import ( + "database/sql" "encoding/json" + "fmt" + "strconv" "time" ) // Task is a row in the task registry (1:1 with an exploration). type Task struct { - ID int64 `json:"id"` - Description string `json:"description"` - Goal string `json:"goal"` - ExplorationID int64 `json:"exploration_id"` - Status string `json:"status"` - Paused bool `json:"paused"` - Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` - ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) - CreatedAt time.Time `json:"created_at"` - CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil + ID int64 `json:"id"` + Description string `json:"description"` + Goal string `json:"goal"` + ExplorationID int64 `json:"exploration_id"` + Status string `json:"status"` + Paused bool `json:"paused"` + Queued bool `json:"queued"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` + // Task-level ordered LLM chain. LLMProfileID remains the compatibility alias + // for ActiveLLMProfileID while older API clients still send one profile id. + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMChainRevision int64 `json:"-"` + LLMFailoverState string `json:"llm_failover_state,omitempty"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` + ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) + CreatedAt time.Time `json:"created_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil // 任务级超时(见 docs/任务级超时与收尾设计.md)。 TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时 FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行 @@ -27,6 +38,23 @@ type Task struct { PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` } +// TaskDeleteResult reports optional related-data cleanup performed in the same +// transaction as the task/exploration delete. +type TaskDeleteResult struct { + AssetsDeleted int64 + AssetsDetached int64 + FindingsDeleted int64 + LLMRecordsDeleted int64 +} + +// TaskDeletePreparation is produced inside the PostgreSQL deletion transaction +// after asset and anchor writers have been excluded. Prepare callbacks may use +// TrafficHosts to stage an external traffic deletion before PostgreSQL commits. +type TaskDeletePreparation struct { + ExplorationID int64 + TrafficHosts []string +} + // IsTerminal reports whether a task status is a terminal (finished) state. // 单一真源,替换散落各处的 done/failed 硬编码判定。 func IsTerminal(status string) bool { @@ -40,6 +68,12 @@ func IsTerminal(status string) bool { // 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min,防止把 planner 打爆。 const MinPlanHeartbeatSeconds = 600 +// MaxTaskSourceCount bounds the amount of live inherited context one task can +// pull into every planner/main-agent prompt. Inheritance is intentionally direct +// only; keeping the fan-in bounded also prevents a single create request from +// multiplying graph and asset-context queries without limit. +const MaxTaskSourceCount = 8 + func normalizeHeartbeat(sec int) int { if sec < MinPlanHeartbeatSeconds { return MinPlanHeartbeatSeconds @@ -48,6 +82,30 @@ func normalizeHeartbeat(sec int) int { } func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { + var ids []int64 + if llmProfileID != nil { + ids = []int64{*llmProfileID} + } + return d.CreateTaskWithOptions(description, goal, TaskCreateOptions{ + LLMProfileIDs: ids, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds, + }) +} + +// TaskCreateOptions contains the task data that must be committed atomically +// with the task/exploration row. +type TaskCreateOptions struct { + SourceTaskIDs []int64 + LLMProfileIDs []int64 + TimeoutSeconds int + PlanHeartbeatSeconds int +} + +// CreateTaskWithOptions creates an exploration, task, direct source relations, +// and the ordered task LLM chain in one transaction. +func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOptions) (*Task, error) { + if len(opts.SourceTaskIDs) > MaxTaskSourceCount { + return nil, fmt.Errorf("too many source tasks: got %d, maximum is %d", len(opts.SourceTaskIDs), MaxTaskSourceCount) + } tx, err := d.Begin() if err != nil { return nil, err @@ -72,24 +130,80 @@ INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, or VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil { return nil, err } - if timeoutSeconds < 0 { - timeoutSeconds = 0 + if opts.TimeoutSeconds < 0 { + opts.TimeoutSeconds = 0 + } + opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds) + var active *int64 + if len(opts.LLMProfileIDs) > 0 { + id := opts.LLMProfileIDs[0] + active = &id + } + t := &Task{ + Description: description, Goal: goal, ExplorationID: expID, + LLMProfileID: active, ActiveLLMProfileID: active, + LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...), + SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...), + TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds, } - planHeartbeatSeconds = normalizeHeartbeat(planHeartbeatSeconds) // <=0→300,<30→30 - t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds} if err := tx.QueryRow(` -INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6) -RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds) +VALUES ($1,$2,$3,$4,$4,$5,$6) +RETURNING id, status, paused, created_at`, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } + if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { + return nil, err + } + if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil { + return nil, err + } + if len(opts.LLMProfileIDs) == 0 { + t.LLMFailoverState = "default" + } else { + t.LLMFailoverState = "ready" + } return t, tx.Commit() } -const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at` +func insertTaskRelations(tx *sql.Tx, taskID int64, sourceIDs []int64) error { + seen := map[int64]bool{} + for _, sourceID := range sourceIDs { + if sourceID <= 0 || sourceID == taskID || seen[sourceID] { + return fmt.Errorf("invalid or duplicate source task id %d", sourceID) + } + seen[sourceID] = true + res, err := tx.Exec(`INSERT INTO task_relations(task_id, source_task_id) +SELECT $1, id FROM tasks WHERE id=$2 AND deleted_at IS NULL`, taskID, sourceID) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n != 1 { + return fmt.Errorf("source task %d not found", sourceID) + } + } + return nil +} + +func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { + seen := map[int64]bool{} + for position, profileID := range profileIDs { + if profileID <= 0 || seen[profileID] { + return fmt.Errorf("invalid or duplicate LLM profile id %d", profileID) + } + seen[profileID] = true + if _, err := tx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position) VALUES ($1,$2,$3)`, taskID, profileID, position); err != nil { + return err + } + } + return nil +} + +const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } return &t, nil @@ -116,7 +230,16 @@ func (d *DB) ListTasks() ([]*Task, error) { } out = append(out, t) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + rows.Close() + for _, task := range out { + if err := d.hydrateTaskContext(task); err != nil { + return nil, err + } + } + return out, nil } // GetTask returns one alive task (nil if not found/deleted). @@ -128,6 +251,9 @@ func (d *DB) GetTask(id int64) (*Task, error) { } return nil, err } + if err := d.hydrateTaskContext(t); err != nil { + return nil, err + } return t, nil } @@ -137,8 +263,6 @@ func (d *DB) SetPaused(id int64, paused bool) error { return err } -// SetQueued persists a task's queued flag (并发上限挂起态). A queued task is held -// (engine not started) until the concurrency reconciler promotes it into a free slot. func (d *DB) SetQueued(id int64, queued bool) error { _, err := d.Exec(`UPDATE tasks SET queued=$1 WHERE id=$2`, queued, id) return err @@ -192,28 +316,132 @@ UPDATE tasks return deadline, err } -// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/ -// activity via ON DELETE CASCADE). The task row is removed first so the -// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete. -// Global assets are NOT touched. +// DeleteTask preserves the historical behavior: remove the task and its +// exploration subgraph while retaining global assets. func (d *DB) DeleteTask(id int64) error { + _, err := d.DeleteTaskCascade(id, false, false) + return err +} + +// DeleteTaskCascade hard-deletes a task and its exploration subgraph +// (nodes/edges/activity via ON DELETE CASCADE). Optional standalone findings and +// assets owned only by this task are deleted in the same transaction; shared +// assets are retained and only have this task id detached. deleteLLMRecords is +// optional to preserve callers of the original two-option API. +func (d *DB) DeleteTaskCascade(id int64, deleteAssets, deleteFindings bool, deleteLLMRecords ...bool) (TaskDeleteResult, error) { + deleteRecords := len(deleteLLMRecords) > 0 && deleteLLMRecords[0] + return d.DeleteTaskCascadePrepared(id, deleteAssets, deleteFindings, deleteRecords, nil) +} + +// DeleteTaskCascadePrepared coordinates reversible external deletion with the +// PostgreSQL cascade. When prepare is non-nil, host ownership is resolved and +// prepare is invoked inside this transaction while asset and anchor writers are +// excluded. The locks remain held through commit, closing the window where a +// host could become shared after its traffic had already been staged. +// +// prepare must only stage reversible work. Any returned error or later database +// error rolls PostgreSQL back; the caller remains responsible for rolling back +// external work that its callback staged successfully. +func (d *DB) DeleteTaskCascadePrepared( + id int64, + deleteAssets, deleteFindings, deleteLLMRecords bool, + prepare func(TaskDeletePreparation) error, +) (TaskDeleteResult, error) { + var result TaskDeleteResult tx, err := d.Begin() if err != nil { - return err + return result, err } defer tx.Rollback() var expID int64 - if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil { - if err.Error() == "sql: no rows in result set" { - return nil // already gone + if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 FOR UPDATE`, id).Scan(&expID); err != nil { + if err == sql.ErrNoRows { + return result, nil // already gone } - return err + return result, err + } + + // SHARE ROW EXCLUSIVE conflicts with every INSERT/UPDATE/DELETE on these + // tables and with another coordinated deletion. Taking both in one fixed order + // prevents phantoms (a distinct asset row for the same host) as well as new + // ownership/anchor references until the deletion transaction commits. + if deleteAssets || prepare != nil { + if _, err := tx.Exec(`LOCK TABLE assets, exploration_anchors IN SHARE ROW EXCLUSIVE MODE`); err != nil { + return result, err + } + } + if prepare != nil { + hosts, err := hostsForTaskDeletion(tx, id, expID) + if err != nil { + return result, err + } + if err := prepare(TaskDeletePreparation{ + ExplorationID: expID, + TrafficHosts: hosts, + }); err != nil { + return result, err + } + } + if deleteAssets { + // Ownership is the union of explicit task_ids and this exploration's anchors. + // An asset is deletable only when no other live task references it through + // either mechanism. This covers legacy anchor-only seeds without destroying + // evidence anchored by another task. + res, err := tx.Exec(` +WITH candidate_assets AS ( + SELECT id FROM assets WHERE $1 = ANY(task_ids) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + WHERE n.exploration_id=$2 +), +deletable AS ( + SELECT a.id + FROM assets a + JOIN candidate_assets c ON c.id=a.id + WHERE NOT EXISTS ( + SELECT 1 FROM tasks t + WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids) + ) AND NOT EXISTS ( + SELECT 1 + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + JOIN tasks t ON t.exploration_id=n.exploration_id + WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL + ) +) +DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID) + if err != nil { + return result, err + } + result.AssetsDeleted, _ = res.RowsAffected() + + res, err = tx.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, id) + if err != nil { + return result, err + } + result.AssetsDetached, _ = res.RowsAffected() + } + if deleteFindings { + res, err := tx.Exec(`DELETE FROM findings WHERE task_id = $1`, id) + if err != nil { + return result, err + } + result.FindingsDeleted, _ = res.RowsAffected() + } + if deleteLLMRecords { + res, err := tx.Exec(`DELETE FROM llm_records WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(id, 10)) + if err != nil { + return result, err + } + result.LLMRecordsDeleted, _ = res.RowsAffected() } if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil { - return err + return result, err } if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil { - return err + return result, err } - return tx.Commit() + return result, tx.Commit() } diff --git a/db/tasks_test.go b/db/tasks_test.go index d407e03..655b06b 100644 --- a/db/tasks_test.go +++ b/db/tasks_test.go @@ -1,6 +1,11 @@ package db -import "testing" +import ( + "fmt" + "sync" + "testing" + "time" +) func TestTaskLifecycleAndDeleteCascade(t *testing.T) { d, err := Open(testDSN(t)) @@ -75,3 +80,406 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) { t.Fatalf("exploration should be deleted, got %d", exps) } } + +func TestTaskDeleteCascadeAssets(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + first, err := d.CreateTask("级联删除测试", "目标A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err := d.CreateTask("共享资产保留测试", "目标B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + exclusiveID, err := assets.UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("delete-%d.example.test", first.ID), + TaskID: first.ID, + }) + if err != nil { + t.Fatal(err) + } + sharedDomain := fmt.Sprintf("shared-%d.example.test", first.ID) + sharedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + if _, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: second.ID}); err != nil { + t.Fatal(err) + } + anchorOnlyDomain := fmt.Sprintf("anchor-only-%d.example.test", first.ID) + anchorOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: anchorOnlyDomain}) + if err != nil { + t.Fatal(err) + } + firstOrigin, err := d.Exploration(first.ExplorationID).OriginFactID() + if err != nil || firstOrigin == 0 { + t.Fatalf("first origin: id=%d err=%v", firstOrigin, err) + } + if err := d.Exploration(first.ExplorationID).Anchor(firstOrigin, anchorOnlyID); err != nil { + t.Fatal(err) + } + protectedDomain := fmt.Sprintf("other-anchor-%d.example.test", first.ID) + protectedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: protectedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + secondOrigin, err := d.Exploration(second.ExplorationID).OriginFactID() + if err != nil || secondOrigin == 0 { + t.Fatalf("second origin: id=%d err=%v", secondOrigin, err) + } + if err := d.Exploration(second.ExplorationID).Anchor(secondOrigin, protectedID); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + }) + + hosts, err := assets.HostsByTask(first.ID) + if err != nil { + t.Fatal(err) + } + hostSet := make(map[string]bool, len(hosts)) + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !hostSet[sharedDomain] { + t.Fatalf("task hosts missing cascade fixtures: %v", hosts) + } + deletableHosts, err := assets.HostsForTaskDeletion(first.ID, first.ExplorationID) + if err != nil { + t.Fatal(err) + } + deletableSet := make(map[string]bool, len(deletableHosts)) + for _, host := range deletableHosts { + deletableSet[host] = true + } + if !deletableSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !deletableSet[anchorOnlyDomain] { + t.Fatalf("exclusive or anchor-only cleanup host missing: %v", deletableHosts) + } + if deletableSet[sharedDomain] || deletableSet[protectedDomain] { + t.Fatalf("shared host was not protected from traffic cleanup: %v", deletableHosts) + } + + findingID, err := d.AddFinding(first.ID, 0, "__task_delete_cascade__", "", SeverityHigh, "summary", "evidence", "tester", []int64{exclusiveID}) + if err != nil { + t.Fatal(err) + } + + result, err := d.DeleteTaskCascade(first.ID, true, true) + if err != nil { + t.Fatal(err) + } + if result.AssetsDeleted < 2 || result.AssetsDetached < 2 { + t.Fatalf("unexpected asset cleanup result: %+v", result) + } + if result.FindingsDeleted != 1 { + t.Fatalf("unexpected finding cleanup result: %+v", result) + } + if finding, err := d.GetFinding(findingID); err != nil || finding != nil { + t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err) + } + remaining, err := assets.GetByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + if err != nil { + t.Fatal(err) + } + if len(remaining) != 2 { + t.Fatalf("only shared and other-anchored assets should remain, got %+v", remaining) + } + byID := make(map[int64]*Asset, len(remaining)) + for _, asset := range remaining { + byID[asset.ID] = asset + } + if shared := byID[sharedID]; shared == nil || len(shared.TaskIDs) != 1 || shared.TaskIDs[0] != second.ID { + t.Fatalf("deleted task should be detached from shared asset: %+v", shared) + } + if protected := byID[protectedID]; protected == nil || len(protected.TaskIDs) != 0 { + t.Fatalf("other task's anchor should preserve the asset after detaching ownership: %+v", protected) + } +} + +func TestTaskRelationsAndLLMFailoverChain(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileIDs := make([]int64, 0, 3) + for i := 0; i < 3; i++ { + id, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("task-chain-%d-%d", suffix, i), Format: "openai", + Model: fmt.Sprintf("model-%d", i), APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + profileIDs = append(profileIDs, id) + } + + sourceA, err := d.CreateTask("source A", "goal A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + sourceB, err := d.CreateTask("source B", "goal B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + child, err := d.CreateTaskWithOptions("child", "new goal", TaskCreateOptions{ + SourceTaskIDs: []int64{sourceA.ID, sourceB.ID}, + LLMProfileIDs: profileIDs, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(child.ID) + _ = d.DeleteTask(sourceA.ID) + _ = d.DeleteTask(sourceB.ID) + for _, id := range profileIDs { + _ = d.DeleteProfile(id) + } + }) + + got, err := d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceA.ID, sourceB.ID}) { + t.Fatalf("unexpected direct sources: %v", got.SourceTaskIDs) + } + if fmt.Sprint(got.LLMProfileIDs) != fmt.Sprint(profileIDs) || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("unexpected initial chain: %+v", got) + } + initialRevision := got.LLMChainRevision + if err := d.ReplaceTaskLLMProfiles(child.ID, profileIDs, profileIDs[0]); err != nil { + t.Fatal(err) + } + staleRevision, err := d.MarkTaskLLMProfileQuotaExhaustedAtRevision(child.ID, profileIDs[0], initialRevision, "late quota from prior generation") + if err != nil { + t.Fatal(err) + } + if !staleRevision.Stale || staleRevision.Advanced || staleRevision.ChainExhausted { + t.Fatalf("old chain revision changed replacement chain: %+v", staleRevision) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("replacement revision/cursor not preserved: %+v", got) + } + + transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota") + if err != nil { + t.Fatal(err) + } + if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] { + t.Fatalf("unexpected first transition: %+v", transition) + } + late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate") + if err != nil { + t.Fatal(err) + } + if late.Advanced || late.NextProfileID == nil || *late.NextProfileID != profileIDs[1] { + t.Fatalf("late failure advanced past current profile: %+v", late) + } + if _, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "quota_exceeded"); err != nil { + t.Fatal(err) + } + last, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "余额不足") + if err != nil { + t.Fatal(err) + } + if !last.Advanced || !last.ChainExhausted || last.NextProfileID != nil { + t.Fatalf("unexpected exhausted transition: %+v", last) + } + duplicateLast, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "late final duplicate") + if err != nil { + t.Fatal(err) + } + if duplicateLast.Advanced || !duplicateLast.ChainExhausted || duplicateLast.NextProfileID != nil { + t.Fatalf("duplicate final failure must be idempotent: %+v", duplicateLast) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("chain exhaustion was not persisted: %+v", got) + } + + // Two requests can observe the same last active profile. Only the transaction + // that clears the cursor may report an advance; the late one is idempotent. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2]}, profileIDs[2]); err != nil { + t.Fatal(err) + } + transitions := make(chan TaskLLMTransition, 2) + errs := make(chan error, 2) + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + transition, markErr := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "concurrent final quota") + if markErr != nil { + errs <- markErr + return + } + transitions <- transition + }() + } + wg.Wait() + close(errs) + close(transitions) + for markErr := range errs { + t.Fatal(markErr) + } + advanced := 0 + for transition := range transitions { + if transition.Advanced { + advanced++ + } + if !transition.ChainExhausted { + t.Fatalf("concurrent final transition must report exhausted chain: %+v", transition) + } + } + if advanced != 1 { + t.Fatalf("last profile advanced %d times, want exactly once", advanced) + } + + // Starting manually from the middle consumes only candidates after that cursor. + // Earlier ready profiles must not be revived by hydration or unrelated deletes. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil { + t.Fatal(err) + } + if step, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "middle quota"); err != nil || step.NextProfileID == nil || *step.NextProfileID != profileIDs[1] { + t.Fatalf("middle cursor did not advance to its successor: transition=%+v err=%v", step, err) + } + if end, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "tail quota"); err != nil || !end.ChainExhausted { + t.Fatalf("tail did not exhaust manual chain: transition=%+v err=%v", end, err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("hydration revived a profile before the manual cursor: %+v", got) + } + unrelatedID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("task-chain-unrelated-%d", suffix), Format: "openai", Model: "other", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + if err := d.DeleteProfile(unrelatedID); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("deleting an unrelated profile revived an exhausted chain: %+v", got) + } + + // A call that finishes after its profile was removed from the chain is stale: + // preserve the new cursor and surface the original provider error upstream. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[1]}, profileIDs[2]); err != nil { + t.Fatal(err) + } + stale, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "obsolete in-flight quota") + if err != nil { + t.Fatalf("removed in-flight profile returned an internal error: %v", err) + } + if stale.Advanced || stale.ChainExhausted || stale.NextProfileID != nil { + t.Fatalf("removed in-flight profile changed the replacement chain: %+v", stale) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[2] { + t.Fatalf("stale failure changed active replacement profile: %+v", got) + } + + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] || got.LLMFailoverState != "ready" || got.LLMFailoverReason != "" { + t.Fatalf("chain reset did not clear failure state: %+v", got) + } + + if err := d.DeleteProfile(profileIDs[0]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[1] { + t.Fatalf("deleting active profile did not select next ready entry: %+v", got) + } + if err := d.DeleteProfile(profileIDs[1]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMFailoverState != "default" { + t.Fatalf("deleting active chain tail must fall back instead of wrapping: %+v", got) + } + + if err := d.DeleteTask(sourceA.ID); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceB.ID}) { + t.Fatalf("source delete did not cascade only its relation: %v", got.SourceTaskIDs) + } +} + +func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + source, err := d.CreateTask("source", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(source.ID) }) + if _, err := d.CreateTaskWithOptions("bad", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID, source.ID}}); err == nil { + t.Fatal("duplicate source task ids should be rejected") + } + + task, err := d.CreateTask("terminal", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + if err := d.SetStatus(task.ID, "done"); err != nil { + t.Fatal(err) + } + if err := d.ReplaceTaskLLMProfiles(task.ID, nil, 0); err == nil { + t.Fatal("terminal task LLM edit should be rejected") + } +} diff --git a/llmrec/llmrec.go b/llmrec/llmrec.go index 37c256d..cf0be46 100644 --- a/llmrec/llmrec.go +++ b/llmrec/llmrec.go @@ -16,6 +16,27 @@ import ( "github.com/Autumn-27/artex/db" ) +type taskIDContextKey struct{} + +// WithTaskID attaches the owning task registry id to an LLM call. Session ids +// are based on exploration ids, which are not interchangeable with task ids. +func WithTaskID(ctx context.Context, taskID string) context.Context { + taskID = strings.TrimSpace(taskID) + if taskID == "" { + return ctx + } + return context.WithValue(ctx, taskIDContextKey{}, taskID) +} + +// TaskIDFrom returns the explicit task registry id attached by the task runtime. +func TaskIDFrom(ctx context.Context) string { + if ctx == nil { + return "" + } + taskID, _ := ctx.Value(taskIDContextKey{}).(string) + return strings.TrimSpace(taskID) +} + // Recorder wraps an llm.Provider and records every completion call. type Recorder struct { inner llm.Provider @@ -70,7 +91,13 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } start := time.Now() session := transcript.SessionIDFrom(ctx) - taskID, worker := parseSession(session) + parsedID, worker := parseSession(session) + taskID := TaskIDFrom(ctx) + if taskID == "" { + // Backward compatibility for non-task callers. For task calls, the task + // runtime always supplies the registry id explicitly. + taskID = parsedID + } // Serialize the request for storage. reqBody := serializeRequest(req) @@ -120,7 +147,9 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } } -// record persists one LLM call to PostgreSQL (async, best-effort). +// record persists one LLM call to PostgreSQL before the provider stream returns. +// Keeping the write inside the owning task operation means task deletion can +// drain calls and then remove records without a late async insert recreating one. func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, reqBody string, start time.Time, text, thinking string, usage llm.Usage, stopReason string, streamErr error) { latency := int(time.Since(start).Milliseconds()) status := "ok" @@ -169,11 +198,9 @@ func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, re ResponseBody: string(respBody), } - go func() { - if err := r.pg.InsertLLMRecord(rec); err != nil { - log.Printf("[llmrec] insert: %v", err) - } - }() + if err := r.pg.InsertLLMRecord(rec); err != nil { + log.Printf("[llmrec] insert: %v", err) + } } // serializeRequest builds a JSON representation of the completion request. diff --git a/llmrec/llmrec_test.go b/llmrec/llmrec_test.go new file mode 100644 index 0000000..e998684 --- /dev/null +++ b/llmrec/llmrec_test.go @@ -0,0 +1,29 @@ +package llmrec + +import ( + "context" + "testing" +) + +func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) { + ctx := WithTaskID(context.Background(), " 42 ") + if got := TaskIDFrom(ctx); got != "42" { + t.Fatalf("TaskIDFrom()=%q want 42", got) + } + if got := TaskIDFrom(WithTaskID(ctx, " ")); got != "42" { + t.Fatalf("blank task id should preserve parent context, got %q", got) + } + if got := TaskIDFrom(nil); got != "" { + t.Fatalf("nil context returned %q", got) + } +} + +func TestParseSessionFallbackIsExplorationScoped(t *testing.T) { + taskID, worker := parseSession("exp12-worker-i99") + if taskID != "12" || worker != "worker" { + t.Fatalf("parseSession()=(%q,%q)", taskID, worker) + } + if taskID, worker := parseSession("not-a-task"); taskID != "" || worker != "" { + t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker) + } +} diff --git a/server/assets.go b/server/assets.go index 62a794e..ac8dcba 100644 --- a/server/assets.go +++ b/server/assets.go @@ -217,7 +217,6 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) { // GET /api/assets // ===================================================================== -// defaultAssetPageSize is the page size used when the caller omits ?limit. const defaultAssetPageSize = 50 func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { @@ -230,8 +229,6 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { typ := q.Get("type") limit, _ := strconv.Atoi(q.Get("limit")) offset, _ := strconv.Atoi(q.Get("offset")) - // limit/offset are page controls, not a cap: `total` always reports the full - // match count so callers can page through everything. if limit <= 0 { limit = defaultAssetPageSize } @@ -351,8 +348,8 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) { Assets []struct { Type string `json:"type"` // root_domain / subdomain - Domain string `json:"domain"` - ICP string `json:"icp"` + Domain string `json:"domain"` + ICP string `json:"icp"` RecordType string `json:"record_type"` RecordValue []string `json:"record_value"` // ip diff --git a/server/chatupload.go b/server/chatupload.go index 7ac55e1..82ae080 100644 --- a/server/chatupload.go +++ b/server/chatupload.go @@ -45,9 +45,11 @@ type chatAttachment struct { // 文件先落这里,前端按返回的 abs 绝对路径写进任务描述) func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) { var sub string + taskScoped := false switch r.URL.Query().Get("scope") { case "task": sub = "tasks" + taskScoped = true case "session": sub = "sessions" case "staging": @@ -61,6 +63,17 @@ func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "非法 id") return } + if taskScoped { + if s.m.ResolveTask(id) == nil { + writeErr(w, 404, "task not found") + return + } + if !s.engine.beginTaskOperation(id) { + writeErr(w, http.StatusConflict, "任务正在删除,无法上传附件") + return + } + defer s.engine.decInflight(id) + } dir := filepath.Join(s.m.dir, sub, id, "uploads") if err := os.MkdirAll(dir, 0o755); err != nil { writeErr(w, 500, "建目录失败: "+err.Error()) diff --git a/server/core_test.go b/server/core_test.go index 49e9417..56b8a18 100644 --- a/server/core_test.go +++ b/server/core_test.go @@ -6,6 +6,8 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "os" + "path/filepath" "strconv" "testing" "time" @@ -92,15 +94,85 @@ func TestCoreTaskLifecyclePG(t *testing.T) { t.Fatalf("list tasks: %d", code) } - // delete → cascade removes the exploration subgraph - code, _ = doRetry("DELETE", "/api/tasks/"+id, nil) + // The task detail header consumes the top-level engine mode while some clients + // read the active-task snapshot. Keep both representations in sync. + code, out = doRetry("GET", "/api/stats?task="+id, nil) + if code != 200 { + t.Fatalf("task stats: %d (%v)", code, out) + } + activeTask, ok := out["active_task"].(map[string]any) + if !ok || activeTask["engine_mode"] != out["engine_mode"] { + t.Fatalf("engine mode mismatch: top=%v active_task=%v", out["engine_mode"], activeTask) + } + + // delete → cascade removes the exploration subgraph and selected related data + taskDir := filepath.Join(m.dir, "tasks", id) + if err := os.MkdirAll(taskDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(taskDir, "artifact.txt"), []byte("test"), 0o644); err != nil { + t.Fatal(err) + } + transcriptBase := "exp" + strconv.FormatInt(expID, 10) + "-main" + transcriptPath := filepath.Join(m.dir, "transcripts", transcriptBase+".jsonl") + sidechainPath := filepath.Join(m.dir, "transcripts", transcriptBase) + if err := os.MkdirAll(filepath.Join(sidechainPath, "subagents"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(transcriptPath, []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(sidechainPath, "subagents", "child.jsonl"), []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + nid, _ := strconv.ParseInt(id, 10, 64) + findingID, err := m.pg.AddFinding(nid, 0, "__core_task_delete__", "", db.SeverityHigh, "summary", "evidence", "tester", nil) + if err != nil { + t.Fatal(err) + } + if err := m.pg.EnsureLLMRecordsTable(); err != nil { + t.Fatal(err) + } + if err := m.pg.InsertLLMRecord(&db.LLMRecord{TaskID: id, SessionID: "delete-test", Status: "ok", RequestBody: "secret"}); err != nil { + t.Fatal(err) + } + // A non-canonical path still resolves to the one canonical task key for the + // barrier, workspace, in-memory registry and DB delete. + code, out = doRetry("DELETE", "/api/tasks/000"+id, map[string]bool{ + "delete_files": true, + "delete_findings": true, + "delete_llm_records": true, + }) if code != 200 { t.Fatalf("delete task: %d", code) } - nid, _ := strconv.ParseInt(id, 10, 64) + if deleted, _ := out["files_deleted"].(bool); !deleted { + t.Fatalf("expected files_deleted response, got %v", out) + } + if deleted, _ := out["findings_deleted"].(float64); deleted != 1 { + t.Fatalf("expected one deleted finding, got %v", out) + } + if deleted, _ := out["llm_records_deleted"].(float64); deleted != 1 { + t.Fatalf("expected one deleted LLM record, got %v", out) + } + if _, err := os.Stat(taskDir); !os.IsNotExist(err) { + t.Fatalf("task directory should be deleted, stat err=%v", err) + } + for _, path := range []string{transcriptPath, sidechainPath} { + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("task transcript should be deleted (%s), stat err=%v", path, err) + } + } if got, _ := m.pg.GetTask(nid); got != nil { t.Fatalf("task should be deleted from PG") } + if finding, err := m.pg.GetFinding(findingID); err != nil || finding != nil { + t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err) + } + var llmRecords int + if err := m.pg.QueryRow(`SELECT count(*) FROM llm_records WHERE task_id=$1`, id).Scan(&llmRecords); err != nil || llmRecords != 0 { + t.Fatalf("task LLM records should be deleted, count=%d err=%v", llmRecords, err) + } var n int m.pg.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, expID).Scan(&n) if n != 0 { diff --git a/server/dto.go b/server/dto.go index 1f7f0cd..9c320c0 100644 --- a/server/dto.go +++ b/server/dto.go @@ -30,21 +30,27 @@ func rawString(raw json.RawMessage) string { // ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status. type TaskDTO struct { - ID string `json:"id"` - ExplorationID int64 `json:"exploration_id"` - Description string `json:"description"` - Goal string `json:"goal"` - Status string `json:"status"` // created | running | paused | done | failed - CreatedAt string `json:"created_at"` - CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc) - CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished - CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished) - LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none) - Paused bool `json:"paused"` - Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption - GoalsTotal int `json:"goals_total"` - GoalsMet int `json:"goals_met"` - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default + ID string `json:"id"` + ExplorationID int64 `json:"exploration_id"` + Description string `json:"description"` + Goal string `json:"goal"` + Status string `json:"status"` // created | running | paused | done | failed + CreatedAt string `json:"created_at"` + CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc) + CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished + CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished) + LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none) + Paused bool `json:"paused"` + Queued bool `json:"queued"` + Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption + GoalsTotal int `json:"goals_total"` + GoalsMet int `json:"goals_met"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default + LLMProfileIDs []int64 `json:"llm_profile_ids"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMFailoverState string `json:"llm_failover_state"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []string `json:"source_task_ids"` } // TokenTotalDTO is a whole-task (all agents) token aggregate. @@ -65,18 +71,30 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO { } func taskDTO(t *Task, status string) TaskDTO { + llmState := t.llmStateSnapshot() + sourceIDs := make([]string, 0, len(t.SourceTaskIDs)) + for _, id := range t.SourceTaskIDs { + sourceIDs = append(sourceIDs, i64s(id)) + } + profileIDs := append(make([]int64, 0, len(llmState.ProfileIDs)), llmState.ProfileIDs...) return TaskDTO{ - ID: t.ID, - ExplorationID: t.ExpID, - Description: t.Description, - Goal: t.Goal, - Status: status, - CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)), - CreatedUnix: t.CreatedAt, - CompletedAt: completedRFC(t.CompletedAt), - CompletedUnix: t.CompletedAt, - Paused: t.Paused, - LLMProfileID: t.LLMProfileID, + ID: t.ID, + ExplorationID: t.ExpID, + Description: t.Description, + Goal: t.Goal, + Status: status, + CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)), + CreatedUnix: t.CreatedAt, + CompletedAt: completedRFC(t.CompletedAt), + CompletedUnix: t.CompletedAt, + Paused: t.Paused, + Queued: t.Queued, + LLMProfileID: llmState.ProfileID, + LLMProfileIDs: profileIDs, + ActiveLLMProfileID: llmState.ActiveID, + LLMFailoverState: llmState.FailoverState, + LLMFailoverReason: llmState.FailoverReason, + SourceTaskIDs: sourceIDs, } } @@ -120,17 +138,19 @@ func trafficDTOs(ex []traffic.ExchangeMeta) []TrafficExchangeDTO { // ---- TaskNode (frontend "TaskNode") — frontier, intents, exploration graph nodes ---- type TaskNodeDTO struct { - ID string `json:"id"` - Type string `json:"type"` // db Node.Kind - Payload string `json:"payload,omitempty"` - Priority int `json:"priority"` - State string `json:"state"` - Origin string `json:"origin"` - TS string `json:"ts"` + ID string `json:"id"` + Type string `json:"type"` // db Node.Kind + Payload string `json:"payload,omitempty"` + Priority int `json:"priority"` + State string `json:"state"` + Origin string `json:"origin"` + TS string `json:"ts"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } func taskNodeDTO(n *db.Node) TaskNodeDTO { - return TaskNodeDTO{ + d := TaskNodeDTO{ ID: i64s(n.ID), Type: n.Kind, Payload: rawString(n.Payload), @@ -139,6 +159,11 @@ func taskNodeDTO(n *db.Node) TaskNodeDTO { Origin: n.Origin, TS: rfc3339(n.CreatedAt), } + if n.SourceTaskID > 0 { + d.SourceTaskID = i64s(n.SourceTaskID) + } + d.Inherited = n.Inherited + return d } func taskNodeDTOs(in []*db.Node) []TaskNodeDTO { @@ -169,6 +194,27 @@ func edgeDTOs(in []db.Edge) []EdgeDTO { return out } +// ---- Coverage asset references ---- + +type CoverageAssetRefDTO struct { + ID int64 `json:"id"` + Kind string `json:"kind"` + State string `json:"state"` + Summary string `json:"summary"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` +} + +func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO { + out := CoverageAssetRefDTO{ + ID: ref.ID, Kind: ref.Kind, State: ref.State, Summary: ref.Summary, Inherited: ref.Inherited, + } + if ref.SourceTaskID > 0 { + out.SourceTaskID = i64s(ref.SourceTaskID) + } + return out +} + // ---- Finding (frontend "Finding") ---- type FindingDTO struct { @@ -186,6 +232,8 @@ type FindingDTO struct { ParamID string `json:"param_id,omitempty"` TaskID string `json:"task_id,omitempty"` TaskDescription string `json:"task_description,omitempty"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` Assets []FindingAssetDTO `json:"assets,omitempty"` TS string `json:"ts"` } @@ -241,7 +289,7 @@ type findingPayload struct { func findingDTO(n *db.Node) FindingDTO { var p findingPayload _ = json.Unmarshal(n.Payload, &p) - return FindingDTO{ + d := FindingDTO{ ID: i64s(n.ID), VulnClass: p.VulnClass, Name: p.Name, @@ -251,6 +299,11 @@ func findingDTO(n *db.Node) FindingDTO { Evidence: rawString(p.Evidence), TS: rfc3339(n.CreatedAt), } + if n.SourceTaskID > 0 { + d.SourceTaskID = i64s(n.SourceTaskID) + } + d.Inherited = n.Inherited + return d } // findingDTOsForTask converts a task's finding nodes to DTOs, stamping each with @@ -260,11 +313,15 @@ func findingDTO(n *db.Node) FindingDTO { // nodes with no row keep the 'pending' default and no finding_id (not editable). // assets pre-resolves the anchored asset rows for label rendering. func findingDTOsForTask(t *Task, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO { + return findingDTOsForOwner(t.ID, t.Description, in, meta, assets) +} + +func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO { out := make([]FindingDTO, 0, len(in)) for _, n := range in { d := findingDTO(n) - d.TaskID = t.ID - d.TaskDescription = t.Description + d.TaskID = taskID + d.TaskDescription = description if m, ok := meta[n.ID]; ok { d.FindingID = i64s(m.ID) d.Status = m.Status @@ -320,16 +377,19 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO { // ---- Activity (frontend "Activity") ---- type ActivityDTO struct { - Seq int64 `json:"seq"` // db Activity.ID - IntentID string `json:"intent_id,omitempty"` // db NodeID - Worker string `json:"worker"` - TS string `json:"ts"` // db CreatedAt - Kind string `json:"kind"` - Tool string `json:"tool,omitempty"` - ToolUseID string `json:"tool_use_id,omitempty"` - IsError bool `json:"is_error"` - Summary string `json:"summary"` - Detail string `json:"detail,omitempty"` + Seq int64 `json:"seq"` // db Activity.ID + IntentID string `json:"intent_id,omitempty"` // db NodeID + Worker string `json:"worker"` + TS string `json:"ts"` // db CreatedAt + Kind string `json:"kind"` + Tool string `json:"tool,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + IsError bool `json:"is_error"` + Summary string `json:"summary"` + Detail string `json:"detail,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` // token usage (set only on kind='result'); used for per-session token totals. InputTokens *int `json:"input_tokens,omitempty"` OutputTokens *int `json:"output_tokens,omitempty"` @@ -342,7 +402,7 @@ func activityDTO(a db.Activity) ActivityDTO { if a.NodeID != nil { intent = i64s(*a.NodeID) } - return ActivityDTO{ + d := ActivityDTO{ Seq: a.ID, IntentID: intent, Worker: a.Worker, @@ -353,11 +413,17 @@ func activityDTO(a db.Activity) ActivityDTO { IsError: a.IsError, Summary: a.Summary, Detail: a.Detail, // list endpoint leaves this empty (lazy) + Metadata: a.Metadata, InputTokens: a.InputTokens, OutputTokens: a.OutputTokens, CacheReadTokens: a.CacheReadTokens, CacheWriteTokens: a.CacheWriteTokens, } + if a.SourceTaskID > 0 { + d.SourceTaskID = i64s(a.SourceTaskID) + } + d.Inherited = a.Inherited + return d } func activityDTOs(in []db.Activity) []ActivityDTO { diff --git a/server/engine.go b/server/engine.go index 205f270..c52e76b 100644 --- a/server/engine.go +++ b/server/engine.go @@ -73,6 +73,13 @@ const ( modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避 ) +// retryableWorkerModelError excludes errors already handled by the task router. +// In particular, a quota error after partial streaming advances the task cursor +// for the next LLM call but must not replay this whole intent on the backup. +func retryableWorkerModelError(reason harness.TerminalReason, err error) bool { + return reason == harness.ReasonModelError && !isTaskLLMRuntimeError(err) +} + // Engine drives the event-driven exploration loop with real LLM agents // (docs §4.3/§4.4): on asset/exploration-graph change (debounced) it wakes the // planner, which reads the route, queries assets, judges goals and emits intents; @@ -86,28 +93,43 @@ type Engine struct { bc *Broadcaster // live activity pub/sub (SSE) - mu sync.RWMutex - planner *agent.Planner - worker *agent.Worker - started sync.Map // taskID -> bool, so Run is idempotent per task - lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat) - paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive) - dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records + mu sync.RWMutex + planner *agent.Planner + worker *agent.Worker + started sync.Map // taskID -> bool, so Run is idempotent per task + lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat) + llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls + paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive) + deleting sync.Map // taskID -> bool, delete barrier (no new task-owned writes) + // deleteWasPaused remembers whether BeginDelete temporarily paused an already + // running task. AbortDelete uses it to restore the exact pre-delete state. + deleteWasPaused sync.Map // taskID -> bool + dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records + + // deleteMu makes installing the delete barrier atomic with registering a new + // task operation. Once BeginDelete returns, every admitted writer is reflected + // in inflight and every later writer is rejected. + deleteMu sync.RWMutex + + // Every long-lived task goroutine (planner, workers and deadline coordinator) + // runs under one task-scoped context. Successful deletion cancels that context, + // waits for all goroutines, then releases every task-level Engine reference. + runtimeMu sync.Mutex + runtimes map[string]*taskRuntime // per-task execution context: each planner.Plan / worker.Execute runs under it, // so pausing can CANCEL an in-flight run (not just skip the next one). Recreated - // on resume since cancelling is one-shot. - // CancelCauseFunc (not CancelFunc): every cancel site attaches a named - // *agent.AbortCause so the run's trace can say WHICH of the ~10 cancel - // paths fired instead of one catch-all "运行被中断" line. + // on resume since cancelling is one-shot. Every cancellation carries a named + // cause so the activity trace can identify the initiating control path. execMu sync.Mutex execCancel map[string]context.CancelCauseFunc execCtx map[string]context.Context - // per-work (per running intent) cancel, so the planner's kill_work can stop a - // single worker without pausing the whole task. Keyed by globally-unique intent id. - workMu sync.Mutex - workCancel map[int64]context.CancelCauseFunc + // Per-work control lets the planner kill a worker and lets the UI pause/cancel + // one intent without pausing the whole task. The done channel closes only after + // runWorkerStep has stopped writing and committed its final state. + workMu sync.Mutex + work map[int64]*workExecution // steerBox queues planner course-corrections for a running work (keyed by intent // id). The worker's PreToolUse hook drains it before its next tool call and hands @@ -124,9 +146,23 @@ type Engine struct { inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain) coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重) - // resolve, if set, returns a task's dedicated planner/worker when it pins a - // specific LLM profile (nil,nil = use the global active pair). Set once at startup. - resolve func(t *Task) (*agent.Planner, *agent.Worker) + // resolve, if set, returns a task's dedicated planner/worker. In legacy mode + // nil,nil falls back to the global pair; an authoritative resolver uses nil,nil + // to mean this task is deliberately unavailable (for example an exhausted chain). + resolve func(t *Task) (*agent.Planner, *agent.Worker) + resolveAuthoritative bool +} + +type taskRuntime struct { + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup +} + +type workExecution struct { + cancel context.CancelCauseFunc + done chan struct{} + action string // user action: pause | cancel } // nextPlannerRound returns the next planner round number for a task (1-based). @@ -139,17 +175,132 @@ func (e *Engine) nextPlannerRound(taskID string) int { // Pause stops a task: marks it paused AND cancels any in-flight planner/worker run // for it (a long worker.Execute would otherwise keep going until it finishes). -// cause says WHY (user button / orchestrator tool / delete / reload) — it surfaces -// verbatim in the cancelled run's trace, so pass the matching agent.Abort* value. func (e *Engine) Pause(taskID string, cause error) { e.paused.Store(taskID, true) e.cancelExec(taskID, cause) } +// BeginDelete installs an execution barrier before task data/files are removed. +// The temporary pause is not a user pause: AbortDelete restores the exact state +// observed here when cleanup fails. +func (e *Engine) BeginDelete(taskID string) bool { + e.deleteMu.Lock() + if _, loaded := e.deleting.LoadOrStore(taskID, true); loaded { + e.deleteMu.Unlock() + return false + } + wasPaused := e.IsPaused(taskID) + e.deleteWasPaused.Store(taskID, wasPaused) + e.paused.Store(taskID, true) + e.deleteMu.Unlock() + e.cancelExec(taskID, agent.AbortTaskDeleted) + return true +} + +func (e *Engine) AbortDelete(taskID string) { + e.deleteMu.Lock() + if !e.IsDeleting(taskID) { + e.deleteMu.Unlock() + return + } + wasPaused := false + if v, ok := e.deleteWasPaused.Load(taskID); ok { + wasPaused, _ = v.(bool) + } + e.deleteWasPaused.Delete(taskID) + e.deleting.Delete(taskID) + if !wasPaused { + e.paused.Delete(taskID) + } + e.deleteMu.Unlock() + if !wasPaused && e.m != nil { + if t, ok := e.m.Task(taskID); ok { + t.Notify() + } + } +} + +func (e *Engine) IsDeleting(taskID string) bool { + _, ok := e.deleting.Load(taskID) + return ok +} + +// registerTaskRoutines reserves count goroutines in the task runtime. Callers +// hold deleteMu for reading so StopTask cannot race WaitGroup.Add with Wait. +func (e *Engine) registerTaskRoutines(parent context.Context, taskID string, count int) *taskRuntime { + e.runtimeMu.Lock() + defer e.runtimeMu.Unlock() + rt := e.runtimes[taskID] + if rt == nil { + ctx, cancel := context.WithCancel(parent) + rt = &taskRuntime{ctx: ctx, cancel: cancel} + e.runtimes[taskID] = rt + } + rt.wg.Add(count) + return rt +} + +func runTaskRoutine(rt *taskRuntime, fn func(context.Context)) { + go func() { + defer rt.wg.Done() + fn(rt.ctx) + }() +} + +// StopTask permanently stops every long-lived goroutine and removes all Engine +// state for a successfully deleted task. The delete barrier remains installed +// until cleanup finishes, so no new task operation can race the teardown. +func (e *Engine) StopTask(taskID string) { + e.deleteMu.Lock() + e.deleting.Store(taskID, true) + e.deleteMu.Unlock() + + e.cancelExec(taskID, agent.AbortTaskDeleted) + e.runtimeMu.Lock() + rt := e.runtimes[taskID] + if rt != nil { + rt.cancel() + } + e.runtimeMu.Unlock() + if rt != nil { + rt.wg.Wait() + } + + e.execMu.Lock() + if cancel := e.execCancel[taskID]; cancel != nil { + cancel(agent.AbortTaskDeleted) + } + delete(e.execCancel, taskID) + delete(e.execCtx, taskID) + e.execMu.Unlock() + + e.runtimeMu.Lock() + if e.runtimes[taskID] == rt { + delete(e.runtimes, taskID) + } + e.runtimeMu.Unlock() + + e.started.Delete(taskID) + e.lastAct.Delete(taskID) + e.llmCalls.Delete(taskID) + e.paused.Delete(taskID) + e.dropCnt.Delete(taskID) + e.plannerRound.Delete(taskID) + e.settling.Delete(taskID) + e.deadline.Delete(taskID) + e.stamped.Delete(taskID) + e.inflight.Delete(taskID) + e.coordStarted.Delete(taskID) + e.deleteWasPaused.Delete(taskID) + + e.deleteMu.Lock() + e.deleting.Delete(taskID) + e.deleteMu.Unlock() +} + // cancelExec cancels a task's current per-task exec context (any in-flight // planner.Plan / worker.Execute), if present. Shared by Pause and the settle -// sequence's hard-drain backstop. cause is recovered by the aborted run via -// context.Cause and printed in its result record. +// sequence's hard-drain backstop. func (e *Engine) cancelExec(taskID string, cause error) { e.execMu.Lock() if cancel := e.execCancel[taskID]; cancel != nil { @@ -206,32 +357,94 @@ func (e *Engine) LastActivity(taskID string) int64 { return 0 } +// BeginLLMCall/EndLLMCall track actual provider calls separately from the +// scheduler's task-operation counter. A task can have live loops while all of +// them are waiting for a trigger; that state must remain idle in the UI. +func (e *Engine) BeginLLMCall(taskID string) { + v, _ := e.llmCalls.LoadOrStore(taskID, new(int64)) + atomic.AddInt64(v.(*int64), 1) +} + +func (e *Engine) EndLLMCall(taskID string) { + if v, ok := e.llmCalls.Load(taskID); ok { + p := v.(*int64) + if atomic.AddInt64(p, -1) <= 0 { + atomic.StoreInt64(p, 0) + } + } +} + +func (e *Engine) ActiveLLMCalls(taskID string) int64 { + if v, ok := e.llmCalls.Load(taskID); ok { + return atomic.LoadInt64(v.(*int64)) + } + return 0 +} + func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) } func NewEngine(m *Manager) *Engine { return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(), execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{}, - workCancel: map[int64]context.CancelCauseFunc{}, steerBox: map[int64][]string{}} + work: map[int64]*workExecution{}, steerBox: map[int64][]string{}, + runtimes: map[string]*taskRuntime{}} } // registerWork records the cancel for the work currently running intentID. func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) { e.workMu.Lock() - e.workCancel[intentID] = cancel + e.work[intentID] = &workExecution{cancel: cancel, done: make(chan struct{})} e.workMu.Unlock() } -// unregisterWork removes and releases the work's cancel (called when Execute returns). -func (e *Engine) unregisterWork(intentID int64) { +// detachWork removes the live control handle once Execute has returned. complete +// must be called after the final intent state write so a waiting cancel handler can +// safely delete the worker's blackboard output without racing a late write. +func (e *Engine) detachWork(intentID int64) (action string, complete func()) { e.workMu.Lock() - if c := e.workCancel[intentID]; c != nil { - delete(e.workCancel, intentID) - c(agent.AbortWorkFinished) // release context resources (no-op if already cancelled) + run := e.work[intentID] + if run != nil { + delete(e.work, intentID) + action = run.action + run.cancel(agent.AbortWorkFinished) // release resources (no-op if already cancelled) } e.workMu.Unlock() e.steerMu.Lock() delete(e.steerBox, intentID) // drop any undelivered steering for a finished work e.steerMu.Unlock() + if run == nil { + return action, func() {} + } + return action, func() { close(run.done) } +} + +// ControlWork requests a user-visible pause or cancellation and waits until the +// worker has fully stopped writing. Cancellation cleanup is performed by the API +// handler after this returns; pause state is committed by runWorkerStep itself. +func (e *Engine) ControlWork(intentID int64, action string) error { + if action != "pause" && action != "cancel" { + return fmt.Errorf("unsupported work action %q", action) + } + e.workMu.Lock() + run := e.work[intentID] + if run == nil { + e.workMu.Unlock() + return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID) + } + if run.action != "" && run.action != action { + e.workMu.Unlock() + return fmt.Errorf("意图 %d 正在执行 %s 操作", intentID, run.action) + } + run.action = action + done := run.done + cause := error(agent.AbortWorkPausedByUser) + if action == "cancel" { + cause = agent.AbortWorkCancelledByUser + } + run.cancel(cause) + e.workMu.Unlock() + <-done + return nil } // SteerWork queues a mid-run course-correction for the work running intentID (the @@ -242,7 +455,7 @@ func (e *Engine) SteerWork(intentID int64, msg string) error { return fmt.Errorf("纠偏消息不能为空") } e.workMu.Lock() - running := e.workCancel[intentID] != nil + running := e.work[intentID] != nil e.workMu.Unlock() if !running { return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID) @@ -307,12 +520,12 @@ func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []s // The work's agent-core session honors ctx cancellation and aborts promptly. func (e *Engine) KillWork(intentID int64) error { e.workMu.Lock() - c := e.workCancel[intentID] + run := e.work[intentID] e.workMu.Unlock() - if c == nil { + if run == nil { return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID) } - c(agent.AbortKilledByPlanner) + run.cancel(agent.AbortKilledByPlanner) return nil } @@ -321,7 +534,7 @@ func (e *Engine) Broadcaster() *Broadcaster { return e.bc } // emitActivity persists one captured step AND fans it out to live subscribers, // from a single point so storage and the SSE stream never diverge. -func (e *Engine) emitActivity(t *Task, r db.Activity) { +func (e *Engine) emitActivity(t *Task, r db.Activity) db.Activity { id, err := e.appendActivity(t, r) if err != nil { // NO LONGER SILENT: dropping a record breaks command↔result pairing in the @@ -346,7 +559,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) { log.Printf("[activity] task %s 丢弃活动记录(该任务累计第 %d 条) worker=%s kind=%s tool=%s tuid=%s reason=%s summary=%q: %v%s", t.ID, n, r.Worker, r.Kind, r.Tool, r.ToolUseID, dropReason(err), preview(r.Summary, 80), err, diag) e.touch(t.ID) - return + return r } r.ID = id if r.CreatedAt.IsZero() { @@ -354,6 +567,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) { } e.bc.Publish(t.ID, r) e.touch(t.ID) + return r } // appendActivity persists one activity row, retrying briefly on write failure. @@ -396,6 +610,15 @@ func (e *Engine) snapshot() (*agent.Planner, *agent.Worker) { // Called once at startup before any task loop runs, so no lock is needed on reads. func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) { e.resolve = fn + e.resolveAuthoritative = false +} + +// SetAuthoritativeAgentResolver installs a resolver whose nil result must not +// fall through to the global provider. Task-level failover chains use this so a +// fully exhausted chain cannot silently bypass its configured boundary. +func (e *Engine) SetAuthoritativeAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) { + e.resolve = fn + e.resolveAuthoritative = true } // snapshotFor returns the planner/worker a task should run on: the task's dedicated @@ -403,7 +626,8 @@ func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worke // pair. The resolver returning nil means "no override" (fall back to global). func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) { if e.resolve != nil { - if p, w := e.resolve(t); p != nil && w != nil { + p, w := e.resolve(t) + if (p != nil && w != nil) || e.resolveAuthoritative { return p, w } } @@ -416,6 +640,14 @@ func (e *Engine) Ready() bool { return p != nil && w != nil } +// ReadyFor reports whether a specific task can resolve a planner/worker pair. +// An explicit task profile chain can be runnable even when no global default +// provider is configured, so task status must not rely on Ready alone. +func (e *Engine) ReadyFor(t *Task) bool { + p, w := e.snapshotFor(t) + return p != nil && w != nil +} + // Mode reports "llm" when ready, else "idle". func (e *Engine) Mode() string { if e.Ready() { @@ -428,15 +660,24 @@ func (e *Engine) Mode() string { // 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() + e.deleteMu.RLock() + if e.IsDeleting(t.ID) { + e.deleteMu.RUnlock() + return + } if _, loaded := e.started.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RUnlock() t.Notify() // already running — just nudge a planning round return } + rt := e.registerTaskRoutines(ctx, t.ID, 1+workers) + e.deleteMu.RUnlock() e.touch(t.ID) - go e.plannerLoop(ctx, t) - workers := e.m.Workers() + runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) }) for i := 0; i < workers; i++ { - go e.workerLoop(ctx, t, fmt.Sprintf("work#%d", i+1)) + name := fmt.Sprintf("work#%d", i+1) + runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, name) }) } e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重) // 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已 @@ -497,6 +738,9 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { if e.IsPaused(t.ID) { return // user-paused: don't plan } + if e.IsDeleting(t.ID) { + return + } // terminal task (goals all met → done, or failed): the run is over. A // resume/nudge — e.g. auto-resume of the active task on restart — must NOT // re-plan (it would burn an LLM round and re-confirm a settled result). @@ -508,10 +752,17 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { if e.isSettling(t.ID) { return } + if !e.beginTaskOperation(t.ID) { + return + } + defer e.decInflight(t.ID) e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) e.touch(t.ID) emit := func(r db.Activity) { e.emitActivity(t, r) } ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline + if ectx.Err() != nil || e.IsDeleting(t.ID) { + return + } log.Printf("[planner] task %s 规划中…(%s 触发)", t.ID, src) // round marker: each Plan() is one planner round; emit a boundary so the // UI can separate rounds in the transcript (kind='round'). @@ -520,10 +771,10 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { // what fired this round (worker done / finding; may be several — debounce // coalesces a burst; empty for time/heartbeat wakes). triggers := t.drainTriggers() - e.incInflight(t.ID) taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) met, reason, err := planner.Plan(ectx, taskIDInt, e.m.assets, t.Store, t.Goal, triggers, emit) - e.decInflight(t.ID) + e.EndLLMCall(t.ID) switch { case err != nil && ectx.Err() == nil: log.Printf("[planner] task %s 规划出错: %v", t.ID, err) @@ -578,6 +829,9 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { } continue // user-paused: don't claim/execute intents } + if e.IsDeleting(t.ID) { + return + } if e.isSettling(t.ID) { if sleepCtx(ctx, 1000*time.Millisecond) { return @@ -590,108 +844,144 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { } continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier } - intent := e.claimNext(t, name) - if intent == nil { - if sleepCtx(ctx, 800*time.Millisecond) { - return - } - continue - } - log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID) - e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) - e.touch(t.ID) - emit := func(r db.Activity) { e.emitActivity(t, r) } - ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline - // per-work child context so the planner's kill_work can stop just this work. - workCtx, workCancel := context.WithCancelCause(ectx) - e.registerWork(intent.ID, workCancel) - // wrap the guard hooks so steer_work can inject a mid-run course-correction - // for THIS intent (drained before the worker's next tool call). - iid := intent.ID - taskEmit := func(a db.Activity) { - nid := iid - a.NodeID, a.Worker = &nid, name - emit(a) - } - workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit) - hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }} - wTaskID, _ := strconv.ParseInt(t.ID, 10, 64) - e.incInflight(t.ID) // 计入在跑,供收尾时序 drain 等待 - reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) - // model_error 收场 → 额外重跑几次(退避后再试)。仅在意图仍属本 work、任务 - // 未暂停/未终止/未取消【且未进入收尾】时重试;否则让位给对应分支处理(收尾期不 - // 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。 - for attempt := 1; attempt <= modelErrorRetries && - reason == harness.ReasonModelError && - workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ { - log.Printf("[worker %s] task %s 意图 #%d model_error 收场,%v 后重试 (%d/%d)", - name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries) - if sleepCtx(workCtx, modelErrorRetryBackoff) { - break // 退避期间被取消(终止/暂停)→ 交给下方分支处理 - } - reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + if !e.beginTaskOperation(t.ID) { + return } + claimed := e.runWorkerStep(ctx, t, name, worker) e.decInflight(t.ID) - // CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's - // ctx was cancelled (planner kill_work) while the TASK ctx kept running; a - // pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER - // unregister would always be true (unregister cancels it) → every completed - // work would be wrongly marked stopped. - killed := workCtx.Err() != nil && ectx.Err() == nil - e.unregisterWork(intent.ID) - // if a pause cancelled this run mid-flight, return the intent to the frontier - // so it is re-claimed (and re-run from scratch) on resume — not marked done. - if ectx.Err() != nil && e.IsPaused(t.ID) { - _ = t.Store.SetIntentState(intent.ID, "open") - continue + if !claimed && sleepCtx(ctx, 800*time.Millisecond) { + return } - // 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾), - // 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。 - if ectx.Err() != nil && e.isSettling(t.ID) { - _ = t.Store.SetIntentState(intent.ID, "exhausted") - log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote) - e.touch(t.ID) - continue - } - // 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义, - // 标 stopped(不是 blocked),别污染已完成任务的意图状态。 - if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) { - _ = t.Store.SetIntentState(intent.ID, "stopped") - log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID) - e.touch(t.ID) - continue - } - // killed by the planner: mark stopped (don't write back results, don't auto-reclaim). - if killed { - _ = t.Store.SetIntentState(intent.ID, "stopped") - log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID) - e.touch(t.ID) - t.Notify() - continue - } - if err != nil { - log.Printf("[worker %s] intent %d: %v", name, intent.ID, err) - } - // terminal分流:撞步数上限 ≠ 完成。max_turns→exhausted(规划者据此知道这个方向 - // 试过但没真正做完、需换角度,而非当成已覆盖永久跳过);出错→blocked;正常→done。 - state := "done" - switch { - case err != nil: - state = "blocked" - case reason == harness.ReasonMaxTurns: - state = "exhausted" - log.Printf("[worker %s] intent %d 撞步数上限(exhausted),本次写回 %s", name, intent.ID, wrote) - case reason == harness.ReasonTimeout: - state = "exhausted" - log.Printf("[worker %s] intent %d 运行超时(exhausted),收尾后写回 %s", name, intent.ID, wrote) - } - _ = t.Store.SetIntentState(intent.ID, state) - log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote) - e.touch(t.ID) - t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id) } } +// runWorkerStep claims and fully settles at most one intent. Its caller holds one +// task-operation admission for the whole claim/execute/state-write sequence so a +// delete cannot observe quiescence between the LLM return and the final DB writes. +func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker *agent.Worker) bool { + intent := e.claimNext(t, name) + if intent == nil { + return false + } + log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID) + e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) + e.touch(t.ID) + emit := func(r db.Activity) { e.emitActivity(t, r) } + ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline + if ectx.Err() != nil || e.IsDeleting(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "open") + return true + } + // per-work child context so the planner's kill_work can stop just this work. + workCtx, workCancel := context.WithCancelCause(ectx) + e.registerWork(intent.ID, workCancel) + // wrap the guard hooks so steer_work can inject a mid-run course-correction + // for THIS intent (drained before the worker's next tool call). + iid := intent.ID + taskEmit := func(a db.Activity) { + nid := iid + a.NodeID, a.Worker = &nid, name + emit(a) + } + workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit) + hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }} + wTaskID, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) + reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + e.EndLLMCall(t.ID) + // model_error 收场 → 额外重跑几次(退避后再试)。仅在意图仍属本 work、任务 + // 未暂停/未终止/未取消【且未进入收尾】时重试;否则让位给对应分支处理(收尾期不 + // 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。 + for attempt := 1; attempt <= modelErrorRetries && + retryableWorkerModelError(reason, err) && + workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ { + log.Printf("[worker %s] task %s 意图 #%d model_error 收场,%v 后重试 (%d/%d)", + name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries) + if sleepCtx(workCtx, modelErrorRetryBackoff) { + break // 退避期间被取消(终止/暂停)→ 交给下方分支处理 + } + e.BeginLLMCall(t.ID) + reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + e.EndLLMCall(t.ID) + } + // Capture kill state before detachWork cancels workCtx. kill = this work's + // ctx was cancelled (planner kill_work) while the TASK ctx kept running; a + // pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER + // unregister would always be true (unregister cancels it) → every completed + // work would be wrongly marked stopped. + killed := workCtx.Err() != nil && ectx.Err() == nil + action, completeWork := e.detachWork(intent.ID) + defer completeWork() + if action == "pause" { + _ = t.Store.SetIntentState(intent.ID, "paused") + log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + if action == "cancel" { + // Leave the node in running until the waiting API handler deletes the intent + // and its direct outputs transactionally. No task Notify is emitted here. + log.Printf("[worker %s] task %s 意图 #%d 已停止,等待取消清理", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + // if a pause cancelled this run mid-flight, return the intent to the frontier + // so it is re-claimed (and re-run from scratch) on resume — not marked done. + if ectx.Err() != nil && e.IsPaused(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "open") + return true + } + // 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾), + // 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。 + if ectx.Err() != nil && e.isSettling(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "exhausted") + log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote) + e.touch(t.ID) + return true + } + // 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义, + // 标 stopped(不是 blocked),别污染已完成任务的意图状态。 + if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) { + _ = t.Store.SetIntentState(intent.ID, "stopped") + log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + // killed by the planner: mark stopped (don't write back results, don't auto-reclaim). + if killed { + _ = t.Store.SetIntentState(intent.ID, "stopped") + log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID) + e.touch(t.ID) + t.Notify() + return true + } + if err != nil { + log.Printf("[worker %s] intent %d: %v", name, intent.ID, err) + } + // terminal分流:撞步数上限 ≠ 完成。max_turns→exhausted(规划者据此知道这个方向 + // 试过但没真正做完、需换角度,而非当成已覆盖永久跳过);出错→blocked;正常→done。 + state := "done" + switch { + case err != nil: + state = "blocked" + case reason == harness.ReasonMaxTurns: + state = "exhausted" + log.Printf("[worker %s] intent %d 撞步数上限(exhausted),本次写回 %s", name, intent.ID, wrote) + case reason == harness.ReasonTimeout: + state = "exhausted" + log.Printf("[worker %s] intent %d 运行超时(exhausted),收尾后写回 %s", name, intent.ID, wrote) + } + if state == "blocked" && isTaskLLMChainExhausted(err) { + _ = t.Store.SetIntentBlockedReason(intent.ID, db.IntentBlockedLLMQuota) + } else { + _ = t.Store.SetIntentState(intent.ID, state) + } + log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote) + e.touch(t.ID) + t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id) + return true +} + func sleepCtx(ctx context.Context, d time.Duration) (done bool) { select { case <-ctx.Done(): diff --git a/server/engine_cancelcause_test.go b/server/engine_cancelcause_test.go index 0669679..5501821 100644 --- a/server/engine_cancelcause_test.go +++ b/server/engine_cancelcause_test.go @@ -7,90 +7,103 @@ import ( "github.com/Autumn-27/artex/agent" ) -// The engine's cancel paths must hand a NAMED cause to the run being cancelled — -// that cause is the whole content of the "运行被中断" line in the trace. These -// exercise the context machinery only (no Manager / DB needed). - func TestPauseCancelsWithNamedCause(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - // the per-work child the worker loop actually runs under - workCtx, workCancel := context.WithCancelCause(ectx) + execCtx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(execCtx) defer workCancel(nil) - e.Pause("t1", agent.AbortPausedByUser) - if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "paused_by_user" { - t.Fatalf("work ctx 拿到的原因 code=%q ok=%v,期望 paused_by_user", code, ok) + t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok) } } func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - workCtx, workCancel := context.WithCancelCause(ectx) + execCtx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(execCtx) e.registerWork(42, workCancel) - if err := e.KillWork(42); err != nil { - t.Fatalf("KillWork: %v", err) + t.Fatal(err) } if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "killed_by_planner" { - t.Fatalf("code=%q ok=%v,期望 killed_by_planner", code, ok) + t.Fatalf("code=%q ok=%v, want killed_by_planner", code, ok) } - // the task ctx must survive — kill_work stops one work, not the whole task - if ectx.Err() != nil { - t.Error("kill_work 不该取消整个任务的 exec ctx") - } - if err := e.KillWork(999); err == nil { - t.Error("对没有在跑的意图 KillWork 应该报错") + if execCtx.Err() != nil { + t.Fatal("kill_work cancelled the entire task context") } + _, complete := e.detachWork(42) + complete() } -func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) { +func TestControlWorkCarriesUserActionCause(t *testing.T) { for _, tc := range []struct { - name string - cause *agent.AbortCause - want string + action string + code string }{ - {"settle", agent.AbortSettleDrainTimeout, "settle_drain_timeout"}, - {"goalMet", agent.AbortGoalMet, "goal_met"}, + {action: "pause", code: "work_paused_by_user"}, + {action: "cancel", code: "work_cancelled_by_user"}, } { - t.Run(tc.name, func(t *testing.T) { + t.Run(tc.action, func(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - e.cancelExec("t1", tc.cause) - if code, _, _, _ := agent.AbortReason(ectx); code != tc.want { - t.Fatalf("code=%q,期望 %q", code, tc.want) + ctx, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + done := make(chan error, 1) + go func() { done <- e.ControlWork(42, tc.action) }() + <-ctx.Done() + if code, _, _, _ := agent.AbortReason(ctx); code != tc.code { + t.Fatalf("code=%q, want %q", code, tc.code) + } + action, complete := e.detachWork(42) + if action != tc.action { + t.Fatalf("action=%q, want %q", action, tc.action) + } + complete() + if err := <-done; err != nil { + t.Fatal(err) } }) } } -// asking for an exec ctx while paused must return an already-cancelled ctx that -// still SAYS why — this is the claim→Execute race guard, easily mistaken for a -// mystery abort. +func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) { + for _, tc := range []struct { + cause *agent.AbortCause + code string + }{ + {agent.AbortSettleDrainTimeout, "settle_drain_timeout"}, + {agent.AbortGoalMet, "goal_met"}, + } { + e := NewEngine(nil) + ctx := e.execContextFor(context.Background(), "t1") + e.cancelExec("t1", tc.cause) + if code, _, _, _ := agent.AbortReason(ctx); code != tc.code { + t.Fatalf("code=%q, want %q", code, tc.code) + } + } +} + func TestExecContextWhilePausedCarriesRaceGuardCause(t *testing.T) { e := NewEngine(nil) e.Pause("t1", agent.AbortPausedByUser) ctx := e.execContextFor(context.Background(), "t1") if ctx.Err() == nil { - t.Fatal("暂停期间不该发出可用的 exec ctx") + t.Fatal("paused task received a live execution context") } if code, _, _, _ := agent.AbortReason(ctx); code != "paused_race_guard" { - t.Fatalf("code=%q,期望 paused_race_guard", code) + t.Fatalf("code=%q, want paused_race_guard", code) } } -// Resume must hand out a live context again (cancelling is one-shot). func TestResumeYieldsFreshContext(t *testing.T) { e := NewEngine(nil) old := e.execContextFor(context.Background(), "t1") e.Pause("t1", agent.AbortPausedByUser) if old.Err() == nil { - t.Fatal("Pause 没取消旧 ctx") + t.Fatal("Pause did not cancel old context") } - e.paused.Delete("t1") // Resume() without the Task/Notify plumbing + e.paused.Delete("t1") if fresh := e.execContextFor(context.Background(), "t1"); fresh.Err() != nil { - t.Fatal("恢复后应拿到全新可用的 exec ctx") + t.Fatal("resume did not create a fresh execution context") } } diff --git a/server/engine_llm_calls_test.go b/server/engine_llm_calls_test.go new file mode 100644 index 0000000..56013e7 --- /dev/null +++ b/server/engine_llm_calls_test.go @@ -0,0 +1,43 @@ +package server + +import ( + "sync" + "testing" +) + +func TestEngineActiveLLMCallsConcurrent(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + const calls = 64 + + var started sync.WaitGroup + var release sync.WaitGroup + var finished sync.WaitGroup + started.Add(calls) + finished.Add(calls) + release.Add(1) + for range calls { + go func() { + defer finished.Done() + e.BeginLLMCall(taskID) + started.Done() + release.Wait() + e.EndLLMCall(taskID) + }() + } + started.Wait() + if got := e.ActiveLLMCalls(taskID); got != calls { + t.Fatalf("active calls = %d, want %d", got, calls) + } + release.Done() + finished.Wait() + if got := e.ActiveLLMCalls(taskID); got != 0 { + t.Fatalf("active calls after completion = %d, want 0", got) + } + + // A defensive extra end must never expose a negative UI count. + e.EndLLMCall(taskID) + if got := e.ActiveLLMCalls(taskID); got != 0 { + t.Fatalf("active calls after extra end = %d, want 0", got) + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 751c61d..7d98e27 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -43,7 +43,20 @@ func (e *Engine) inflightCounter(taskID string) *int64 { v, _ := e.inflight.LoadOrStore(taskID, new(int64)) return v.(*int64) } -func (e *Engine) incInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), 1) } + +// beginTaskOperation atomically registers a task-owned operation unless deletion +// has already installed its barrier. The delete handler can therefore wait for +// inflight==0 without a check-then-start race recreating files after cleanup. +func (e *Engine) beginTaskOperation(taskID string) bool { + e.deleteMu.RLock() + defer e.deleteMu.RUnlock() + if e.IsDeleting(taskID) { + return false + } + atomic.AddInt64(e.inflightCounter(taskID), 1) + return true +} + func (e *Engine) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) } func (e *Engine) inflightCount(taskID string) int64 { return atomic.LoadInt64(e.inflightCounter(taskID)) @@ -106,10 +119,18 @@ func (e *Engine) startDeadlineCoordinator(ctx context.Context, t *Task) { if t == nil || t.TimeoutSeconds <= 0 { return } - if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RLock() + if e.IsDeleting(t.ID) { + e.deleteMu.RUnlock() return } - go e.deadlineCoordinator(ctx, t) + if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RUnlock() + return + } + rt := e.registerTaskRoutines(ctx, t.ID, 1) + e.deleteMu.RUnlock() + runTaskRoutine(rt, func(loopCtx context.Context) { e.deadlineCoordinator(loopCtx, t) }) } // deadlineCoordinator waits until the task's absolute deadline, then runs the settle @@ -170,6 +191,10 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。 met := e.runFinalPlannerRound(ctx, t) + if !e.beginTaskOperation(t.ID) { + return + } + defer e.decInflight(t.ID) // ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done, // 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。 @@ -192,6 +217,9 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // task-timeout planner words (final goal judgment; no new intents). Waits for the // LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout. func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { + if e.IsDeleting(t.ID) { + return false + } planner, _ := e.snapshotFor(t) for planner == nil { if sleepCtx(ctx, deadlinePollInterval) { @@ -200,17 +228,24 @@ func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { if isTerminalStatus(e.m.TaskStatus(t.ID)) { return false } + if e.IsDeleting(t.ID) { + return false + } planner, _ = e.snapshotFor(t) } + // 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。 + fctx := e.clockCtx(ctx, t, true) + if !e.beginTaskOperation(t.ID) { + return false + } + defer e.decInflight(t.ID) emit := func(r db.Activity) { e.emitActivity(t, r) } e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))}) - // 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。 - fctx := e.clockCtx(ctx, t, true) - e.incInflight(t.ID) tTaskID, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) met, reason, err := planner.Plan(fctx, tTaskID, e.m.assets, t.Store, t.Goal, t.drainTriggers(), emit) - e.decInflight(t.ID) + e.EndLLMCall(t.ID) if err != nil { log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err) } else if met { diff --git a/server/goals.go b/server/goals.go index 41c4e99..3d83708 100644 --- a/server/goals.go +++ b/server/goals.go @@ -10,7 +10,6 @@ import ( "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" - "github.com/Autumn-27/norma/llm" ) type goalSpec struct { @@ -27,25 +26,36 @@ type goalSpec struct { // // 异步(goroutine)所以调用方立即返回,两条路径行为一致:秒建任务、后台拆目标。 func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) { + if !s.engine.beginTaskOperation(t.ID) { + return + } s.seed(t, seedText) if seedFirstIntent { s.seedFirstIntent(t) } - // 并发上限:满了就挂起排队(不做目标拆解、不启动引擎),由 reconcileConcurrency 补位启动。 if s.queueIfAtCapacity(t) { + s.engine.decInflight(t.ID) return } - go s.startTaskEngine(t) + go func() { + defer s.engine.decInflight(t.ID) + s.startTaskEngine(t) + }() } -// startTaskEngine does the actual running work: 第0轮目标拆解 + 启动引擎循环。 -// Shared by fresh launches and by the concurrency reconciler promoting a queued task. func (s *Server) startTaskEngine(t *Task) { + ctx := s.engine.execContextFor(s.ctx, t.ID) + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", Summary: "第 0 轮目标拆解"}) - goals := s.createGoals(s.ctx, t, func(r db.Activity) { + goals := s.createGoals(ctx, t, func(r db.Activity) { s.engine.emitActivity(t, r) }) + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } for _, g := range goals { summary := g.Text if g.VulnClass != "" { @@ -53,37 +63,26 @@ func (s *Server) startTaskEngine(t *Task) { } s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary}) } + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } s.engine.Run(s.ctx, t) } -// occupiesSlot reports whether a task counts against the concurrency cap: it is -// live (not terminal), not queued, and not paused. A just-created / decomposing -// task counts too (it will run momentarily). -func (s *Server) occupiesSlot(t *Task) bool { - if t == nil || t.Queued || isTerminalStatus(t.Status) { - return false - } - return !t.Paused && !s.engine.IsPaused(t.ID) +func (s *Server) occupiesConcurrencySlot(t *Task) bool { + return t != nil && !t.Queued && !t.Paused && !s.engine.IsPaused(t.ID) && !isTerminalStatus(t.Status) } -// runningTaskCount counts tasks occupying a concurrency slot, excluding excludeID -// (pass the candidate's id when deciding admission so it doesn't count itself). func (s *Server) runningTaskCount(excludeID string) int { - n := 0 - for _, t := range s.m.List() { - if t.ID == excludeID { - continue - } - if s.occupiesSlot(t) { - n++ + count := 0 + for _, task := range s.m.List() { + if task.ID != excludeID && s.occupiesConcurrencySlot(task) { + count++ } } - return n + return count } -// queueIfAtCapacity holds a task in the queued state when the concurrency cap is -// enabled and already full. Returns true if the task was queued (caller must not -// start it). No-op (false) when the feature is off or a slot is free. func (s *Server) queueIfAtCapacity(t *Task) bool { s.concMu.Lock() defer s.concMu.Unlock() @@ -93,52 +92,49 @@ func (s *Server) queueIfAtCapacity(t *Task) bool { } if err := s.m.SetTaskQueued(t.ID, true); err != nil { log.Printf("[concurrency] task %s 置排队失败: %v", t.ID, err) - return false // 落库失败宁可放它跑,也别静默丢任务 + return false } - t.Queued = true - s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", + s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: fmt.Sprintf("已排队:达到并发上限 %d,等待空位后自动开始", limit)}) - log.Printf("[concurrency] task %s 排队(并发上限 %d)", t.ID, limit) return true } -// reconcileConcurrency promotes queued tasks into free slots. Runs on every -// scheduler tick. When the feature is disabled it releases ALL queued tasks; -// when enabled it starts the oldest-queued tasks up to (limit - running). func (s *Server) reconcileConcurrency() { s.concMu.Lock() defer s.concMu.Unlock() - var queued []*Task - for _, t := range s.m.List() { - if t.Queued && !isTerminalStatus(t.Status) { - queued = append(queued, t) + queued := []*Task{} + for _, task := range s.m.List() { + if task.Queued && !isTerminalStatus(task.Status) { + queued = append(queued, task) } } - if len(queued) == 0 { - return - } - // 最早排队的先启动。 sort.Slice(queued, func(i, j int) bool { return queued[i].CreatedAt < queued[j].CreatedAt }) - enabled, limit := s.m.ConcurrencyLimit() - slots := len(queued) // feature off → 全部放行 + slots := len(queued) if enabled { - if !s.engine.Ready() { - return // 引擎未就绪(无 LLM):继续挂起,别空跑 - } slots = limit - s.runningTaskCount("") } - for _, t := range queued { + for _, task := range queued { if slots <= 0 { break } - if err := s.m.SetTaskQueued(t.ID, false); err != nil { - log.Printf("[concurrency] task %s 解除排队失败: %v", t.ID, err) + // Keep FIFO order, but do not consume a concurrency slot for a task + // whose explicit chain/global provider is unavailable. It will be retried + // after the user configures or resets its LLM chain. + if task.Paused || !s.engine.ReadyFor(task) { continue } - t.Queued = false - log.Printf("[concurrency] task %s 出队启动", t.ID) - go s.startTaskEngine(t) + if !s.engine.beginTaskOperation(task.ID) { + continue + } + if err := s.m.SetTaskQueued(task.ID, false); err != nil { + s.engine.decInflight(task.ID) + continue + } + go func(t *Task) { + defer s.engine.decInflight(t.ID) + s.startTaskEngine(t) + }(task) slots-- } } @@ -153,7 +149,6 @@ func (s *Server) reviveTask(t *Task) { if t == nil { return } - // 排队中的任务不复活启动——它本就等待并发空位,交给 reconcileConcurrency。 if t.Queued { return } @@ -184,24 +179,25 @@ func (s *Server) reviveTask(t *Task) { // Returns the seeded specs so callers can emit activity records for them. // emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI. func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec { + if t == nil { + return nil + } // Use the SAME LLM the task runs on (its pinned profile, else the active profile), // NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars, // so FromEnv returned empty and every task silently fell back to the crude rule // splitter (which shredded URLs / made meaningless 2-way splits). var specs []goalSpec - if prov, ok := s.goalsProvider(t); ok { - var as *db.AssetStore - if s.m != nil { - as = s.m.Assets() - } - taskID, _ := strconv.ParseInt(t.ID, 10, 64) - // DecomposeGoals persists each decomposed goal via the set_goals tool straight - // into t.Store; it returns the specs it wrote so we can emit activity + spot the - // empty (LLM error / no provider) case below. No write-back needed here anymore. - for _, g := range agent.DecomposeGoals(ctx, prov, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) { - if strings.TrimSpace(g.Text) != "" { - specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) - } + var as *db.AssetStore + if s.m != nil { + as = s.m.Assets() + } + taskID, _ := strconv.ParseInt(t.ID, 10, 64) + s.engine.BeginLLMCall(t.ID) + decomposed := agent.DecomposeGoalsWithProvider(ctx, s.agentsForTask(t).runtime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) + s.engine.EndLLMCall(t.ID) + for _, g := range decomposed { + if strings.TrimSpace(g.Text) != "" { + specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) } } if len(specs) == 0 { @@ -220,23 +216,3 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity } return specs } - -// goalsProvider resolves the provider for goal decomposition by the standard -// precedence: the goals agent's own binding → the task's pinned profile → the -// global active provider (the very instance the engine runs on, so decomposition -// shares its rate limiter, gets recorded, and follows LLM failover). -func (s *Server) goalsProvider(t *Task) (llm.Provider, bool) { - var pin *int64 - if t != nil { - pin = t.LLMProfileID - } - if eff := s.effectiveProfileForAgent("goals", pin); eff != nil { - if prov, cfg, ok := s.providerForProfile(*eff); ok { - return s.poolForBinding(*eff, prov, cfg), true - } - } - s.cfgMu.Lock() - prov, on := s.llmProv, s.llmOn - s.cfgMu.Unlock() - return prov, on && prov != nil -} diff --git a/server/goals_test.go b/server/goals_test.go new file mode 100644 index 0000000..057d85b --- /dev/null +++ b/server/goals_test.go @@ -0,0 +1,12 @@ +package server + +import ( + "context" + "testing" +) + +func TestCreateGoalsNilTask(t *testing.T) { + if got := (&Server{}).createGoals(context.Background(), nil, nil); len(got) != 0 { + t.Fatalf("nil task produced goals: %+v", got) + } +} diff --git a/server/inheritance_api_test.go b/server/inheritance_api_test.go new file mode 100644 index 0000000..df9fe4f --- /dev/null +++ b/server/inheritance_api_test.go @@ -0,0 +1,100 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Autumn-27/artex/db" +) + +func TestInheritedActivityDetailAndRelationDeletion(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + + source, err := m.CreateTaskWithOptions("detail source", "source goal", db.TaskCreateOptions{}) + if err != nil { + t.Fatal(err) + } + sourceID := mustTaskID(t, source.ID) + current, err := m.CreateTaskWithOptions("detail current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{sourceID}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(current.ID, DeleteTaskOptions{}) + _, _ = m.DeleteTask(source.ID, DeleteTaskOptions{}) + }) + + plannerStep, err := source.Store.AppendActivity(db.Activity{ + Worker: "planner", Kind: "text", Summary: "private planner", Detail: "private planner detail", + }) + if err != nil { + t.Fatal(err) + } + intentID, err := source.Store.AddIntent(map[string]any{"summary": "completed source work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := source.Store.SetIntentState(intentID, "done"); err != nil { + t.Fatal(err) + } + workerStep, err := source.Store.AppendActivity(db.Activity{ + NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "shared result", Detail: "shared worker detail", + }) + if err != nil { + t.Fatal(err) + } + + s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir()) + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatal(err) + } + fetchDetail := func(stepID int64) string { + req := httptest.NewRequest(http.MethodGet, + fmt.Sprintf("/api/exploration/activity/%d?task=%s", stepID, current.ID), nil) + req.Header.Set("Authorization", "Bearer "+token) + rec := httptest.NewRecorder() + s.Handler().ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("activity detail %d: status=%d body=%s", stepID, rec.Code, rec.Body.String()) + } + var body struct { + Detail string `json:"detail"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + return body.Detail + } + if got := fetchDetail(plannerStep); got != "" { + t.Fatalf("source planner transcript leaked through API: %q", got) + } + if got := fetchDetail(workerStep); got != "shared worker detail" { + t.Fatalf("terminal inherited worker detail missing: %q", got) + } + + if _, err := m.DeleteTask(source.ID, DeleteTaskOptions{}); err != nil { + t.Fatal(err) + } + currentAfter, ok := m.Task(current.ID) + if !ok || len(currentAfter.SourceTaskIDs) != 0 { + t.Fatalf("deleted source remained in live task DTO state: %+v", currentAfter) + } +} + +func mustTaskID(t *testing.T, raw string) int64 { + t.Helper() + var id int64 + if _, err := fmt.Sscan(raw, &id); err != nil || id <= 0 { + t.Fatalf("bad task id %q: %v", raw, err) + } + return id +} diff --git a/server/inheritance_dto_test.go b/server/inheritance_dto_test.go new file mode 100644 index 0000000..078d233 --- /dev/null +++ b/server/inheritance_dto_test.go @@ -0,0 +1,117 @@ +package server + +import ( + "encoding/json" + "testing" + "time" + + "github.com/Autumn-27/artex/db" +) + +func TestInheritedDTOProvenance(t *testing.T) { + createdAt := time.Date(2026, time.August, 19, 12, 0, 0, 0, time.UTC) + node := &db.Node{ + ID: 41, + Kind: db.KindFinding, + Payload: json.RawMessage(`{"vulnclass":"xss","severity":"high","summary":"source finding"}`), + State: "confirmed", + CreatedAt: createdAt, + SourceTaskID: 7, + Inherited: true, + } + + nodeDTO := taskNodeDTO(node) + if !nodeDTO.Inherited || nodeDTO.SourceTaskID != "7" { + t.Fatalf("node provenance lost: %+v", nodeDTO) + } + finding := findingDTOsForOwner("7", "source", []*db.Node{node}, nil, nil) + if len(finding) != 1 || !finding[0].Inherited || finding[0].SourceTaskID != "7" || finding[0].TaskID != "7" { + t.Fatalf("finding provenance lost: %+v", finding) + } + + intentID := int64(44) + activity := activityDTO(db.Activity{ + ID: 45, + NodeID: &intentID, + CreatedAt: createdAt, + SourceTaskID: 7, + Inherited: true, + }) + if !activity.Inherited || activity.SourceTaskID != "7" || activity.IntentID != "44" { + t.Fatalf("activity provenance lost: %+v", activity) + } + assetRef := coverageAssetRefDTO(db.AssetRef{ + ID: 46, Kind: db.KindFact, State: "confirmed", Summary: "source fact", SourceTaskID: 7, Inherited: true, + }) + if !assetRef.Inherited || assetRef.SourceTaskID != "7" { + t.Fatalf("asset ref provenance lost: %+v", assetRef) + } +} + +func TestInheritedIntentResultOnlyAllowsHistory(t *testing.T) { + for _, state := range []string{"done", "blocked", "exhausted", "stopped"} { + if !inheritedIntentResult(state) { + t.Errorf("terminal state %q should be inherited history", state) + } + } + for _, state := range []string{"open", "running", ""} { + if inheritedIntentResult(state) { + t.Errorf("live state %q must not enter inherited history", state) + } + } +} + +func TestInheritedGraphSnapshotFiltersLiveIntentsAndEdges(t *testing.T) { + nodes := []*db.Node{ + {ID: 1, Kind: db.KindFact, State: "confirmed"}, + {ID: 2, Kind: db.KindIntent, State: "running"}, + {ID: 3, Kind: db.KindIntent, State: "done"}, + {ID: 4, Kind: db.KindFinding, State: "confirmed"}, + } + edges := []db.Edge{ + {From: 1, Rel: db.RelDerivedFrom, To: 2}, + {From: 1, Rel: db.RelDerivedFrom, To: 3}, + {From: 3, Rel: db.RelYields, To: 4}, + } + + gotNodes, gotEdges := inheritedGraphSnapshot(nodes, edges, 77) + if len(gotNodes) != 3 { + t.Fatalf("expected fact, terminal intent and finding; got %+v", gotNodes) + } + for _, node := range gotNodes { + if node.ID == 2 { + t.Fatal("running inherited intent leaked into graph snapshot") + } + if !node.Inherited || node.SourceTaskID != 77 { + t.Fatalf("source provenance missing: %+v", node) + } + } + if len(gotEdges) != 2 { + t.Fatalf("edge referencing hidden live intent was not removed: %+v", gotEdges) + } + for _, edge := range gotEdges { + if edge.From == 2 || edge.To == 2 { + t.Fatalf("live intent id leaked through edge: %+v", edge) + } + } +} + +func TestFindingProvenanceInTask(t *testing.T) { + task := &Task{ID: "10", SourceTaskIDs: []int64{7, 8}} + + local := int64(10) + if source, inherited, allowed := findingProvenanceInTask(task, &local); !allowed || inherited || source != "" { + t.Fatalf("local finding: source=%q inherited=%v allowed=%v", source, inherited, allowed) + } + sourceID := int64(7) + if source, inherited, allowed := findingProvenanceInTask(task, &sourceID); !allowed || !inherited || source != "7" { + t.Fatalf("direct source finding: source=%q inherited=%v allowed=%v", source, inherited, allowed) + } + unrelated := int64(6) + if _, _, allowed := findingProvenanceInTask(task, &unrelated); allowed { + t.Fatal("unrelated finding must not be readable through task context") + } + if _, _, allowed := findingProvenanceInTask(task, nil); allowed { + t.Fatal("detached finding must not be readable through task context") + } +} diff --git a/server/manager.go b/server/manager.go index 78e403d..5f8dc36 100644 --- a/server/manager.go +++ b/server/manager.go @@ -2,6 +2,7 @@ package server import ( "encoding/json" + "errors" "fmt" "log" "os" @@ -12,38 +13,44 @@ import ( "sync" "time" - actool "github.com/Autumn-27/norma/tool" "github.com/Autumn-27/artex/agent" pgdb "github.com/Autumn-27/artex/db" "github.com/Autumn-27/artex/enrich" "github.com/Autumn-27/artex/guard" "github.com/Autumn-27/artex/intercept" "github.com/Autumn-27/artex/traffic" + actool "github.com/Autumn-27/norma/tool" ) // Task is one engagement: a description + goal + its own exploration store, // sharing the process-wide asset store. ID is the PG task id as a string; ExpID // is the exploration the task owns. type Task struct { - ID string `json:"id"` - ExpID int64 `json:"exploration_id"` - Description string `json:"description"` - Goal string `json:"goal"` - CreatedAt int64 `json:"created_at"` - CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成 - Paused bool `json:"paused"` - Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 - ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录) - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置 - Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态;空/其它则由运行态推导) + ID string `json:"id"` + ExpID int64 `json:"exploration_id"` + Description string `json:"description"` + Goal string `json:"goal"` + CreatedAt int64 `json:"created_at"` + CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成 + Paused bool `json:"paused"` + Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 + ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录) + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置 + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMFailoverState string `json:"llm_failover_state,omitempty"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` + Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态;空/其它则由运行态推导) // 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。 - TimeoutSeconds int `json:"timeout_seconds"` - PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒) - FirstRunAt int64 `json:"first_run_at,omitempty"` - DeadlineAt int64 `json:"deadline_at,omitempty"` - Store *pgdb.ExplorationStore `json:"-"` - Guard *guard.Guard `json:"-"` - notify chan struct{} + TimeoutSeconds int `json:"timeout_seconds"` + PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒) + FirstRunAt int64 `json:"first_run_at,omitempty"` + DeadlineAt int64 `json:"deadline_at,omitempty"` + Store *pgdb.ExplorationStore `json:"-"` + Guard *guard.Guard `json:"-"` + notify chan struct{} + llmMu sync.RWMutex // pendingTriggers accumulates the concrete changes (worker done / finding) that // fired planning rounds since the last one consumed them. The debounce coalesces @@ -52,14 +59,74 @@ type Task struct { pendingTriggers []agent.TriggerEvent } +type taskLLMState struct { + ProfileID *int64 + ProfileIDs []int64 + ActiveID *int64 + FailoverState string + FailoverReason string +} + +func cloneInt64Ptr(value *int64) *int64 { + if value == nil { + return nil + } + copy := *value + return © +} + +func (t *Task) llmStateSnapshot() taskLLMState { + t.llmMu.RLock() + defer t.llmMu.RUnlock() + return taskLLMState{ + ProfileID: cloneInt64Ptr(t.LLMProfileID), + ProfileIDs: append(make([]int64, 0, len(t.LLMProfileIDs)), t.LLMProfileIDs...), + ActiveID: cloneInt64Ptr(t.ActiveLLMProfileID), + FailoverState: t.LLMFailoverState, + FailoverReason: t.LLMFailoverReason, + } +} + +func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, state, reason string) { + t.llmMu.Lock() + t.LLMProfileID = cloneInt64Ptr(profileID) + t.ActiveLLMProfileID = cloneInt64Ptr(activeID) + t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...) + t.LLMFailoverState = state + t.LLMFailoverReason = reason + t.llmMu.Unlock() +} + +// DeleteTaskOptions controls cleanup of data stored outside the task's own +// exploration graph. All options default to false for backward compatibility. +type DeleteTaskOptions struct { + DeleteAssets bool `json:"delete_assets"` + DeleteTraffic bool `json:"delete_traffic"` + DeleteFiles bool `json:"delete_files"` + DeleteFindings bool `json:"delete_findings"` + DeleteLLMRecords bool `json:"delete_llm_records"` +} + +// DeleteTaskResult makes destructive cleanup auditable to API callers. +type DeleteTaskResult struct { + Deleted string `json:"deleted"` + AssetsDeleted int64 `json:"assets_deleted"` + AssetsDetached int64 `json:"assets_detached"` + TrafficDeleted int64 `json:"traffic_deleted"` + FilesDeleted bool `json:"files_deleted"` + FindingsDeleted int64 `json:"findings_deleted"` + LLMRecordsDeleted int64 `json:"llm_records_deleted"` + CleanupWarning string `json:"cleanup_warning,omitempty"` +} + // Manager owns the PostgreSQL data source (asset graph + every task's exploration // graph + config) and the in-memory set of task handles. type Manager struct { dir string pg *pgdb.DB assets *pgdb.AssetStore - traffic *traffic.Traffic // process-wide recording proxy (may be nil) - enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP) + traffic *traffic.Traffic // process-wide recording proxy (may be nil) + enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP) interceptor *intercept.Interceptor // user-configured tool-call interception rules mu sync.RWMutex @@ -529,7 +596,10 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID, Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued, CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef, - LLMProfileID: pt.LLMProfileID, + LLMProfileID: pt.LLMProfileID, + LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID, + LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason, + SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...), TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1), @@ -539,7 +609,17 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter // CreateTask creates a task + its exploration and makes it active. // timeoutSeconds is the task-level wall-clock budget (0 = 不限时). func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { - pt, err := m.pg.CreateTask(description, goal, llmProfileID, timeoutSeconds, planHeartbeatSeconds) + var ids []int64 + if llmProfileID != nil { + ids = []int64{*llmProfileID} + } + return m.CreateTaskWithOptions(description, goal, pgdb.TaskCreateOptions{ + LLMProfileIDs: ids, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds, + }) +} + +func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.TaskCreateOptions) (*Task, error) { + pt, err := m.pg.CreateTaskWithOptions(description, goal, opts) if err != nil { return nil, err } @@ -551,6 +631,35 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time return t, nil } +// ReplaceTaskLLMProfiles resets a non-terminal task's ordered provider chain and +// mirrors the committed state onto the live task handle. +func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activeProfileID int64) (int64, error) { + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return 0, err + } + if err := m.pg.ReplaceTaskLLMProfiles(n, profileIDs, activeProfileID); err != nil { + return 0, err + } + pt, err := m.pg.GetTask(n) + if err != nil || pt == nil { + return 0, err + } + m.mu.Lock() + if task := m.tasks[id]; task != nil { + task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason) + } + m.mu.Unlock() + if task, ok := m.Task(id); ok { + reopened, reopenErr := task.Store.ReopenIntentsByBlockedReason(pgdb.IntentBlockedLLMQuota) + if reopenErr != nil { + return reopened, reopenErr + } + return reopened, nil + } + return 0, nil +} + // LoadExisting rebuilds in-memory task handles from the PG task registry. func (m *Manager) LoadExisting() []*Task { pts, err := m.pg.ListTasks() @@ -706,17 +815,123 @@ func (m *Manager) SetTaskStatus(id, status string) error { return nil } -// DeleteTask removes a task (and its exploration subgraph) from PG + memory. -func (m *Manager) DeleteTask(id string) error { +// DeleteTask removes a task and optionally its related global data. Traffic has +// no task-id column, so related exchanges are resolved by exact hosts from the +// task's asset rows. Files are staged before the database operation; traffic is +// staged while PostgreSQL excludes asset/anchor writers. Both are restored on a +// database failure and purged only after its commit. +func (m *Manager) DeleteTask(id string, opts DeleteTaskOptions) (DeleteTaskResult, error) { + result := DeleteTaskResult{Deleted: id} n, err := strconv.ParseInt(id, 10, 64) if err != nil { - return err + return result, err } - if err := m.pg.DeleteTask(n); err != nil { - return err + registered, err := m.pg.GetTask(n) + if err != nil { + return result, err } + if registered == nil { + m.forgetTask(id, n) + return result, nil + } + + var fileStage *taskFileDeleteStage + if opts.DeleteFiles { + fileStage, err = stageTaskFiles(m.dir, id, registered.ExplorationID) + if err != nil { + return result, err + } + result.FilesDeleted = fileStage.deleted + } + + var trafficStage *traffic.HostDeleteStage + var prepare func(pgdb.TaskDeletePreparation) error + if opts.DeleteTraffic && m.traffic != nil { + prepare = func(p pgdb.TaskDeletePreparation) error { + if len(p.TrafficHosts) == 0 { + return nil + } + trafficStage, err = m.traffic.StageDeleteHostsExact(p.TrafficHosts) + if err != nil { + return err + } + result.TrafficDeleted = trafficStage.Deleted() + return nil + } + } + + dbResult, err := m.pg.DeleteTaskCascadePrepared( + n, opts.DeleteAssets, opts.DeleteFindings, opts.DeleteLLMRecords, prepare, + ) + if err != nil { + return result, rollbackTaskDelete(err, trafficStage, fileStage) + } + result.AssetsDeleted = dbResult.AssetsDeleted + result.AssetsDetached = dbResult.AssetsDetached + result.FindingsDeleted = dbResult.FindingsDeleted + result.LLMRecordsDeleted = dbResult.LLMRecordsDeleted + + // PostgreSQL is now authoritative: finalize the staged external deletion and + // forget the live task even if a final purge reports an error. Such errors are + // typed so the HTTP layer can still tear down the task runtime instead of + // incorrectly reviving a task whose database row is already gone. + var finalizeErrs []error + if trafficStage != nil { + if err := trafficStage.Commit(); err != nil { + finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize traffic deletion: %w", err)) + } + } + if fileStage != nil { + if err := fileStage.commit(); err != nil { + finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize task file deletion: %w", err)) + } + } + m.forgetTask(id, n) + if err := errors.Join(finalizeErrs...); err != nil { + return result, &taskDeleteCommittedError{err: err} + } + return result, nil +} + +// taskDeleteCommittedError means PostgreSQL deletion succeeded but purging one +// of the recoverable staging directories failed. The task must stay deleted. +type taskDeleteCommittedError struct{ err error } + +func (e *taskDeleteCommittedError) Error() string { + return "task deletion committed; external cleanup incomplete: " + e.err.Error() +} + +func (e *taskDeleteCommittedError) Unwrap() error { return e.err } + +func rollbackTaskDelete(cause error, trafficStage *traffic.HostDeleteStage, fileStage *taskFileDeleteStage) error { + errs := []error{cause} + // Reverse the preparation order. Both restorations are attempted even if the + // first one fails, and errors.Join preserves the original PostgreSQL error. + if trafficStage != nil { + if err := trafficStage.Rollback(); err != nil { + errs = append(errs, fmt.Errorf("restore traffic after task delete failure: %w", err)) + } + } + if fileStage != nil { + if err := fileStage.rollback(); err != nil { + errs = append(errs, fmt.Errorf("restore task files after task delete failure: %w", err)) + } + } + return errors.Join(errs...) +} + +func (m *Manager) forgetTask(id string, numericID int64) { m.mu.Lock() delete(m.tasks, id) + for _, task := range m.tasks { + kept := task.SourceTaskIDs[:0] + for _, sourceID := range task.SourceTaskIDs { + if sourceID != numericID { + kept = append(kept, sourceID) + } + } + task.SourceTaskIDs = kept + } if m.active == id { m.active = "" for _, t := range m.tasks { @@ -725,7 +940,128 @@ func (m *Manager) DeleteTask(id string) error { } } m.mu.Unlock() - return nil +} + +type stagedTaskPath struct { + source string + staged string +} + +type taskFileDeleteStage struct { + stageDir string + moves []stagedTaskPath + deleted bool + done bool +} + +// stageTaskFiles atomically renames the task workspace and owned transcripts to +// a same-filesystem staging directory. The trailing dash in the transcript +// prefix is significant: exploration 12 must not match exploration 123. +func stageTaskFiles(dataDir, taskID string, explorationID int64) (*taskFileDeleteStage, error) { + stage := &taskFileDeleteStage{} + var targets []string + taskDir := filepath.Join(dataDir, "tasks", taskID) + if _, err := os.Lstat(taskDir); err == nil { + targets = append(targets, taskDir) + } else if !os.IsNotExist(err) { + return nil, err + } + + transcriptDir := filepath.Join(dataDir, "transcripts") + entries, err := os.ReadDir(transcriptDir) + if err != nil { + if !os.IsNotExist(err) { + return nil, err + } + entries = nil + } + prefix := fmt.Sprintf("exp%d-", explorationID) + for _, entry := range entries { + name := entry.Name() + if !strings.HasPrefix(name, prefix) || (!entry.IsDir() && !strings.HasSuffix(name, ".jsonl")) { + continue + } + targets = append(targets, filepath.Join(transcriptDir, name)) + } + if len(targets) == 0 { + stage.done = true + return stage, nil + } + + parent := filepath.Join(dataDir, ".delete-staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return nil, err + } + stage.stageDir, err = os.MkdirTemp(parent, "task-"+taskID+"-") + if err != nil { + return nil, err + } + for _, source := range targets { + staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source))) + if err := os.Rename(source, staged); err != nil { + cause := fmt.Errorf("stage task file %s: %w", source, err) + if restoreErr := stage.rollback(); restoreErr != nil { + return nil, errors.Join(cause, fmt.Errorf("restore partially staged task files: %w", restoreErr)) + } + return nil, cause + } + stage.moves = append(stage.moves, stagedTaskPath{source: source, staged: staged}) + } + stage.deleted = true + return stage, nil +} + +func (s *taskFileDeleteStage) commit() error { + if s == nil || s.done { + return nil + } + err := os.RemoveAll(s.stageDir) + s.done = true + return err +} + +func (s *taskFileDeleteStage) rollback() error { + if s == nil || s.done { + return nil + } + var errs []error + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.MkdirAll(filepath.Dir(move.source), 0o755); err != nil { + errs = append(errs, fmt.Errorf("create restore parent for %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove task file stage: %w", err)) + } + } + s.done = true + return errors.Join(errs...) +} + +// deleteTaskFiles retains the standalone helper contract used by focused tests. +func deleteTaskFiles(dataDir, taskID string, explorationID int64) (bool, error) { + stage, err := stageTaskFiles(dataDir, taskID, explorationID) + if err != nil { + return false, err + } + deleted := stage.deleted + if err := stage.commit(); err != nil { + return deleted, err + } + return deleted, nil } func (m *Manager) Task(id string) (*Task, bool) { diff --git a/server/manager_delete_files_test.go b/server/manager_delete_files_test.go new file mode 100644 index 0000000..48de983 --- /dev/null +++ b/server/manager_delete_files_test.go @@ -0,0 +1,280 @@ +package server + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + pgdb "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +func TestSeedAssociatesTargetAssetWithTask(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + task, err := m.CreateTask("seed ownership", "seed ownership", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true}) + }) + + host := fmt.Sprintf("seed-%s.example.test", task.ID) + s := &Server{m: m} + s.seed(task, "https://"+host) + taskID := mustTaskID(t, task.ID) + assets, err := m.Assets().QueryByTask(taskID, "", 20, 0) + if err != nil { + t.Fatal(err) + } + found := false + for _, asset := range assets { + if asset.Domain == host || asset.URL == "https://"+host { + found = true + break + } + } + if !found { + t.Fatalf("seed target was not associated with task %s: %+v", task.ID, assets) + } +} + +func TestDeleteTaskFilesRemovesOnlyOwnedWorkspaceAndTranscripts(t *testing.T) { + dataDir := t.TempDir() + + removed := []struct { + path string + isDir bool + }{ + {path: filepath.Join("tasks", "77"), isDir: true}, + {path: filepath.Join("transcripts", "exp12-main.jsonl")}, + {path: filepath.Join("transcripts", "exp12-planner.jsonl")}, + {path: filepath.Join("transcripts", "exp12-worker-i9.jsonl")}, + {path: filepath.Join("transcripts", "exp12-main"), isDir: true}, + {path: filepath.Join("transcripts", "exp12-worker-i9"), isDir: true}, + } + kept := []struct { + path string + isDir bool + }{ + {path: filepath.Join("tasks", "78"), isDir: true}, + {path: filepath.Join("transcripts", "exp123-main.jsonl")}, + {path: filepath.Join("transcripts", "exp123-main"), isDir: true}, + {path: filepath.Join("transcripts", "exp1-main.jsonl")}, + {path: filepath.Join("transcripts", "exp12main.jsonl")}, + {path: filepath.Join("transcripts", "exp12-main.jsonl.bak")}, + {path: filepath.Join("transcripts", "exp12-main.txt")}, + {path: filepath.Join("transcripts", "conv-12.jsonl")}, + {path: filepath.Join("transcripts", "unrelated-session"), isDir: true}, + } + + for _, fixture := range append(removed, kept...) { + createDeleteFixture(t, dataDir, fixture.path, fixture.isDir) + } + + deleted, err := deleteTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatalf("delete task files: %v", err) + } + if !deleted { + t.Fatal("expected deletion to be reported") + } + for _, fixture := range removed { + assertPathMissing(t, filepath.Join(dataDir, fixture.path)) + } + for _, fixture := range kept { + assertPathExists(t, filepath.Join(dataDir, fixture.path)) + } +} + +func TestDeleteTaskFilesReportsTranscriptOnlyDeletion(t *testing.T) { + dataDir := t.TempDir() + path := filepath.Join("transcripts", "exp44-main.jsonl") + createDeleteFixture(t, dataDir, path, false) + + deleted, err := deleteTaskFiles(dataDir, "missing-task", 44) + if err != nil { + t.Fatalf("delete task files: %v", err) + } + if !deleted { + t.Fatal("expected transcript deletion to be reported") + } + assertPathMissing(t, filepath.Join(dataDir, path)) +} + +func TestDeleteTaskFilesMissingTargetsIsIdempotent(t *testing.T) { + deleted, err := deleteTaskFiles(t.TempDir(), "77", 12) + if err != nil { + t.Fatalf("delete missing task files: %v", err) + } + if deleted { + t.Fatal("missing targets must not be reported as deleted") + } +} + +func TestStageTaskFilesRollbackRestoresWorkspaceAndTranscripts(t *testing.T) { + dataDir := t.TempDir() + paths := []string{ + filepath.Join("tasks", "77", "notes.txt"), + filepath.Join("transcripts", "exp12-main.jsonl"), + filepath.Join("transcripts", "exp12-worker-i9", "sidechain.jsonl"), + } + for _, path := range paths { + createDeleteFixture(t, dataDir, path, false) + } + + stage, err := stageTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatal(err) + } + if !stage.deleted { + t.Fatal("expected files to be staged") + } + for _, path := range paths { + assertPathMissing(t, filepath.Join(dataDir, path)) + } + if err := stage.rollback(); err != nil { + t.Fatal(err) + } + for _, path := range paths { + assertPathExists(t, filepath.Join(dataDir, path)) + } +} + +func TestStageTaskFilesRollbackReportsRestoreFailure(t *testing.T) { + dataDir := t.TempDir() + workspace := filepath.Join(dataDir, "tasks", "77") + createDeleteFixture(t, dataDir, filepath.Join("tasks", "77"), true) + + stage, err := stageTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(workspace, []byte("conflict"), 0o644); err != nil { + t.Fatal(err) + } + cause := errors.New("forced PostgreSQL delete failure") + err = rollbackTaskDelete(cause, nil, stage) + if err == nil || !errors.Is(err, cause) || !strings.Contains(err.Error(), "destination already exists") { + t.Fatalf("rollback err=%v, want original and restore errors", err) + } + if _, statErr := os.Stat(stage.stageDir); statErr != nil { + t.Fatalf("staging was removed after failed restore: %v", statErr) + } +} + +func TestManagerDeleteTaskRestoresFilesAndTrafficWhenPostgresDeleteFails(t *testing.T) { + dataDir := t.TempDir() + m, err := NewManager(dataDir, "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + t.Cleanup(func() { _ = m.Close() }) + + task, err := m.CreateTask("delete rollback", "delete rollback", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + taskID := mustTaskID(t, task.ID) + triggerName := fmt.Sprintf("test_fail_task_delete_%d", taskID) + functionName := triggerName + "_fn" + cleanupDB := func() { + _, _ = m.pg.Exec(fmt.Sprintf(`DROP TRIGGER IF EXISTS %s ON tasks`, triggerName)) + _, _ = m.pg.Exec(fmt.Sprintf(`DROP FUNCTION IF EXISTS %s()`, functionName)) + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true, DeleteTraffic: true, DeleteFiles: true}) + } + t.Cleanup(cleanupDB) + + host := fmt.Sprintf("delete-rollback-%d.example.test", taskID) + if _, err := m.Assets().UpsertHTTPService(pgdb.UpsertHTTPServiceReq{URL: "https://" + host, TaskID: taskID}); err != nil { + t.Fatal(err) + } + createDeleteFixture(t, dataDir, filepath.Join("tasks", task.ID), true) + transcriptPath := filepath.Join(dataDir, "transcripts", fmt.Sprintf("exp%d-main.jsonl", task.ExpID)) + createDeleteFixture(t, dataDir, filepath.Join("transcripts", filepath.Base(transcriptPath)), false) + + tr, err := traffic.Open(filepath.Join(dataDir, "traffic"), "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + m.traffic = tr + trafficTree := filepath.Join(dataDir, "traffic", host, "GET", "1-0001") + if err := os.MkdirAll(trafficTree, 0o755); err != nil { + t.Fatal(err) + } + trafficMarker := filepath.Join(trafficTree, "request.http") + if err := os.WriteFile(trafficMarker, []byte("original traffic"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, "1-0001", 1, host, "GET", "/", "https://"+host+"/", 200, "text/plain", 0, 0, host+"/GET/1-0001"); err != nil { + t.Fatal(err) + } + + if _, err := m.pg.Exec(fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $body$ +BEGIN RAISE EXCEPTION 'forced task delete failure'; END $body$`, functionName)); err != nil { + t.Fatal(err) + } + if _, err := m.pg.Exec(fmt.Sprintf(`CREATE TRIGGER %s BEFORE DELETE ON tasks +FOR EACH ROW WHEN (OLD.id = %d) EXECUTE FUNCTION %s()`, triggerName, taskID, functionName)); err != nil { + t.Fatal(err) + } + + _, err = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteTraffic: true, DeleteFiles: true}) + if err == nil || !strings.Contains(err.Error(), "forced task delete failure") { + t.Fatalf("delete err=%v, want injected PostgreSQL failure", err) + } + if got, err := m.pg.GetTask(taskID); err != nil || got == nil { + t.Fatalf("task row was lost after failed delete: task=%+v err=%v", got, err) + } + assertPathExists(t, filepath.Join(dataDir, "tasks", task.ID)) + assertPathExists(t, transcriptPath) + if got, err := os.ReadFile(trafficMarker); err != nil || string(got) != "original traffic" { + t.Fatalf("traffic tree was not restored: content=%q err=%v", got, err) + } + var trafficRows int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&trafficRows); err != nil || trafficRows != 1 { + t.Fatalf("traffic index was not restored: rows=%d err=%v", trafficRows, err) + } +} + +func createDeleteFixture(t *testing.T, root, relative string, isDir bool) { + t.Helper() + path := filepath.Join(root, relative) + if isDir { + if err := os.MkdirAll(filepath.Join(path, "subagents"), 0o755); err != nil { + t.Fatalf("create fixture directory %s: %v", path, err) + } + if err := os.WriteFile(filepath.Join(path, "subagents", "agent-test.jsonl"), []byte("test"), 0o644); err != nil { + t.Fatalf("create fixture sidechain %s: %v", path, err) + } + return + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatalf("create fixture parent %s: %v", path, err) + } + if err := os.WriteFile(path, []byte("test"), 0o644); err != nil { + t.Fatalf("create fixture file %s: %v", path, err) + } +} + +func assertPathMissing(t *testing.T, path string) { + t.Helper() + if _, err := os.Lstat(path); !os.IsNotExist(err) { + t.Fatalf("expected %s to be removed, stat err=%v", path, err) + } +} + +func assertPathExists(t *testing.T, path string) { + t.Helper() + if _, err := os.Lstat(path); err != nil { + t.Fatalf("expected %s to remain: %v", path, err) + } +} diff --git a/server/orchestration.go b/server/orchestration.go index 011b0f6..9ed4446 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -199,12 +199,13 @@ func (s *Server) toolListTasks() actool.CoreTool { if t.ParentRef != "" { row["parent_ref"] = t.ParentRef } - if t.LLMProfileID == nil { + llmState := t.llmStateSnapshot() + if llmState.ProfileID == nil { row["llm_profile"] = "(激活配置)" - } else if n, ok := profName[*t.LLMProfileID]; ok { + } else if n, ok := profName[*llmState.ProfileID]; ok { row["llm_profile"] = n } else { - row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID) + row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *llmState.ProfileID) } out = append(out, row) } @@ -305,6 +306,7 @@ func (s *Server) toolPauseTask() actool.CoreTool { return actool.Errorf("task 不存在: " + a.TaskID), nil } s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator) + s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask) return actool.Text("task paused: " + a.TaskID), nil }) } @@ -419,7 +421,7 @@ func (s *Server) deriveTaskStatus(t *Task) string { return t.Status case t.Paused || s.engine.IsPaused(t.ID): return "paused" - case s.engine.Ready() && s.engine.Started(t.ID): + case s.engine.ReadyFor(t) && s.engine.Started(t.ID): return "running" } return "created" diff --git a/server/server.go b/server/server.go index 2fe0d1d..2f44d26 100644 --- a/server/server.go +++ b/server/server.go @@ -12,6 +12,7 @@ import ( "net/url" "path/filepath" "regexp" + "sort" "strconv" "strings" "sync" @@ -51,14 +52,12 @@ type Server struct { cfgMu sync.Mutex mainAgent *agent.MainAgent // nil when no LLM provider is configured chatAgent *agent.ChatAgent // conversational runner for the chat page; nil w/o LLM - llmCfg agent.Config // current LLM config (key not exposed) - llmOn bool - llmProf string // active LLM profile name (for llmrec tagging) // llmProv is the fully-decorated global provider (recorder + failover chain) - // applyLLM installed. Kept so one-off callers (goal decomposition) reuse the - // same instance instead of building a private one that skips the rate limiter, - // the recorder and failover. + // installed by applyLLM. Task routers use it after an explicit chain is cleared. llmProv llm.Provider + llmCfg agent.Config // current LLM config (key not exposed) + llmOn bool + llmProf string // active LLM profile name (for llmrec tagging) // chatBusy guards the per-task main-agent run: the chat handler launches the // agent on the server's background ctx (not the request ctx) and returns @@ -97,18 +96,25 @@ type Server struct { // provByProfile caches ONE provider per LLM profile id so every agent bound or // pinned to the same profile shares a single provider instance — hence one rate - // limiter. Separate from profAgents (whole planner/worker pairs). The globally-active - // provider (built in applyLLM) lives OUTSIDE this cache, so binding an agent to the - // currently-active profile builds a second instance — harmless (that binding equals - // leaving it unset), just not deduped. Cleared alongside profAgents on profile edits. + // limiter. Separate from profAgents (whole planner/worker pairs). applyLLM also + // resolves the persisted global-active profile through this cache, so global + // fallback and task chains do not accidentally double the configured request rate. + // Cleared alongside profAgents on profile edits, then repopulated by active reapply. provCacheMu sync.Mutex provByProfile map[int64]*provEntry + provCacheGen uint64 // llmHealth is the process-wide circuit-breaker state for LLM failover (轮询). // It deliberately lives OUTSIDE the provider caches: rebuilding the chain // (saving an unrelated profile, flipping a setting) must not erase what we // learned about which backends are out of credit / rate-limited. llmHealth *llmpool.Registry + + // Explicit task chains use one task-scoped dynamic provider shared by goals, + // planner, workers, and the main agent. Bundles are stable; their runtime reads + // the persisted current profile at every new LLM call. + taskAgentMu sync.Mutex + taskAgents map[string]*taskAgentBundle } // profBundle is a planner/worker pair built from one LLM profile. @@ -143,14 +149,17 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi chatCancel: map[string]context.CancelCauseFunc{}, triggerQ: map[string][]triggeredRun{}, triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{}, profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{}, - provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg)} - // per-task LLM: a task pinned to a specific profile runs on that profile's - // dedicated planner/worker; unpinned tasks fall back to the global active pair. - s.engine.SetAgentResolver(func(t *Task) (*agent.Planner, *agent.Worker) { - if t.LLMProfileID == nil { + provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg), + taskAgents: map[string]*taskAgentBundle{}} + // Every task uses a stable task router. An empty explicit chain is resolved by + // that router through Agent bindings and then the global provider, so adding a + // first chain to a running task takes effect on its very next LLM call. + s.engine.SetAuthoritativeAgentResolver(func(t *Task) (*agent.Planner, *agent.Worker) { + if !s.taskRuntimeAvailable(t, "planner", "worker") { return nil, nil } - return s.agentsForProfile(*t.LLMProfileID) + b := s.agentsForTask(t) + return b.pl, b.wk }) // Wire DB-stored prompt templates into the agents (新版方案 §3.3 / §5a). With no // override row, agents keep their built-in defaults — behavior is unchanged. @@ -297,7 +306,7 @@ func (s *Server) loadLLMConfig() (agent.Config, bool) { } // saveLLMConfig persists the LLM config as the active "default" profile in PG. -func (s *Server) saveLLMConfig(cfg agent.Config) { +func (s *Server) saveLLMConfig(cfg agent.Config) error { format := cfg.Provider() if format != "anthropic" { format = "openai" @@ -322,9 +331,10 @@ func (s *Server) saveLLMConfig(cfg agent.Config) { ContextWindowK: cfg.ContextWindowK, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, Priority: priority, PoolExclude: poolExclude, }) - if err == nil { - _ = s.m.pg.SetActiveProfile(newID) + if err != nil { + return err } + return s.m.pg.SetActiveProfile(newID) } // reapplyActiveProfile hot-reloads the engine from the active DB profile so that @@ -380,15 +390,26 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent // applyLLM (re)builds the planner/worker/main-agent from cfg and installs them on // the running engine as the GLOBAL active pair. Safe to call at runtime (UI configures LLM). func (s *Server) applyLLM(cfg agent.Config) error { - prov, err := cfg.NewProvider() - if err != nil { - return err + var prov llm.Provider + // A persisted active profile must use the same cached provider as task chains + // and Agent bindings, otherwise each path owns a separate rate limiter. + if active, _ := s.m.pg.ActiveProfile(); active != nil { + if activeCfg, ok := s.loadProfileConfig(active.ID); ok && activeCfg == cfg { + prov, _, _ = s.providerForProfile(active.ID) + } + } + if prov == nil { + var err error + prov, err = cfg.NewProvider() + if err != nil { + return err + } + // Non-persisted/env configs do not have a profile cache key. + s.cfgMu.Lock() + profName := s.llmProf + s.cfgMu.Unlock() + prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) } - // Wrap provider with LLM call recorder (persists request/response to PG). - s.cfgMu.Lock() - profName := s.llmProf - s.cfgMu.Unlock() - prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) // LLM 轮询(默认关):把激活配置包进故障转移链,当前配置不可用时自动切下一个。 // 只影响「走全局激活配置」的这条路径——agent 绑定 / 任务 pin 的走 providerForProfile, // 默认仍然独占该配置(见 poolForBinding)。关闭或无备选时返回原 provider,行为不变。 @@ -414,10 +435,11 @@ func (s *Server) applyLLM(cfg agent.Config) error { s.chatAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) s.chatAgent.SetWebSearch(s.m.WebSearchOpts()) s.chatAgent.SetGuard(s.chatGuard()) - s.llmCfg = cfg s.llmProv = prov + s.llmCfg = cfg s.llmOn = true s.cfgMu.Unlock() + s.invalidateTaskAgents() // wake the active task so a task created while idle starts exploring. if t := s.m.ActiveTask(); t != nil { @@ -466,6 +488,7 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) s.provCacheMu.Unlock() return e.prov, e.cfg, true } + generation := s.provCacheGen s.provCacheMu.Unlock() cfg, ok := s.loadProfileConfig(id) if !ok { @@ -481,6 +504,10 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, s.m.LLMRecordEnabled) } s.provCacheMu.Lock() + if generation != s.provCacheGen { + s.provCacheMu.Unlock() + return s.providerForProfile(id) + } if e := s.provByProfile[id]; e != nil { // lost the race → keep the winner prov, cfg = e.prov, e.cfg } else { @@ -571,8 +598,10 @@ func (s *Server) invalidateProfileAgents() { s.profChatAgents = map[int64]*agent.ChatAgent{} s.profMu.Unlock() s.provCacheMu.Lock() + s.provCacheGen++ s.provByProfile = map[int64]*provEntry{} s.provCacheMu.Unlock() + s.invalidateTaskAgents() } func (s *Server) mainAgentRef() *agent.MainAgent { @@ -619,6 +648,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /api/workspace/upload", s.wsUpload) mux.HandleFunc("GET /api/tasks/{id}/scope", s.taskScopeList) mux.HandleFunc("POST /api/tasks/{id}/control", s.control) + mux.HandleFunc("PUT /api/tasks/{id}/llm", s.updateTaskLLMProfiles) + mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/control", s.controlIntent) mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/rerun", s.rerunIntent) // 重跑单条 blocked/exhausted/stopped 意图 mux.HandleFunc("POST /api/tasks/{id}/intents/rerun-blocked", s.rerunBlocked) // 批量重跑本任务全部 blocked 意图 mux.HandleFunc("POST /api/active", s.setActive) @@ -821,6 +852,7 @@ func (s *Server) stats(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, out) // no task selected: only global fields return } + out["llm_configured"] = s.engine.ReadyFor(t) st, _ := t.Store.Stats() out["exploration"] = st @@ -842,27 +874,31 @@ func (s *Server) stats(w http.ResponseWriter, r *http.Request) { } last := s.engine.LastActivity(t.ID) paused := s.engine.IsPaused(t.ID) - running := s.engine.Ready() && s.engine.Started(t.ID) && !paused - // stalled: running but no activity for a while and nothing in flight. - stalled := running && inFlight == 0 && last > 0 && time.Now().Unix()-last > 120 + activeCalls := s.engine.ActiveLLMCalls(t.ID) + running := s.engine.ReadyFor(t) && s.engine.Started(t.ID) && !paused + // A live event-driven engine with no work is healthy idle. Only a persisted + // running intent without any corresponding LLM call can be considered stalled. + stalled := running && activeCalls == 0 && inFlight > 0 && last > 0 && time.Now().Unix()-last > 120 // engine_mode follows the spec enum (exploring|paused|stalled|idle). + engineMode := "idle" switch { case paused: - out["engine_mode"] = "paused" + engineMode = "paused" case stalled: - out["engine_mode"] = "stalled" - case running: - out["engine_mode"] = "exploring" - default: - out["engine_mode"] = "idle" + engineMode = "stalled" + case running && activeCalls > 0: + engineMode = "exploring" } + out["engine_mode"] = engineMode out["active_task"] = map[string]any{ "id": t.ID, "description": t.Description, "goal": t.Goal, "running": running, "paused": paused, + "engine_mode": engineMode, "in_flight": inFlight, + "llm_in_flight": activeCalls, "last_activity": last, "stalled": stalled, "goals_total": len(goals), @@ -882,18 +918,7 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { goalCounts, _ := s.m.PG().GoalCountsAll() // goal progress per exploration, one query dtos := make([]TaskDTO, 0, len(list)) for _, t := range list { - status := "created" - switch { - case isTerminalStatus(t.Status): // 持久化终态优先(done/failed/timeout) - status = t.Status - case t.Queued: // 因并发上限挂起、等待空位(尚未启动) - status = "queued" - case t.Paused || s.engine.IsPaused(t.ID): - status = "paused" - case s.engine.Ready() && s.engine.Started(t.ID): - status = "running" - } - dto := taskDTO(t, status) + dto := taskDTO(t, s.resolvedTaskStatus(t)) dto.Tokens = tokenTotalDTO(toks[t.ExpID]) // prefer the live in-memory heartbeat (fresher) and fall back to the // persisted max activity time (survives restarts) for run-duration display. @@ -909,6 +934,28 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"tasks": dtos, "active": active}) } +func (s *Server) resolvedTaskStatus(t *Task) string { + if t == nil { + return "created" + } + switch { + case isTerminalStatus(t.Status): + return t.Status + case t.Queued: + return "queued" + case t.Paused || s.engine.IsPaused(t.ID): + return "paused" + case t.llmStateSnapshot().FailoverState == "chain_exhausted" && s.engine.Started(t.ID): + // LLM readiness is an execution dependency, not lifecycle state. Keeping + // an exhausted task running lets the user edit/reset its chain directly. + return "running" + case s.engine.ReadyFor(t) && s.engine.Started(t.ID): + return "running" + default: + return "created" + } +} + func (s *Server) setActive(w http.ResponseWriter, r *http.Request) { var req struct { ID string `json:"id"` @@ -946,6 +993,9 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { switch req.Action { case "pause": s.engine.Pause(t.ID, agent.AbortPausedByUser) + // Main Agent turns run outside Engine's planner/worker context so they need + // their own cancellation. New turns are rejected while the task is paused. + s.cancelTaskChat(t.ID, agent.AbortChatPausedWithTask) case "resume": s.engine.Run(s.ctx, t) // ensure loops are alive, then un-pause + nudge s.engine.Resume(t) @@ -962,6 +1012,99 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"id": t.ID, "paused": paused}) } +// controlIntent pauses, resumes, or cancels one local worker intent. A running +// worker is always stopped before state mutation/cleanup, preventing tool output +// that arrives after the user action from recreating deleted blackboard records. +func (s *Server) controlIntent(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + if !s.engine.beginTaskOperation(t.ID) { + writeErr(w, 409, "任务正在删除,无法控制意图") + return + } + defer s.engine.decInflight(t.ID) + + iid, err := strconv.ParseInt(r.PathValue("iid"), 10, 64) + if err != nil || iid <= 0 { + writeErr(w, 400, "bad intent id") + return + } + var req struct { + Action string `json:"action"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeErr(w, 400, "bad json: "+err.Error()) + return + } + if req.Action != "pause" && req.Action != "resume" && req.Action != "cancel" { + writeErr(w, 400, "action must be pause|resume|cancel") + return + } + + node, err := t.Store.GetNode(iid) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if node == nil { + if inherited, sourceErr := t.Store.GetNodeWithSources(iid); sourceErr == nil && inherited != nil && inherited.Inherited { + writeErr(w, 409, "继承意图为只读,不能控制") + } else { + writeErr(w, 404, "intent not found") + } + return + } + if node.Kind != db.KindIntent { + writeErr(w, 409, "node is not an intent") + return + } + + switch req.Action { + case "pause": + if node.State != "running" { + writeErr(w, 409, "仅运行中的意图可以暂停") + return + } + if err := s.engine.ControlWork(iid, "pause"); err != nil { + writeErr(w, 409, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"id": iid, "state": "paused"}) + case "resume": + if node.State != "paused" { + writeErr(w, 409, "仅已暂停的意图可以恢复") + return + } + if err := t.Store.SetIntentState(iid, "open"); err != nil { + writeErr(w, 500, err.Error()) + return + } + t.Notify() + writeJSON(w, 200, map[string]any{"id": iid, "state": "open"}) + case "cancel": + if node.State != "running" && node.State != "paused" { + writeErr(w, 409, "仅运行中或已暂停的意图可以取消") + return + } + if node.State == "running" { + if err := s.engine.ControlWork(iid, "cancel"); err != nil { + writeErr(w, 409, err.Error()) + return + } + } + deleted, err := t.Store.CancelIntent(iid) + if err != nil { + writeErr(w, 409, err.Error()) + return + } + t.Notify() + writeJSON(w, 200, map[string]any{"id": iid, "state": "cancelled", "deleted": deleted}) + } +} + // rerunIntent 重跑一条没跑成功的意图(blocked/exhausted/stopped):把它置回 open,worker // 会重新认领、从头再跑(已写回图谱的 fact/finding/asset 保留);若任务已终态/暂停则顺带复活。 // 用于「出错的 work 点击继续运行」——网络/LLM 抖动导致 blocked 后可一键重试。 @@ -1063,11 +1206,24 @@ func (s *Server) setLLM(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "api_key required") return } + // Validate provider construction before persisting it as the active profile. + if _, err := cfg.NewProvider(); err != nil { + writeErr(w, 400, "provider init failed: "+err.Error()) + return + } + if err := s.saveLLMConfig(cfg); err != nil { + writeErr(w, 500, "persist provider failed: "+err.Error()) + return + } + s.invalidateProfileAgents() + if _, ok := s.loadLLMConfig(); !ok { + writeErr(w, 500, "saved provider is unavailable") + return + } if err := s.applyLLM(cfg); err != nil { writeErr(w, 400, "provider init failed: "+err.Error()) return } - s.saveLLMConfig(cfg) // persist to DB log.Printf("[engine] LLM configured via UI: %s / %s", cfg.Provider(), cfg.Model) s.getLLM(w, r) } @@ -1116,12 +1272,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { } type createTaskReq struct { - Description string `json:"description"` - Goal string `json:"goal"` - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 - TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时 - 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 轮)。 + Description string `json:"description"` + Goal string `json:"goal"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效 + SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务 + TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时 + 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 轮)。 } func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { @@ -1133,18 +1291,39 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { if strings.TrimSpace(req.Description) == "" { req.Description = "未命名任务" } - // validate the pinned profile up front (must exist + be usable), so a bad id fails - // task creation instead of silently falling back to the active profile at run time. - if req.LLMProfileID != nil { - if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok { - writeErr(w, 400, "所选 LLM 配置不存在或未设置 API Key") - return - } + if len(req.LLMProfileIDs) == 0 && req.LLMProfileID != nil { + req.LLMProfileIDs = []int64{*req.LLMProfileID} + } + if err := s.validateTaskProfileIDs(req.LLMProfileIDs); err != nil { + writeErr(w, 400, err.Error()) + return } if req.TimeoutSeconds < 0 { req.TimeoutSeconds = 0 } - t, err := s.m.CreateTask(req.Description, req.Goal, req.LLMProfileID, req.TimeoutSeconds, req.PlanHeartbeatSeconds) + if len(req.SourceTaskIDs) > db.MaxTaskSourceCount { + writeErr(w, 400, fmt.Sprintf("关联任务最多选择 %d 个", db.MaxTaskSourceCount)) + return + } + sourceIDs := make([]int64, 0, len(req.SourceTaskIDs)) + seenSources := map[int64]bool{} + for _, raw := range req.SourceTaskIDs { + id, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64) + if err != nil || id <= 0 || seenSources[id] { + writeErr(w, 400, "关联任务 id 无效或重复") + return + } + if _, ok := s.m.Task(strconv.FormatInt(id, 10)); !ok { + writeErr(w, 400, fmt.Sprintf("关联任务 #%d 不存在", id)) + return + } + seenSources[id] = true + sourceIDs = append(sourceIDs, id) + } + t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{ + SourceTaskIDs: sourceIDs, LLMProfileIDs: req.LLMProfileIDs, + TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds, + }) if err != nil { writeErr(w, 500, err.Error()) return @@ -1153,7 +1332,73 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { // 共享的建后流程(seed + 种子意图 + 后台目标分解 + engine.Run),与 spawn_task 复用同一段。 // launchTask 内部异步,不阻塞 UI —— 目标分解在后台可见地进行。 s.launchTask(t, req.Description+" "+req.Goal, req.SeedFirstIntent != nil && *req.SeedFirstIntent) - writeJSON(w, 201, t) + writeJSON(w, 201, taskDTO(t, s.resolvedTaskStatus(t))) +} + +func (s *Server) validateTaskProfileIDs(ids []int64) error { + seen := map[int64]bool{} + for _, id := range ids { + if id <= 0 || seen[id] { + return fmt.Errorf("LLM 配置 id 无效或重复") + } + seen[id] = true + if _, ok := s.loadProfileConfig(id); !ok { + return fmt.Errorf("LLM 配置 #%d 不存在或未设置 API Key", id) + } + } + return nil +} + +func (s *Server) updateTaskLLMProfiles(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + if isTerminalStatus(t.Status) { + writeErr(w, 409, "已结束的任务不能修改 LLM 配置") + return + } + before := t.llmStateSnapshot() + var req struct { + LLMProfileIDs []int64 `json:"llm_profile_ids"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeErr(w, 400, "bad json: "+err.Error()) + return + } + if err := s.validateTaskProfileIDs(req.LLMProfileIDs); err != nil { + writeErr(w, 400, err.Error()) + return + } + active := int64(0) + if req.ActiveLLMProfileID != nil { + active = *req.ActiveLLMProfileID + } + reopened, err := s.m.ReplaceTaskLLMProfiles(t.ID, req.LLMProfileIDs, active) + if err != nil { + if strings.Contains(err.Error(), "terminal") { + writeErr(w, 409, err.Error()) + } else { + writeErr(w, 400, err.Error()) + } + return + } + // Profile edits affect only subsequent LLM calls. Existing in-flight calls + // retain their provider; a nudge lets an idle running task resume promptly. + t.Notify() + llmState := t.llmStateSnapshot() + result := map[string]any{ + "id": t.ID, "llm_profile_ids": llmState.ProfileIDs, + "active_llm_profile_id": llmState.ActiveID, + "llm_failover_state": llmState.FailoverState, "reopened_intents": reopened, + } + if !sameOptionalID(before.ActiveID, llmState.ActiveID) { + event := s.emitManualTaskLLMSwitch(t, before.ActiveID, llmState.ActiveID) + result["switch_event"] = activityDTO(event) + } + writeJSON(w, 200, result) } var ( @@ -1241,12 +1486,13 @@ func (s *Server) seed(t *Task, text string) { } var rootID int64 if as := s.m.Assets(); as != nil { + taskID, _ := strconv.ParseInt(t.ID, 10, 64) if net.ParseIP(host) != nil { - rootID, _ = as.UpsertIP(db.UpsertIPReq{IP: host}) + rootID, _ = as.UpsertIP(db.UpsertIPReq{IP: host, TaskID: taskID}) } else if scheme == "https" || scheme == "http" { - rootID, _ = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: u}) + rootID, _ = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: u, TaskID: taskID}) } else { - rootID, _ = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host}) + rootID, _ = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: taskID}) } } // anchor the seeded assets to this task's begin root as lineage/provenance @@ -1286,7 +1532,7 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "task not found") return } - writeJSON(w, 200, t) + writeJSON(w, 200, taskDTO(t, s.resolvedTaskStatus(t))) } // taskCoverage returns a task's rough asset test coverage (denominator/tested/backlog). @@ -1302,7 +1548,7 @@ func (s *Server) taskCoverage(w http.ResponseWriter, r *http.Request) { return } taskID, _ := strconv.ParseInt(t.ID, 10, 64) - cov, err := as.TaskCoverage(taskID, t.ExpID) + cov, err := as.TaskCoverageWithSources(taskID) if err != nil { writeErr(w, 500, err.Error()) return @@ -1345,22 +1591,23 @@ func (s *Server) taskAssetRefs(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "需要 asset_id") return } - refs, err := t.Store.AssetRefs(assetID) + refs, err := t.Store.AssetRefsWithSources(assetID) if err != nil { writeErr(w, 500, err.Error()) return } - intents := []db.AssetRef{} - facts := []db.AssetRef{} - findings := []db.AssetRef{} + intents := []CoverageAssetRefDTO{} + facts := []CoverageAssetRefDTO{} + findings := []CoverageAssetRefDTO{} for _, ref := range refs { + dto := coverageAssetRefDTO(ref) switch ref.Kind { case "intent": - intents = append(intents, ref) + intents = append(intents, dto) case "fact": - facts = append(facts, ref) + facts = append(facts, dto) case "finding": - findings = append(findings, ref) + findings = append(findings, dto) } } writeJSON(w, 200, map[string]any{"intents": intents, "facts": facts, "findings": findings}) @@ -1379,7 +1626,7 @@ func (s *Server) taskScopeList(w http.ResponseWriter, r *http.Request) { return } taskID, _ := strconv.ParseInt(t.ID, 10, 64) - rows, err := as.ListTaskScope(taskID) + rows, err := as.ListTaskScopeWithSources(taskID) if err != nil { writeErr(w, 500, err.Error()) return @@ -1443,7 +1690,11 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, []any{}) return } - f, _ := t.Store.ListByKind(db.KindFinding, 200) + f, err := t.Store.ListByKind(db.KindFinding, 200) + if err != nil { + writeErr(w, 500, err.Error()) + return + } tid, _ := strconv.ParseInt(t.ID, 10, 64) meta, err := s.m.pg.FindingMetaByNodeID(tid) if err != nil { @@ -1454,7 +1705,42 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) { aidSet = append(aidSet, m.AssetIDs...) } assets := s.resolveAssetIDs(aidSet) - writeJSON(w, 200, findingDTOsForTask(t, f, meta, assets)) + out := findingDTOsForTask(t, f, meta, assets) + + // Related-task findings are a live, read-only view. Provenance metadata lets + // the task UI suppress mutation affordances while retaining stable ids. + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + nodes, listErr := source.Store.ListByKind(db.KindFinding, 200) + if listErr != nil { + writeErr(w, 500, listErr.Error()) + return + } + for _, node := range nodes { + node.SourceTaskID = source.Task.TaskID + node.Inherited = true + } + sourceMeta, metaErr := s.m.pg.FindingMetaByNodeID(source.Task.TaskID) + if metaErr != nil { + log.Printf("[findings] source_task=%d meta: %v", source.Task.TaskID, metaErr) + } + var sourceAssetIDs []int64 + for _, item := range sourceMeta { + sourceAssetIDs = append(sourceAssetIDs, item.AssetIDs...) + } + out = append(out, findingDTOsForOwner( + i64s(source.Task.TaskID), + source.Task.Description, + nodes, + sourceMeta, + s.resolveAssetIDs(sourceAssetIDs), + )...) + } + writeJSON(w, 200, out) } // normFilter maps the frontend's "all" sentinel (and empty) to "" so the DB layer @@ -1529,7 +1815,41 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) { return } assets := s.resolveAssetIDs(f.AssetIDs) - writeJSON(w, 200, findingFromDB(f, assets)) + dto := findingFromDB(f, assets) + if contextTaskID := strings.TrimSpace(r.URL.Query().Get("context_task")); contextTaskID != "" { + contextTask := s.m.ResolveTask(contextTaskID) + if contextTask == nil { + writeErr(w, 404, "context task not found") + return + } + sourceTaskID, inherited, allowed := findingProvenanceInTask(contextTask, f.TaskID) + if !allowed { + writeErr(w, 404, "finding not available in task context") + return + } + dto.SourceTaskID = sourceTaskID + dto.Inherited = inherited + } + writeJSON(w, 200, dto) +} + +func findingProvenanceInTask(contextTask *Task, findingTaskID *int64) (sourceTaskID string, inherited, allowed bool) { + if contextTask == nil || findingTaskID == nil { + return "", false, false + } + contextID, err := strconv.ParseInt(contextTask.ID, 10, 64) + if err != nil { + return "", false, false + } + if *findingTaskID == contextID { + return "", false, true + } + for _, sourceID := range contextTask.SourceTaskIDs { + if sourceID == *findingTaskID { + return i64s(sourceID), true, true + } + } + return "", false, false } // findingLineage returns the exploration sub-DAG from the task root to this @@ -1699,17 +2019,108 @@ func (s *Server) intents(w http.ResponseWriter, r *http.Request) { return } // Paged form: ?before= (or ?page as a marker) → {items, has_more} so the - // worker session list can reach past the old fixed 300 boundary on scroll. + // worker session list can reach past the old fixed 300 boundary on scroll. It + // also includes immutable intent results from directly related tasks; those + // rows never participate in this task's frontier or claim path. before := int64(atoiDefault(q.Get("before"), 0)) - in, hasMore, err := t.Store.ListByKindPage(db.KindIntent, before, limit) + in, hasMore, err := taskIntentHistoryPage(t.Store, before, limit, false, 0) if err != nil { log.Printf("[intents] task=%s before=%d limit=%d: %v", t.ID, before, limit, err) writeErr(w, 500, err.Error()) return } + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + items, sourceMore, sourceErr := taskIntentHistoryPage(source.Store, before, limit, true, source.Task.TaskID) + if sourceErr != nil { + writeErr(w, 500, sourceErr.Error()) + return + } + in = append(in, items...) + hasMore = hasMore || sourceMore + } + sort.Slice(in, func(i, j int) bool { return in[i].ID > in[j].ID }) + if len(in) > limit { + in = in[:limit] + hasMore = true + } writeJSON(w, 200, map[string]any{"items": taskNodeDTOs(in), "has_more": hasMore}) } +func inheritedIntentResult(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false + } +} + +// inheritedGraphSnapshot exposes immutable source-task history without leaking +// work that is still open or running. Edges are filtered with the nodes so the +// response cannot contain dangling ids that reveal hidden live intent topology. +func inheritedGraphSnapshot(nodes []*db.Node, edges []db.Edge, sourceTaskID int64) ([]*db.Node, []db.Edge) { + visible := make(map[int64]struct{}, len(nodes)) + filteredNodes := make([]*db.Node, 0, len(nodes)) + for _, node := range nodes { + if node == nil || (node.Kind == db.KindIntent && !inheritedIntentResult(node.State)) { + continue + } + node.SourceTaskID = sourceTaskID + node.Inherited = true + visible[node.ID] = struct{}{} + filteredNodes = append(filteredNodes, node) + } + filteredEdges := make([]db.Edge, 0, len(edges)) + for _, edge := range edges { + if _, ok := visible[edge.From]; !ok { + continue + } + if _, ok := visible[edge.To]; !ok { + continue + } + filteredEdges = append(filteredEdges, edge) + } + return filteredNodes, filteredEdges +} + +// taskIntentHistoryPage returns one newest-first page for an exploration. +// A source may have many open/running intents, so keep paging until enough +// historical results are collected instead of leaking its frontier into the UI. +func taskIntentHistoryPage(store *db.ExplorationStore, before int64, limit int, inherited bool, sourceTaskID int64) ([]*db.Node, bool, error) { + if !inherited { + return store.ListByKindPage(db.KindIntent, before, limit) + } + batch := max(limit, 300) + cursor := before + out := make([]*db.Node, 0, limit+1) + for { + page, more, err := store.ListByKindPage(db.KindIntent, cursor, batch) + if err != nil { + return nil, false, err + } + for _, node := range page { + if !inheritedIntentResult(node.State) { + continue + } + node.SourceTaskID = sourceTaskID + node.Inherited = true + out = append(out, node) + if len(out) > limit { + return out[:limit], true, nil + } + } + if !more || len(page) == 0 { + return out, false, nil + } + cursor = page[len(page)-1].ID + } +} + // explorationGraph returns the whole exploration chain (task graph) as nodes+edges. func (s *Server) explorationGraph(w http.ResponseWriter, r *http.Request) { t := s.m.ResolveTask(r.URL.Query().Get("task")) @@ -1717,8 +2128,36 @@ func (s *Server) explorationGraph(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"nodes": []any{}, "edges": []any{}}) return } - nodes, _ := t.Store.Nodes(2000) - edges, _ := t.Store.Edges(5000) + nodes, err := t.Store.Nodes(2000) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + edges, err := t.Store.Edges(5000) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + sourceNodes, nodeErr := source.Store.Nodes(2000) + if nodeErr != nil { + writeErr(w, 500, nodeErr.Error()) + return + } + sourceEdges, edgeErr := source.Store.Edges(5000) + if edgeErr != nil { + writeErr(w, 500, edgeErr.Error()) + return + } + sourceNodes, sourceEdges = inheritedGraphSnapshot(sourceNodes, sourceEdges, source.Task.TaskID) + nodes = append(nodes, sourceNodes...) + edges = append(edges, sourceEdges...) + } writeJSON(w, 200, map[string]any{"nodes": taskNodeDTOs(nodes), "edges": edgeDTOs(edges)}) } @@ -1737,7 +2176,16 @@ func (s *Server) activity(w http.ResponseWriter, r *http.Request) { intentPtr = &n } } - items, cursor, err := t.Store.ActivityList(intentPtr, since, limit) + var ( + items []db.Activity + cursor int64 + err error + ) + if intentPtr != nil { + items, cursor, err = t.Store.ActivityListWithSources(*intentPtr, since, limit) + } else { + items, cursor, err = t.Store.ActivityList(nil, since, limit) + } if err != nil { log.Printf("[activity] task=%s since=%d limit=%d intent=%v: %v", t.ID, since, limit, intentPtr, err) writeErr(w, 500, err.Error()) @@ -1765,6 +2213,38 @@ func parseActivitySession(sess string) (db.ActivitySessionFilter, bool) { return db.ActivitySessionFilter{}, false } +// activitySessionStore resolves a worker session against the current task and +// its direct sources. Planner/main always stay local; inherited intent sessions +// are immutable history and use the source exploration only for reads. +func activitySessionStore(t *Task, filter db.ActivitySessionFilter) (*db.ExplorationStore, int64, error) { + if filter.NodeID == nil { + return t.Store, 0, nil + } + node, err := t.Store.GetNodeWithSources(*filter.NodeID) + if err != nil { + return nil, 0, err + } + if node == nil || node.Kind != db.KindIntent { + return nil, 0, nil + } + if !node.Inherited { + return t.Store, 0, nil + } + if !inheritedIntentResult(node.State) { + return nil, 0, nil + } + sources, err := t.Store.DirectSourceStores() + if err != nil { + return nil, 0, err + } + for _, source := range sources { + if source.Task.TaskID == node.SourceTaskID { + return source.Store, source.Task.TaskID, nil + } + } + return nil, 0, nil +} + // activityHistory serves one reverse-paginated page of a session's activity history. // The latest page (no ?before) opens a session; ?before= pulls the older page on // scroll-up. snapshot_cursor is the TASK-level max id at query time — the client uses @@ -1785,18 +2265,43 @@ func (s *Server) activityHistory(w http.ResponseWriter, r *http.Request) { } before := int64(atoiDefault(q.Get("before"), 0)) limit := min(atoiDefault(q.Get("limit"), 200), 500) // cap so one request can't pull an unbounded slice + store, sourceTaskID, err := activitySessionStore(t, filter) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if store == nil { + writeErr(w, 404, "session not found") + return + } + // The browser tails only this task's broadcaster. Keep the SSE join cursor + // local even when the opened worker transcript comes from a related task. snapshot, err := t.Store.ActivityMaxID() if err != nil { log.Printf("[activity/history] task=%s session=%s snapshot: %v", t.ID, sess, err) writeErr(w, 500, err.Error()) return } - items, hasMore, err := t.Store.ActivityPage(filter, before, limit) + var items []db.Activity + var hasMore bool + if sourceTaskID > 0 { + // Source sessions are readable only while the intent remains terminal. The + // DB query also removes model reasoning/accounting rows from inherited data. + items, hasMore, err = store.ActivityPageForTerminalIntent(*filter.NodeID, before, limit) + } else { + items, hasMore, err = store.ActivityPage(filter, before, limit) + } if err != nil { log.Printf("[activity/history] task=%s session=%s before=%d limit=%d: %v", t.ID, sess, before, limit, err) writeErr(w, 500, err.Error()) return } + if sourceTaskID > 0 { + for i := range items { + items[i].SourceTaskID = sourceTaskID + items[i].Inherited = true + } + } earliest := before if len(items) > 0 { earliest = items[0].ID @@ -2072,7 +2577,7 @@ func (s *Server) activityDetail(w http.ResponseWriter, r *http.Request) { return } seq, _ := strconv.ParseInt(r.PathValue("seq"), 10, 64) - d, err := t.Store.ActivityDetail(seq) + d, err := taskActivityDetail(t, seq) if err != nil { writeErr(w, 500, err.Error()) return @@ -2080,6 +2585,14 @@ func (s *Server) activityDetail(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"detail": d}) } +// taskActivityDetail keeps the legacy local-task behavior (including thinking +// rows), while inherited details are restricted to terminal worker intents. +// This prevents a guessed global activity id from exposing source planner/main +// or in-flight transcripts. +func taskActivityDetail(t *Task, seq int64) (string, error) { + return t.Store.ActivityDetailWithSources(seq) +} + func (s *Server) getTraffic(w http.ResponseWriter, r *http.Request) { tr := s.m.Traffic() if tr == nil { @@ -2423,6 +2936,14 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "no active task") return } + if s.engine.IsDeleting(t.ID) { + writeErr(w, 409, "任务正在删除,无法发送新消息") + return + } + if s.engine.IsPaused(t.ID) { + writeErr(w, 409, "任务已暂停,无法发送新消息") + return + } var req struct { Message string `json:"message"` Attachments []chatAttachment `json:"attachments,omitempty"` // 方式1 上传的文件(路径相对任务工作目录) @@ -2431,29 +2952,36 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, err.Error()) return } + // Admission is serialized with deletion's chat cancellation. Mark every chat + // turn busy before its first activity/file write, including rule-mode turns. + // If deletion wins the race, the second barrier check rejects this request. + s.chatMu.Lock() + if s.engine.IsDeleting(t.ID) || s.engine.IsPaused(t.ID) { + s.chatMu.Unlock() + writeErr(w, 409, "任务正在删除或已暂停,无法发送新消息") + return + } + if s.chatBusy[t.ID] { + s.chatMu.Unlock() + writeErr(w, 409, "主 Agent 正在处理上一条消息,请稍候") + return + } + ctx, cancel := context.WithCancelCause(s.ctx) + s.chatBusy[t.ID] = true + s.chatCancel[t.ID] = cancel + s.chatMu.Unlock() + // Persist + broadcast the human turn so the 主 Agent 编排会话 survives page // reloads and updates live: the conversation lives in the activity stream as // worker="mainagent" (the per-task activity table, replayed via SSE). With // attachments, the activity's Detail carries {text, attachments} so the transcript // renders attachment cards. s.engine.emitActivity(t, userActivityWithAttachments("mainagent", req.Message, req.Attachments)) - if ma := s.mainAgentRef(); ma != nil { - // Serialize per task: one main-agent run at a time so concurrent messages - // don't corrupt the shared exp-main transcript. If a prior turn is still - // running, tell the caller to wait rather than starting a racing run. - s.chatMu.Lock() - if s.chatBusy[t.ID] { - s.chatMu.Unlock() - writeErr(w, 409, "主 Agent 正在处理上一条消息,请稍候") - return - } - // Cancellable context so stopChat can abort this turn without affecting - // the server's root context (mirrors the conversation stop pattern). - ctx, cancel := context.WithCancelCause(s.ctx) - s.chatBusy[t.ID] = true - s.chatCancel[t.ID] = cancel - s.chatMu.Unlock() - + var ma *agent.MainAgent + if s.taskRuntimeAvailable(t, "mainagent") { + ma = s.agentsForTask(t).main + } + if ma != nil { // Run the agent on a per-turn cancellable ctx (NOT r.Context()): the turn can // take minutes (multi-tool loop), and binding it to the request lifecycle meant // a page reload / proxy timeout cancelled it mid-run ("context canceled"). The @@ -2461,11 +2989,7 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // handler returns immediately and the browser never needs to hold the request. go func() { defer func() { - cancel(agent.AbortChatTurnFinished) - s.chatMu.Lock() - delete(s.chatBusy, t.ID) - delete(s.chatCancel, t.ID) - s.chatMu.Unlock() + s.finishTaskChat(t.ID, cancel) }() // emit every step (thinking/tool_use/tool_result/text/result) so the main // agent session shows its work live, like worker/planner. The final answer @@ -2477,7 +3001,10 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // taskDir = agent 的工作目录(CWD),与 chatUpload 落盘、ensureRunDir 一致。 taskDir := filepath.Join(s.m.dir, "tasks", t.ID) agentMsg := composeAgentMessage(req.Message, req.Attachments, taskDir) - if _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, agentMsg, emit, t.Notify, resume, t.NotifyGoal); err != nil && ctx.Err() == nil { + s.engine.BeginLLMCall(t.ID) + _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, agentMsg, emit, t.Notify, resume, t.NotifyGoal) + s.engine.EndLLMCall(t.ID) + if err != nil && ctx.Err() == nil { s.engine.emitActivity(t, db.Activity{Worker: "mainagent", Kind: "text", IsError: true, Summary: "(主 Agent 出错:" + err.Error() + ")"}) } }() @@ -2486,9 +3013,28 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { } reply := s.fallbackChat(t, req.Message) s.engine.emitActivity(t, db.Activity{Worker: "mainagent", Kind: "text", Summary: reply}) + s.finishTaskChat(t.ID, cancel) writeJSON(w, 200, map[string]any{"reply": reply, "mode": "rule"}) } +func (s *Server) cancelTaskChat(taskID string, cause error) bool { + s.chatMu.Lock() + cancel := s.chatCancel[taskID] + if cancel != nil { + cancel(cause) + } + s.chatMu.Unlock() + return cancel != nil +} + +func (s *Server) finishTaskChat(taskID string, cancel context.CancelCauseFunc) { + cancel(agent.AbortChatTurnFinished) + s.chatMu.Lock() + delete(s.chatBusy, taskID) + delete(s.chatCancel, taskID) + s.chatMu.Unlock() +} + // stopChat aborts the in-flight main-agent turn for a task (manual stop button). // Mirrors pgStopConversation: cancels only the current turn; planner/workers are // unaffected and continue running. The user can send a new message immediately. @@ -2498,14 +3044,10 @@ func (s *Server) stopChat(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "task not found") return } - s.chatMu.Lock() - cancel := s.chatCancel[t.ID] - s.chatMu.Unlock() - if cancel == nil { + if !s.cancelTaskChat(t.ID, agent.AbortChatStoppedByUser) { writeJSON(w, 200, map[string]any{"status": "idle"}) return } - cancel(agent.AbortChatStoppedByUser) writeJSON(w, 200, map[string]any{"status": "stopping"}) } diff --git a/server/server_mgmt.go b/server/server_mgmt.go index eef78bb..9732b9d 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -5,6 +5,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "log" @@ -66,61 +67,97 @@ func decode(r *http.Request, v any) error { return json.NewDecoder(r.Body).Decod // ---------- tasks (delete) ---------- -// pgDeleteTask removes a task and, per opt-in flags in the (optional) JSON body, -// its associated data: assets (task-owned), findings, recorded traffic (by the -// task's hosts) and the per-task working directory. Empty body → all flags false → -// only the task + its exploration subgraph are removed (legacy behavior). +const taskDeleteDrainTimeout = 10 * time.Second + +func canonicalTaskID(raw string) (string, bool) { + n, err := strconv.ParseInt(raw, 10, 64) + if err != nil || n <= 0 { + return "", false + } + return strconv.FormatInt(n, 10), true +} + func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) { - id := r.PathValue("id") - var opts struct { - Assets bool `json:"assets"` - Findings bool `json:"findings"` - Traffic bool `json:"traffic"` - Files bool `json:"files"` + id, ok := canonicalTaskID(r.PathValue("id")) + if !ok { + writeErr(w, http.StatusBadRequest, "任务 id 无效") + return } - _ = json.NewDecoder(r.Body).Decode(&opts) // optional body; ignore EOF/parse errors - - result := map[string]any{"deleted": id} - nid, _ := strconv.ParseInt(id, 10, 64) - as := s.m.Assets() - - // Traffic first: it needs the task's hosts, which come from assets that a later - // asset-delete would remove. Traffic is host-keyed (no per-task tag), so this - // purges ALL recorded traffic for those hosts, not just this task's. - if opts.Traffic && nid > 0 && as != nil { - if tr := s.m.Traffic(); tr != nil { - if hosts, err := as.HostsByTask(nid); err == nil && len(hosts) > 0 { - if n, err := tr.DeleteHostsExact(hosts); err == nil { - result["traffic_deleted"] = n - } - } + var opts DeleteTaskOptions + if err := decode(r, &opts); err != nil && err != io.EOF { + writeErr(w, 400, "invalid JSON: "+err.Error()) + return + } + if !s.engine.BeginDelete(id) { + writeErr(w, http.StatusConflict, "任务正在删除") + return + } + deleted := false + defer func() { + if !deleted { + s.engine.AbortDelete(id) } - } - if opts.Findings && nid > 0 { - if n, err := s.m.pg.DeleteFindingsByTask(nid); err == nil { - result["findings_deleted"] = n - } - } - if opts.Assets && nid > 0 && as != nil { - if n, err := as.DeleteByTaskID(nid); err == nil { - result["assets_deleted"] = n - } - } + }() - s.engine.Pause(id, agent.AbortTaskDeleted) // cancel running loops before removing - if err := s.m.DeleteTask(id); err != nil { - writeErr(w, 500, err.Error()) + // Main-agent runs use a separate context from planner/workers. Cancel it, then + // wait until both execution domains have returned before removing transcripts. + s.cancelTaskChat(id, agent.AbortTaskDeleted) + drainCtx, cancelDrain := context.WithTimeout(r.Context(), taskDeleteDrainTimeout) + defer cancelDrain() + if err := s.waitTaskQuiescent(drainCtx, id); err != nil { + writeErr(w, http.StatusConflict, "任务仍有运行中的 Agent,删除已取消") return } - // Files last: the per-task working dir /tasks/ (worker artifacts + - // uploads). Confined to the work dir; id is numeric so no traversal risk. - if opts.Files && nid > 0 { - if err := os.RemoveAll(filepath.Join(s.m.dir, "tasks", id)); err == nil { - result["files_deleted"] = true + result, err := s.m.DeleteTask(id, opts) + if err != nil { + var committed *taskDeleteCommittedError + if errors.As(err, &committed) { + // PostgreSQL is already gone. Complete runtime teardown and return the + // auditable counts together with the post-commit cleanup warning. + s.engine.StopTask(id) + s.taskAgentMu.Lock() + delete(s.taskAgents, id) + s.taskAgentMu.Unlock() + deleted = true + writeCommittedTaskDelete(w, result, err) + return + } + writeErr(w, 500, err.Error()) + return + } + // Manager has removed the task from the registry, so no new API operation can + // resolve it. Now stop and join every task-owned Engine goroutine and clear its + // lifecycle maps before releasing the delete barrier. + s.engine.StopTask(id) + s.taskAgentMu.Lock() + delete(s.taskAgents, id) + s.taskAgentMu.Unlock() + deleted = true + writeJSON(w, 200, result) +} + +func writeCommittedTaskDelete(w http.ResponseWriter, result DeleteTaskResult, err error) { + result.CleanupWarning = err.Error() + writeJSON(w, http.StatusOK, result) +} + +func (s *Server) waitTaskQuiescent(ctx context.Context, taskID string) error { + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + s.chatMu.Lock() + chatBusy := s.chatBusy[taskID] + s.chatMu.Unlock() + if s.engine.inflightCount(taskID) == 0 && !chatBusy { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: } } - writeJSON(w, 200, result) } // ---------- agents ---------- @@ -352,6 +389,9 @@ func (s *Server) pgSaveAgentConfig(w http.ResponseWriter, r *http.Request) { if profileChanged { s.invalidateProfileAgents() } + // Task bundles capture operational Agent settings (turn/time budgets, tools, + // web search). Rebuild them even when no global profile is configured. + s.invalidateTaskAgents() // rebuild the live agents so the new max_turns/run_seconds/binding apply without a restart. s.cfgMu.Lock() cfg, on := s.llmCfg, s.llmOn @@ -396,8 +436,8 @@ func (s *Server) pgGetAgent(w http.ResponseWriter, r *http.Request) { } writeJSON(w, 200, map[string]any{ "agent": agentDTO(a), "prompt": cur, "variables": vars, "versions": vers, - "visibility": map[string]any{"mcp": mcp, "skill": sk}, - "llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选 + "visibility": map[string]any{"mcp": mcp, "skill": sk}, + "llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选 "wrapup_prompt": a.WrapupPrompt, // 已保存的收尾提示词(空=用内置默认) "wrapup_default": agent.WrapupDefault(a.Key), // 内置默认(供占位/恢复默认) @@ -1607,7 +1647,9 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) { writeErr(w, 500, err.Error()) return } - // Editing a profile rebuilds any task pinned to it on its next round. + // Editing a profile rebuilds any task pinned to it on its next round. Reapply + // the active profile too: the global fallback and explicit task chains must + // repopulate the same provider-cache entry and therefore share one limiter. s.invalidateProfileAgents() // Hot-apply. Not just when the ACTIVE profile changed: with failover on, any // profile's key/model/priority/pool_exclude edit reshapes the chain, so the @@ -1622,16 +1664,50 @@ func (s *Server) pgDeleteProfile(w http.ResponseWriter, r *http.Request) { return } id, _ := pathInt(r, "id") - if err := pg.DeleteProfile(id); err != nil { - writeErr(w, 500, err.Error()) + if err := pg.DeleteProfileContext(r.Context(), id); err != nil { + switch { + case errors.Is(err, db.ErrActiveLLMProfileDelete): + writeErr(w, 409, "当前激活的 LLM 配置不能删除,请先激活其他配置") + case errors.Is(err, db.ErrLLMProfileReferencesChanged): + writeErr(w, 409, "LLM 配置正在被任务或会话修改,请重试") + case errors.Is(err, context.DeadlineExceeded): + writeErr(w, 409, "等待 LLM 配置引用释放超时,请重试") + case errors.Is(err, db.ErrLLMProfileNotFound): + writeErr(w, 404, err.Error()) + default: + writeErr(w, 500, err.Error()) + } return } s.invalidateProfileAgents() // drop cached agents for the removed profile - s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away) - s.reapplyActiveProfile() // the deleted profile may have been in the failover chain + s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away) + s.restoreTasksAfterProfileDelete(pg) + // Cache invalidation also removed the active profile's shared provider entry. + // Reapply after task-state sync so the wake-up observes the post-delete chain. + s.reapplyActiveProfile() writeJSON(w, 200, map[string]any{"deleted": id}) } +func (s *Server) restoreTasksAfterProfileDelete(pg *db.DB) { + for _, task := range s.m.List() { + if taskID, err := strconv.ParseInt(task.ID, 10, 64); err == nil { + if pt, err := pg.GetTask(taskID); err == nil && pt != nil { + s.syncTaskLLMState(pt) + // Deleting the active entry may advance the task to a ready successor, + // or clear the explicit chain and restore its Agent/global fallback. + // Resume only intents that were blocked by the exhausted chain; a + // paused task remains paused and terminal tasks remain immutable. + if !isTerminalStatus(task.Status) && s.taskRuntimeAvailable(task, "planner", "worker") { + if _, reopenErr := task.Store.ReopenIntentsByBlockedReason(db.IntentBlockedLLMQuota); reopenErr != nil { + log.Printf("[llm-profile] task %s reopen quota-blocked intents after profile delete: %v", task.ID, reopenErr) + } + } + task.Notify() + } + } + } +} + func (s *Server) pgActivateProfile(w http.ResponseWriter, r *http.Request) { pg := s.pg(w) if pg == nil { diff --git a/server/task_delete_barrier_test.go b/server/task_delete_barrier_test.go new file mode 100644 index 0000000..4b95192 --- /dev/null +++ b/server/task_delete_barrier_test.go @@ -0,0 +1,220 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" +) + +func TestEngineDeleteBarrierRejectsNewTaskOperations(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + + if !e.beginTaskOperation(taskID) { + t.Fatal("first task operation should be admitted") + } + if !e.BeginDelete(taskID) { + t.Fatal("first delete should install the barrier") + } + if !e.IsDeleting(taskID) { + t.Fatal("delete barrier should be visible") + } + if !e.IsPaused(taskID) { + t.Fatal("delete should temporarily pause task execution") + } + if e.beginTaskOperation(taskID) { + t.Fatal("operation started after BeginDelete returned") + } + if got := e.inflightCount(taskID); got != 1 { + t.Fatalf("existing operation must remain drainable, inflight=%d", got) + } + e.decInflight(taskID) + + if e.BeginDelete(taskID) { + t.Fatal("duplicate delete must not acquire the barrier") + } + e.AbortDelete(taskID) + if e.IsPaused(taskID) { + t.Fatal("aborted delete should restore a previously running task") + } + if !e.beginTaskOperation(taskID) { + t.Fatal("aborting deletion should reopen operation admission") + } + e.decInflight(taskID) +} + +func TestAbortDeletePreservesExistingPause(t *testing.T) { + e := NewEngine(nil) + e.Pause("42", agent.AbortPausedByUser) + if !e.BeginDelete("42") { + t.Fatal("delete barrier should be installed") + } + e.AbortDelete("42") + if !e.IsPaused("42") { + t.Fatal("aborted delete must preserve a pre-existing user pause") + } +} + +func TestStopTaskCancelsRuntimeAndClearsLifecycleState(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + e.deleteMu.RLock() + rt := e.registerTaskRoutines(context.Background(), taskID, 1) + e.started.Store(taskID, true) + e.lastAct.Store(taskID, int64(1)) + e.paused.Store(taskID, true) + e.plannerRound.Store(taskID, 3) + e.coordStarted.Store(taskID, true) + e.deleteMu.RUnlock() + + done := make(chan struct{}) + runTaskRoutine(rt, func(ctx context.Context) { + <-ctx.Done() + close(done) + }) + e.StopTask(taskID) + + select { + case <-done: + default: + t.Fatal("task runtime was not cancelled before StopTask returned") + } + if e.Started(taskID) || e.IsPaused(taskID) || e.IsDeleting(taskID) { + t.Fatalf("task lifecycle maps were not cleared: started=%v paused=%v deleting=%v", + e.Started(taskID), e.IsPaused(taskID), e.IsDeleting(taskID)) + } + if _, ok := e.lastAct.Load(taskID); ok { + t.Fatal("last activity state was retained") + } + if _, ok := e.plannerRound.Load(taskID); ok { + t.Fatal("planner round state was retained") + } + if _, ok := e.coordStarted.Load(taskID); ok { + t.Fatal("deadline coordinator state was retained") + } + e.runtimeMu.Lock() + _, retained := e.runtimes[taskID] + e.runtimeMu.Unlock() + if retained { + t.Fatal("task runtime retained after StopTask") + } +} + +func TestCanonicalTaskID(t *testing.T) { + for raw, want := range map[string]string{"42": "42", "00042": "42", "+42": "42"} { + if got, ok := canonicalTaskID(raw); !ok || got != want { + t.Fatalf("canonicalTaskID(%q)=(%q,%v), want (%q,true)", raw, got, ok, want) + } + } + for _, raw := range []string{"", "0", "-1", "abc", "42x"} { + if got, ok := canonicalTaskID(raw); ok { + t.Fatalf("canonicalTaskID(%q)=(%q,true), want invalid", raw, got) + } + } +} + +func TestCancelChatThenWaitTaskQuiescent(t *testing.T) { + e := NewEngine(nil) + chatCtx, chatCancel := context.WithCancelCause(context.Background()) + s := &Server{ + engine: e, + chatBusy: map[string]bool{"7": true}, + chatCancel: map[string]context.CancelCauseFunc{"7": chatCancel}, + } + if !e.beginTaskOperation("7") { + t.Fatal("task operation should be admitted") + } + + cancelObserved := make(chan struct{}) + releaseChat := make(chan struct{}) + go func() { + <-chatCtx.Done() + close(cancelObserved) + <-releaseChat + s.finishTaskChat("7", chatCancel) + }() + if !s.cancelTaskChat("7", agent.AbortTaskDeleted) { + t.Fatal("expected active chat to be cancelled") + } + <-cancelObserved + + waitDone := make(chan error, 1) + waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second) + defer waitCancel() + go func() { waitDone <- s.waitTaskQuiescent(waitCtx, "7") }() + + select { + case err := <-waitDone: + t.Fatalf("wait returned while chat and engine operation were active: %v", err) + case <-time.After(20 * time.Millisecond): + } + close(releaseChat) + select { + case err := <-waitDone: + t.Fatalf("wait returned while engine operation was active: %v", err) + case <-time.After(20 * time.Millisecond): + } + e.decInflight("7") + if err := <-waitDone; err != nil { + t.Fatalf("wait for quiescence: %v", err) + } +} + +func TestWaitTaskQuiescentHonorsContext(t *testing.T) { + s := &Server{ + engine: NewEngine(nil), + chatBusy: map[string]bool{"9": true}, + chatCancel: map[string]context.CancelCauseFunc{}, + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := s.waitTaskQuiescent(ctx, "9"); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected deadline exceeded, got %v", err) + } +} + +func TestTaskChatUploadRejectsDeleteBarrierBeforeCreatingDirectory(t *testing.T) { + dataDir := t.TempDir() + m := &Manager{dir: dataDir, tasks: map[string]*Task{"7": {ID: "7"}}} + s := &Server{m: m, engine: NewEngine(m)} + if !s.engine.BeginDelete("7") { + t.Fatal("delete barrier should be installed") + } + + req := httptest.NewRequest(http.MethodPost, "/api/chat/upload?scope=task&id=7", nil) + rec := httptest.NewRecorder() + s.chatUpload(rec, req) + if rec.Code != http.StatusConflict { + t.Fatalf("expected 409 while deleting, got %d: %s", rec.Code, rec.Body.String()) + } + if _, err := os.Stat(filepath.Join(dataDir, "tasks", "7")); !os.IsNotExist(err) { + t.Fatalf("upload must not recreate task directory, stat err=%v", err) + } +} + +func TestCommittedTaskDeleteReturnsSuccessWithCleanupWarning(t *testing.T) { + rec := httptest.NewRecorder() + result := DeleteTaskResult{Deleted: "7", FilesDeleted: true} + err := &taskDeleteCommittedError{err: errors.New("remove staged task files: permission denied")} + + writeCommittedTaskDelete(rec, result, err) + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String()) + } + var got DeleteTaskResult + if decodeErr := json.NewDecoder(rec.Body).Decode(&got); decodeErr != nil { + t.Fatal(decodeErr) + } + if got.Deleted != "7" || !got.FilesDeleted || !strings.Contains(got.CleanupWarning, "cleanup incomplete") { + t.Fatalf("unexpected response: %+v", got) + } +} diff --git a/server/task_llm.go b/server/task_llm.go new file mode 100644 index 0000000..c1dbe88 --- /dev/null +++ b/server/task_llm.go @@ -0,0 +1,462 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "log" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/llmrec" + "github.com/Autumn-27/norma/llm" + "github.com/Autumn-27/norma/memory" + "github.com/Autumn-27/norma/transcript" +) + +type taskAgentBundle struct { + runtime *taskLLMRuntime // goal decomposition runtime + plannerRuntime *taskLLMRuntime + workerRuntime *taskLLMRuntime + mainRuntime *taskLLMRuntime + pl *agent.Planner + wk *agent.Worker + main *agent.MainAgent +} + +type llmAuditProfile struct { + ID int64 `json:"id"` + Name string `json:"name"` + Format string `json:"format"` + Model string `json:"model"` +} + +type llmTransitionAudit struct { + Mode string `json:"mode"` // automatic | manual | exhausted + Reason string `json:"reason"` + Previous *llmAuditProfile `json:"previous,omitempty"` + Next *llmAuditProfile `json:"next,omitempty"` +} + +type llmActivityMetadata struct { + LLMTransition llmTransitionAudit `json:"llm_transition"` +} + +type taskLLMRuntime struct { + s *Server + taskID string + agentKey string +} + +type taskLLMError struct { + taskID string + chainExhausted bool + cause error +} + +func (e *taskLLMError) Error() string { + if e.chainExhausted { + return fmt.Sprintf("task %s LLM profile chain exhausted: %v", e.taskID, e.cause) + } + return e.cause.Error() +} + +func (e *taskLLMError) Unwrap() error { return e.cause } + +func isTaskLLMRuntimeError(err error) bool { + var target *taskLLMError + return errors.As(err, &target) +} + +func isTaskLLMChainExhausted(err error) bool { + var target *taskLLMError + return errors.As(err, &target) && target.chainExhausted +} + +// isQuotaExhaustedError is intentionally strict. A generic 429, auth error, +// network failure, or 5xx does not rotate providers; the response must explicitly +// identify quota, credits, billing balance, or payment exhaustion. +func isQuotaExhaustedError(err error) bool { + return err != nil && agent.IsQuotaExhaustedMessage(err.Error()) +} + +type taskLLMSelection struct { + task *Task + profileID int64 + revision int64 + provider llm.Provider +} + +type taskLLMStreamHooks struct { + current func() (taskLLMSelection, error) + exhaust func(taskLLMSelection, error) (db.TaskLLMTransition, error) + transition func(taskLLMSelection, db.TaskLLMTransition, error) +} + +func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) { + taskNum, err := parseTaskID(r.taskID) + if err != nil { + return nil, 0, 0, nil, err + } + pt, err := r.s.m.pg.GetTask(taskNum) + if err != nil { + return nil, 0, 0, nil, err + } + if pt == nil { + return nil, 0, 0, nil, fmt.Errorf("task %s not found", r.taskID) + } + r.s.syncTaskLLMState(pt) + t, _ := r.s.m.Task(r.taskID) + if len(pt.LLMProfileIDs) == 0 { + prov, _, ok := r.s.fallbackProviderForAgent(r.agentKey) + if !ok { + return t, 0, pt.LLMChainRevision, nil, fmt.Errorf("task %s has no available fallback LLM provider", r.taskID) + } + return t, 0, pt.LLMChainRevision, prov, nil + } + if pt.ActiveLLMProfileID == nil { + return t, 0, pt.LLMChainRevision, nil, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")} + } + prov, _, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID) + if !ok { + return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, nil, fmt.Errorf("LLM profile #%d is missing or invalid", *pt.ActiveLLMProfileID) + } + return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, prov, nil +} + +func parseTaskID(id string) (int64, error) { + n, err := strconv.ParseInt(id, 10, 64) + if err != nil || n <= 0 { + return 0, fmt.Errorf("invalid task id %q", id) + } + return n, nil +} + +func (r *taskLLMRuntime) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + ctx = llmrec.WithTaskID(ctx, r.taskID) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + task, profileID, revision, provider, err := r.current() + return taskLLMSelection{task: task, profileID: profileID, revision: revision, provider: provider}, err + }, + exhaust: func(selection taskLLMSelection, cause error) (db.TaskLLMTransition, error) { + taskNum, _ := parseTaskID(r.taskID) + transition, err := r.s.m.pg.MarkTaskLLMProfileQuotaExhaustedAtRevision(taskNum, selection.profileID, selection.revision, cause.Error()) + if err != nil { + return transition, err + } + if pt, getErr := r.s.m.pg.GetTask(taskNum); getErr == nil && pt != nil { + r.s.syncTaskLLMState(pt) + } + return transition, nil + }, + transition: func(selection taskLLMSelection, transition db.TaskLLMTransition, cause error) { + r.s.emitTaskLLMTransition(selection.task, transition, cause) + }, + } + return streamTaskLLM(ctx, r.taskID, req, hooks) +} + +func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) iter.Seq2[llm.StreamEvent, error] { + return func(yield func(llm.StreamEvent, error) bool) { + for { + selection, err := hooks.current() + if err != nil { + yield(llm.StreamEvent{}, err) + return + } + committed := false + var pending []llm.StreamEvent + var streamErr error + for event, err := range selection.provider.Stream(ctx, req) { + if err != nil { + streamErr = err + break + } + if !committed && !streamEventCommitsOutput(event) { + pending = append(pending, event) + continue + } + if !committed { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + pending = nil + committed = true + } + if !yield(event, nil) { + return + } + } + if streamErr == nil { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + return + } + // profileID=0 means the explicit chain was cleared while this stable task + // bundle was still in use. Agent/global fallback errors follow the legacy + // behavior and never mutate task failover state. + if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + yield(llm.StreamEvent{}, streamErr) + return + } + transition, markErr := hooks.exhaust(selection, streamErr) + if markErr != nil { + cause := fmt.Errorf("mark profile quota exhausted after %v: %w", streamErr, markErr) + if committed { + // Output may already have driven tool execution. Report the persistence + // failure, but classify it as router-handled so the worker does not + // replay the entire intent and duplicate those side effects. + yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, cause: cause}) + } else { + yield(llm.StreamEvent{}, cause) + } + return + } + if transition.Advanced && !transition.Stale && hooks.transition != nil { + hooks.transition(selection, transition, streamErr) + } + if committed || (!transition.Stale && transition.NextProfileID == nil) { + yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: streamErr}) + return + } + // No event reached the caller, so replaying the same logical request on + // the next profile cannot duplicate model output or tool execution. + } + } +} + +func streamEventCommitsOutput(event llm.StreamEvent) bool { + switch event.Type { + case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON: + return event.Text != "" + case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop: + return true + default: + return false + } +} + +func (r *taskLLMRuntime) CompactionWindow() int { + taskNum, err := parseTaskID(r.taskID) + if err != nil { + return (agent.Config{}).CompactionWindow() + } + chain, err := r.s.m.pg.TaskLLMProfiles(taskNum) + if err != nil { + return (agent.Config{}).CompactionWindow() + } + if len(chain) == 0 { + if _, cfg, ok := r.s.fallbackProviderForAgent(r.agentKey); ok { + return cfg.CompactionWindow() + } + return (agent.Config{}).CompactionWindow() + } + minimum := 0 + for _, entry := range chain { + if cfg, ok := r.s.loadProfileConfig(entry.ProfileID); ok { + window := cfg.CompactionWindow() + if minimum == 0 || window < minimum { + minimum = window + } + } + } + if minimum == 0 { + return (agent.Config{}).CompactionWindow() + } + return minimum +} + +// fallbackProviderForAgent resolves the pre-existing precedence used by tasks +// without an explicit chain: the Agent's saved binding first, then the current +// global provider (which may come from the default profile or environment). +func (s *Server) fallbackProviderForAgent(agentKey string) (llm.Provider, agent.Config, bool) { + if id := s.effectiveProfileForAgent(agentKey, nil); id != nil { + if prov, cfg, ok := s.providerForProfile(*id); ok { + return s.poolForBinding(*id, prov, cfg), cfg, true + } + } + s.cfgMu.Lock() + defer s.cfgMu.Unlock() + if !s.llmOn || s.llmProv == nil { + return nil, agent.Config{}, false + } + return s.llmProv, s.llmCfg, true +} + +func (s *Server) taskRuntimeAvailable(t *Task, agentKeys ...string) bool { + if t == nil { + return false + } + state := t.llmStateSnapshot() + if len(state.ProfileIDs) > 0 { + if state.ActiveID == nil { + return false + } + _, _, ok := s.providerForProfile(*state.ActiveID) + return ok + } + s.cfgMu.Lock() + globallyReady := s.llmOn && s.llmProv != nil + s.cfgMu.Unlock() + if globallyReady { + return true + } + for _, key := range agentKeys { + if _, _, ok := s.fallbackProviderForAgent(key); !ok { + return false + } + } + return len(agentKeys) > 0 +} + +func (s *Server) invalidateTaskAgents() { + s.taskAgentMu.Lock() + s.taskAgents = map[string]*taskAgentBundle{} + s.taskAgentMu.Unlock() +} + +func (s *Server) agentsForTask(t *Task) *taskAgentBundle { + s.taskAgentMu.Lock() + defer s.taskAgentMu.Unlock() + if bundle := s.taskAgents[t.ID]; bundle != nil { + return bundle + } + goalRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "goals"} + plannerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "planner"} + workerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "worker"} + mainRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "mainagent"} + tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) + window := workerRuntime.CompactionWindow() + wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker")) + wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow) + wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) + wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory"))) + wk.SetWebSearch(s.webSearchFor("worker")) + pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner")) + pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow) + pl.SetKillWork(s.engine.KillWork) + pl.SetSteerWork(s.engine.SteerWork) + pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + pl.SetWebSearch(s.webSearchFor("planner")) + main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent")) + main.SetCompactionWindowResolver(mainRuntime.CompactionWindow) + main.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + main.SetWebSearch(s.webSearchFor("mainagent")) + bundle := &taskAgentBundle{ + runtime: goalRuntime, plannerRuntime: plannerRuntime, workerRuntime: workerRuntime, + mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, + } + s.taskAgents[t.ID] = bundle + return bundle +} + +func (s *Server) syncTaskLLMState(pt *db.Task) { + if pt == nil { + return + } + id := fmt.Sprintf("%d", pt.ID) + s.m.mu.Lock() + if task := s.m.tasks[id]; task != nil { + task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason) + } + s.m.mu.Unlock() +} + +func (s *Server) emitTaskLLMTransition(t *Task, transition db.TaskLLMTransition, cause error) { + if t == nil { + return + } + previous := s.llmAuditProfile(transition.PreviousProfileID) + var next *llmAuditProfile + if transition.NextProfileID != nil { + next = s.llmAuditProfile(*transition.NextProfileID) + } + mode := "automatic" + kind := "llm_switch" + summary := fmt.Sprintf("%s 额度不足", llmAuditProfileLabel(previous)) + if transition.NextProfileID != nil { + summary += fmt.Sprintf(",后续调用切换到 %s", llmAuditProfileLabel(next)) + } else { + mode = "exhausted" + kind = "llm_failover" + summary += ",配置链已耗尽" + } + metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{ + Mode: mode, Reason: cause.Error(), Previous: previous, Next: next, + }}) + s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: kind, IsError: transition.ChainExhausted, Summary: summary, Detail: cause.Error(), Metadata: metadata}) + log.Printf("[llm-failover] task %s: %s", t.ID, summary) +} + +func (s *Server) llmAuditProfile(id int64) *llmAuditProfile { + if id <= 0 || s.m == nil || s.m.pg == nil { + return nil + } + p, err := s.m.pg.ProfileByID(id) + if err != nil || p == nil { + return &llmAuditProfile{ID: id, Name: fmt.Sprintf("配置 #%d", id)} + } + return &llmAuditProfile{ID: p.ID, Name: p.Name, Format: p.Format, Model: p.Model} +} + +func llmAuditProfileLabel(profile *llmAuditProfile) string { + if profile == nil { + return "默认配置" + } + name := profile.Name + if name == "" { + name = fmt.Sprintf("配置 #%d", profile.ID) + } + detail := []string{} + if profile.Format != "" { + detail = append(detail, profile.Format) + } + if profile.Model != "" { + detail = append(detail, profile.Model) + } + if len(detail) == 0 { + return name + } + return fmt.Sprintf("%s(%s)", name, strings.Join(detail, " / ")) +} + +func sameOptionalID(a, b *int64) bool { + if a == nil || b == nil { + return a == nil && b == nil + } + return *a == *b +} + +func (s *Server) emitManualTaskLLMSwitch(t *Task, previousID, nextID *int64) db.Activity { + previous := (*llmAuditProfile)(nil) + next := (*llmAuditProfile)(nil) + if previousID != nil { + previous = s.llmAuditProfile(*previousID) + } + if nextID != nil { + next = s.llmAuditProfile(*nextID) + } + summary := fmt.Sprintf("用户手动将任务 LLM 从 %s 切换到 %s", llmAuditProfileLabel(previous), llmAuditProfileLabel(next)) + metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{ + Mode: "manual", Reason: "用户手动切换任务 LLM", Previous: previous, Next: next, + }}) + return s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "llm_switch", Summary: summary, Detail: summary, Metadata: metadata}) +} diff --git a/server/task_llm_test.go b/server/task_llm_test.go new file mode 100644 index 0000000..faf1cfd --- /dev/null +++ b/server/task_llm_test.go @@ -0,0 +1,507 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "sync" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/norma/harness" + "github.com/Autumn-27/norma/llm" +) + +type scriptedLLMProvider struct { + events []llm.StreamEvent + err error + calls int +} + +func (p *scriptedLLMProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + return func(yield func(llm.StreamEvent, error) bool) { + p.calls++ + for _, event := range p.events { + if !yield(event, nil) { + return + } + } + if p.err != nil { + yield(llm.StreamEvent{}, p.err) + } + } +} + +func collectTaskLLMStream(seq iter.Seq2[llm.StreamEvent, error]) ([]llm.StreamEvent, error) { + var events []llm.StreamEvent + var streamErr error + for event, err := range seq { + if err != nil { + streamErr = err + break + } + events = append(events, event) + } + return events, streamErr +} + +func TestIsQuotaExhaustedError(t *testing.T) { + t.Parallel() + positive := []string{ + `openai: code=insufficient_quota status=429`, + `provider error: quota_exceeded`, + `grpc code=RESOURCE_EXHAUSTED`, + `You exceeded your current quota, please check your plan and billing details.`, + `billing_hard_limit_reached`, + `billing_not_active`, + `HTTP status 402 Payment Required`, + `账户余额不足,请充值`, + `credit balance is too low`, + } + for _, message := range positive { + if !isQuotaExhaustedError(errors.New(message)) { + t.Errorf("expected quota classification for %q", message) + } + } + + negative := []string{ + `HTTP status 429: rate limit exceeded`, + `RESOURCE_EXHAUSTED: rate limit exceeded`, + `HTTP status 429: quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`, + `RESOURCE_EXHAUSTED: TPM quota exceeded`, + `HTTP status 401: invalid api key`, + `HTTP status 401: insufficient_quota`, + `HTTP status 403: forbidden`, + `HTTP status 403: billing_hard_limit_reached`, + `HTTP status 500: internal server error`, + `HTTP status 503: quota_exceeded`, + `dial tcp: network is unreachable`, + `context length exceeded`, + `invalid request parameter`, + } + for _, message := range negative { + if isQuotaExhaustedError(errors.New(message)) { + t.Errorf("unexpected quota classification for %q", message) + } + } + if isQuotaExhaustedError(nil) { + t.Error("nil error must not be classified as quota exhaustion") + } +} + +func TestTaskDTOUsesEmptyArrays(t *testing.T) { + t.Parallel() + dto := taskDTO(&Task{ID: "7"}, "created") + raw, err := json.Marshal(dto) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + t.Fatal(err) + } + for _, field := range []string{"source_task_ids", "llm_profile_ids"} { + value, ok := payload[field].([]any) + if !ok || len(value) != 0 { + t.Fatalf("%s must serialize as [], payload=%s", field, raw) + } + } +} + +func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) { + t.Parallel() + task := &Task{ID: "7"} + var wg sync.WaitGroup + for writer := 0; writer < 2; writer++ { + wg.Add(1) + go func(offset int) { + defer wg.Done() + for i := 0; i < 500; i++ { + id := int64(offset*1000 + i + 1) + task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "") + } + }(writer) + } + for reader := 0; reader < 4; reader++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 500; i++ { + state := task.llmStateSnapshot() + if len(state.ProfileIDs) > 0 { + state.ProfileIDs[0] = -1 + } + _ = taskDTO(task, "running") + } + }() + } + wg.Wait() + state := task.llmStateSnapshot() + if len(state.ProfileIDs) != 2 || state.ProfileIDs[0] <= 0 { + t.Fatalf("snapshot mutation leaked into task state: %+v", state) + } +} + +func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) { + t.Parallel() + quota := errors.New("openai: status 402: insufficient_quota") + first := &scriptedLLMProvider{err: quota} + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "fallback"}}} + active := int64(11) + transitions := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(selection taskLLMSelection, _ error) (db.TaskLLMTransition, error) { + if selection.profileID != 11 { + t.Fatalf("unexpected exhausted profile %d", selection.profileID) + } + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil { + t.Fatal(err) + } + if len(events) != 1 || events[0].Text != "fallback" { + t.Fatalf("unexpected fallback events: %+v", events) + } + if first.calls != 1 || second.calls != 1 || transitions != 1 { + t.Fatalf("calls first=%d second=%d transitions=%d", first.calls, second.calls, transitions) + } +} + +func TestTaskLLMStreamRetriesAfterNonContentStartEvent(t *testing.T) { + t.Parallel() + quota := errors.New("anthropic: status 402: insufficient credit balance") + first := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SEMessageStart}}, + err: quota, + } + second := &scriptedLLMProvider{events: []llm.StreamEvent{ + {Type: llm.SEMessageStart}, + {Type: llm.SETextDelta, Text: "fallback"}, + }} + active := int64(11) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil { + t.Fatal(err) + } + if len(events) != 2 || events[0].Type != llm.SEMessageStart || events[1].Text != "fallback" { + t.Fatalf("non-content event from failed attempt leaked or fallback missing: %+v", events) + } + if first.calls != 1 || second.calls != 1 { + t.Fatalf("calls first=%d second=%d", first.calls, second.calls) + } +} + +func TestTaskLLMStreamRetriesStalePreStreamFailureAgainstReplacementChain(t *testing.T) { + t.Parallel() + first := &scriptedLLMProvider{err: errors.New("status 402: insufficient_quota")} + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "replacement"}}} + current := first + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + profileID := int64(11) + if current == second { + profileID = 22 + } + return taskLLMSelection{profileID: profileID, provider: current}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + current = second + return db.TaskLLMTransition{PreviousProfileID: 11, Stale: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil || len(events) != 1 || events[0].Text != "replacement" { + t.Fatalf("stale call did not continue on replacement chain: events=%+v err=%v", events, err) + } +} + +func TestTaskLLMStreamDoesNotSwitchForOrdinaryErrors(t *testing.T) { + t.Parallel() + tests := []string{ + "openai: status 429: rate limit exceeded", + "openai: status 401: invalid api key", + "openai: status 500: internal server error", + "dial tcp: network is unreachable", + } + for _, message := range tests { + t.Run(message, func(t *testing.T) { + provider := &scriptedLLMProvider{err: errors.New(message)} + exhausted := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + exhausted++ + return db.TaskLLMTransition{}, nil + }, + } + _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err == nil || err.Error() != message { + t.Fatalf("error=%v, want %q", err, message) + } + if exhausted != 0 { + t.Fatalf("ordinary error changed failover state %d times", exhausted) + } + }) + } +} + +func TestTaskLLMStreamAdvancesAfterMidStreamQuotaWithoutReplay(t *testing.T) { + t.Parallel() + quota := errors.New("anthropic: status 402: credit balance is too low") + first := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "partial"}}, + err: quota, + } + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "next-call"}}} + active := int64(11) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if len(events) != 1 || events[0].Text != "partial" || err == nil { + t.Fatalf("mid-stream result events=%+v err=%v", events, err) + } + if first.calls != 1 || second.calls != 0 { + t.Fatalf("mid-stream call was replayed: first=%d second=%d", first.calls, second.calls) + } + + nextEvents, nextErr := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if nextErr != nil || len(nextEvents) != 1 || nextEvents[0].Text != "next-call" || second.calls != 1 { + t.Fatalf("next call did not use fallback: events=%+v err=%v calls=%d", nextEvents, nextErr, second.calls) + } +} + +func TestTaskLLMStreamMarkFailureAfterToolStartSkipsWorkerReplay(t *testing.T) { + t.Parallel() + quota := errors.New("openai: status 402: insufficient_quota") + markFailure := errors.New("database temporarily unavailable") + provider := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SEToolUseStart}}, + err: quota, + } + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + return db.TaskLLMTransition{}, markFailure + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if len(events) != 1 || events[0].Type != llm.SEToolUseStart || err == nil { + t.Fatalf("committed tool stream events=%+v err=%v", events, err) + } + if !isTaskLLMRuntimeError(err) { + t.Fatalf("committed mark failure must be router-classified: %T %v", err, err) + } + if retryableWorkerModelError(harness.ReasonModelError, err) { + t.Fatal("worker must not replay after a committed tool stream when quota persistence fails") + } + if provider.calls != 1 { + t.Fatalf("provider calls=%d, want 1", provider.calls) + } +} + +func TestTaskLLMRuntimeErrorSkipsWorkerReplay(t *testing.T) { + partialStreamErr := &taskLLMError{ + taskID: "7", + chainExhausted: false, + cause: errors.New("insufficient_quota after partial output"), + } + if retryableWorkerModelError(harness.ReasonModelError, partialStreamErr) { + t.Fatal("worker must not replay an intent after the task router handled a mid-stream quota error") + } + if !retryableWorkerModelError(harness.ReasonModelError, errors.New("temporary network error")) { + t.Fatal("ordinary model errors should retain the existing worker retry behavior") + } +} + +func TestEngineReadyForTaskOverride(t *testing.T) { + t.Parallel() + e := &Engine{} + if e.Ready() { + t.Fatal("engine without a global provider must not be globally ready") + } + e.SetAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + return &agent.Planner{}, &agent.Worker{} + }) + if !e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("task-specific provider chain must make that task ready") + } +} + +func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) { + t.Parallel() + engine := NewEngine(nil) + engine.started.Store("7", true) + task := &Task{ID: "7", Status: "running"} + profileID := int64(11) + task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted") + s := &Server{engine: engine} + if got := s.resolvedTaskStatus(task); got != "running" { + t.Fatalf("exhausted started status=%q, want running", got) + } + task.Paused = true + if got := s.resolvedTaskStatus(task); got != "paused" { + t.Fatalf("paused exhausted status=%q, want paused", got) + } +} + +func TestAuthoritativeTaskResolverDoesNotFallBackWhenChainUnavailable(t *testing.T) { + t.Parallel() + e := &Engine{} + e.UseLLM(&agent.Planner{}, &agent.Worker{}) + if !e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("global provider should initially make task ready") + } + e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + return nil, nil + }) + if e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("unavailable authoritative task chain must not fall back to global provider") + } +} + +func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) { + t.Parallel() + provider := &scriptedLLMProvider{err: errors.New("billing_not_active")} + transitions := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + return db.TaskLLMTransition{PreviousProfileID: 11, Advanced: true, ChainExhausted: true}, nil + }, + transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, + } + + _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if !isTaskLLMChainExhausted(err) { + t.Fatalf("expected chain-exhausted error, got %v", err) + } + if provider.calls != 1 || transitions != 1 { + t.Fatalf("calls=%d transitions=%d", provider.calls, transitions) + } +} + +func TestProfileDeleteRestoresQuotaBlockedIntentWhenFallbackAvailable(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + + profileID, err := m.pg.SaveProfile(&db.LLMProfile{ + Name: fmt.Sprintf("delete-quota-fallback-%d", time.Now().UnixNano()), + Format: "openai", + Model: "test-model", + APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + task, err := m.CreateTaskWithOptions("quota fallback", "resume intent", db.TaskCreateOptions{ + LLMProfileIDs: []int64{profileID}, + }) + if err != nil { + _ = m.pg.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) + _ = m.pg.DeleteProfile(profileID) + }) + + intentID, err := task.Store.AddIntent(map[string]any{"action": "retry after fallback"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + taskID, err := parseTaskID(task.ID) + if err != nil { + t.Fatal(err) + } + if transition, err := m.pg.MarkTaskLLMProfileQuotaExhausted(taskID, profileID, "insufficient_quota"); err != nil || !transition.ChainExhausted { + t.Fatalf("exhaust task chain: transition=%+v err=%v", transition, err) + } + if err := task.Store.SetIntentBlockedReason(intentID, db.IntentBlockedLLMQuota); err != nil { + t.Fatal(err) + } + task.Paused = true + if err := m.SetTaskPaused(task.ID, true); err != nil { + t.Fatal(err) + } + + // Removing the final explicit entry restores the legacy/global fallback. The + // post-delete sync must reopen only quota-blocked work without unpausing it. + if err := m.pg.DeleteProfile(profileID); err != nil { + t.Fatal(err) + } + s := &Server{ + m: m, + llmOn: true, + llmProv: &scriptedLLMProvider{}, + provByProfile: map[int64]*provEntry{}, + } + s.restoreTasksAfterProfileDelete(m.pg) + + node, err := task.Store.GetNode(intentID) + if err != nil { + t.Fatal(err) + } + if node == nil || node.State != "open" || node.BlockedReason != "" { + t.Fatalf("quota-blocked intent was not reopened: %+v", node) + } + if !task.Paused { + t.Fatal("profile deletion must not resume a paused task") + } + state := task.llmStateSnapshot() + if len(state.ProfileIDs) != 0 || state.FailoverState != "default" { + t.Fatalf("task did not return to fallback state: %+v", state) + } +} diff --git a/traffic/traffic.go b/traffic/traffic.go index ff1d5a3..49675d2 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -10,6 +10,7 @@ import ( "database/sql" "encoding/hex" "encoding/json" + "errors" "fmt" "io/fs" "log" @@ -24,9 +25,9 @@ import ( "sync/atomic" "time" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" - "github.com/Autumn-27/artex/db" mproxy "github.com/lqqyt2423/go-mitmproxy/proxy" _ "modernc.org/sqlite" ) @@ -471,28 +472,205 @@ func (t *Traffic) DeleteHost(host string) (int64, error) { // so picking "api.example.com" never sweeps "api.example.com.cn". Returns rows // deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass). func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) { - t.wmu.Lock() - defer t.wmu.Unlock() - var deleted int64 - removed := make(map[string]bool) - for _, h := range hosts { - res, err := t.db.Exec(`DELETE FROM exchanges WHERE host=?`, h) - if err != nil { - return deleted, err - } - n, _ := res.RowsAffected() - deleted += n - if !removed[h] { - os.RemoveAll(filepath.Join(t.dir, sanitize(h))) - removed[h] = true - } + stage, err := t.StageDeleteHostsExact(hosts) + if err != nil { + return 0, err } - if deleted > 0 { - _ = t.gcBlobs() + deleted := stage.Deleted() + if err := stage.Commit(); err != nil { + return deleted, err } return deleted, nil } +type stagedTrafficPath struct { + source string + staged string +} + +// HostDeleteStage keeps the SQLite delete transaction open and moves matching +// host trees to a same-filesystem staging directory. The Traffic write lock is +// held until Commit or Rollback, so no recorder can recreate a staged tree or a +// blob reference while a task deletion is being coordinated with PostgreSQL. +type HostDeleteStage struct { + traffic *Traffic + tx *sql.Tx + stageDir string + moves []stagedTrafficPath + deleted int64 + done bool +} + +// Deleted reports the number of exchange rows selected by this stage. +func (s *HostDeleteStage) Deleted() int64 { + if s == nil { + return 0 + } + return s.deleted +} + +// StageDeleteHostsExact prepares a reversible exact-host deletion. Callers must +// finish every successful stage with Commit or Rollback. +func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error) { + t.wmu.Lock() + stage := &HostDeleteStage{traffic: t} + fail := func(cause error) (*HostDeleteStage, error) { + if rollbackErr := stage.rollbackLocked(); rollbackErr != nil { + return nil, errors.Join(cause, fmt.Errorf("restore staged traffic: %w", rollbackErr)) + } + return nil, cause + } + + unique := make([]string, 0, len(hosts)) + seen := make(map[string]struct{}, len(hosts)) + for _, host := range hosts { + if _, ok := seen[host]; ok { + continue + } + seen[host] = struct{}{} + unique = append(unique, host) + } + + tx, err := t.db.Begin() + if err != nil { + t.wmu.Unlock() + stage.done = true + return nil, err + } + stage.tx = tx + for _, h := range unique { + res, err := tx.Exec(`DELETE FROM exchanges WHERE host=?`, h) + if err != nil { + return fail(err) + } + n, _ := res.RowsAffected() + stage.deleted += n + } + + stagedSources := make(map[string]struct{}, len(unique)) + for _, host := range unique { + source := filepath.Join(t.dir, sanitize(host)) + if _, duplicatePath := stagedSources[source]; duplicatePath { + continue + } + stagedSources[source] = struct{}{} + if _, err := os.Lstat(source); err != nil { + if os.IsNotExist(err) { + continue + } + return fail(fmt.Errorf("inspect traffic tree for %q: %w", host, err)) + } + if stage.stageDir == "" { + parent := filepath.Join(t.dir, "_delete_staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return fail(fmt.Errorf("create traffic staging directory: %w", err)) + } + stage.stageDir, err = os.MkdirTemp(parent, "hosts-") + if err != nil { + return fail(fmt.Errorf("create traffic stage: %w", err)) + } + } + staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source))) + if err := os.Rename(source, staged); err != nil { + return fail(fmt.Errorf("stage traffic tree for %q: %w", host, err)) + } + stage.moves = append(stage.moves, stagedTrafficPath{source: source, staged: staged}) + } + return stage, nil +} + +// Rollback restores both the SQLite rows and every staged host tree. +func (s *HostDeleteStage) Rollback() error { + if s == nil || s.done { + return nil + } + return s.rollbackLocked() +} + +func (s *HostDeleteStage) rollbackLocked() error { + if s == nil || s.done { + return nil + } + var errs []error + if s.tx != nil { + if err := s.tx.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) { + errs = append(errs, fmt.Errorf("rollback traffic index: %w", err)) + } + } + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) + } + } + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(errs...) +} + +// Commit makes the staged deletion permanent. SQLite is committed only after +// the caller has committed its PostgreSQL task deletion. +func (s *HostDeleteStage) Commit() error { + if s == nil || s.done { + return nil + } + if err := s.tx.Commit(); err != nil { + // A failed SQLite commit normally rolls the transaction back. Restore the + // trees so the traffic store remains internally consistent and recoverable. + restoreErr := s.restoreTreesLocked() + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(fmt.Errorf("commit traffic index deletion: %w", err), restoreErr) + } + var errs []error + if s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove staged traffic trees: %w", err)) + } + } + if err := s.traffic.gcBlobs(); err != nil { + errs = append(errs, fmt.Errorf("garbage collect traffic blobs: %w", err)) + } + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(errs...) +} + +func (s *HostDeleteStage) restoreTreesLocked() error { + var errs []error + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) + } + } + return errors.Join(errs...) +} + // gcBlobs removes content-addressed blobs in _blobs that no remaining exchange // tree references. References appear as "@blob sha256:" lines inside the // tree's request.http/response.http files (see bodyOrBlob). Best-effort: walk diff --git a/traffic/traffic_test.go b/traffic/traffic_test.go index 25d95b9..497bcf5 100644 --- a/traffic/traffic_test.go +++ b/traffic/traffic_test.go @@ -162,6 +162,152 @@ VALUES(?,?,?,?,?,?,?,?,?,?,?)`, } } +func TestDeleteHostsExactReportsTreeRemovalFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "api.example.com" + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, + "1-0001", 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/1-0001"); err != nil { + t.Fatal(err) + } + + notDir := filepath.Join(dir, "not-a-directory") + if err := os.WriteFile(notDir, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + tr.dir = notDir + + n, err := tr.DeleteHostsExact([]string{host}) + if err == nil { + t.Fatal("DeleteHostsExact returned nil after traffic tree removal failed") + } + if n != 0 { + t.Fatalf("deleted=%d, want 0 after atomic rollback", n) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("rolled-back index count=%d err=%v, want 1", count, err) + } +} + +func TestDeleteHostsExactRollsBackWholeIndexBatch(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + for i, host := range []string{"a.example.com", "b.example.com"} { + id := fmt.Sprintf("1-%04d", i+1) + if err := os.MkdirAll(filepath.Join(dir, host, "GET", id), 0o755); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + } + if _, err := tr.DB().Exec(`CREATE TRIGGER fail_second_host BEFORE DELETE ON exchanges +WHEN OLD.host='b.example.com' BEGIN SELECT RAISE(ABORT, 'forced delete failure'); END`); err != nil { + t.Fatal(err) + } + + if n, err := tr.DeleteHostsExact([]string{"a.example.com", "b.example.com"}); err == nil || n != 0 { + t.Fatalf("DeleteHostsExact failure=(%d,%v), want (0,error)", n, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges`).Scan(&count); err != nil || count != 2 { + t.Fatalf("rolled-back index count=%d err=%v, want 2", count, err) + } + for _, host := range []string{"a.example.com", "b.example.com"} { + if _, err := os.Stat(filepath.Join(dir, host)); err != nil { + t.Fatalf("tree %s changed despite index rollback: %v", host, err) + } + } +} + +func TestStageDeleteHostsExactRollbackRestoresIndexAndTree(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "rollback.example.com" + const id = "1-0001" + tree := filepath.Join(dir, host, "GET", id) + if err := os.MkdirAll(tree, 0o755); err != nil { + t.Fatal(err) + } + marker := filepath.Join(tree, "request.http") + if err := os.WriteFile(marker, []byte("original"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + if stage.Deleted() != 1 { + t.Fatalf("staged deleted=%d, want 1", stage.Deleted()) + } + if _, err := os.Stat(filepath.Join(dir, host)); !os.IsNotExist(err) { + t.Fatalf("host tree was not staged: %v", err) + } + if err := stage.Rollback(); err != nil { + t.Fatal(err) + } + if got, err := os.ReadFile(marker); err != nil || string(got) != "original" { + t.Fatalf("restored tree content=%q err=%v", got, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("restored index count=%d err=%v, want 1", count, err) + } +} + +func TestStageDeleteHostsExactRollbackReportsRestoreFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "conflict.example.com" + hostDir := filepath.Join(dir, host) + if err := os.MkdirAll(hostDir, 0o755); err != nil { + t.Fatal(err) + } + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + // Simulate an out-of-band conflicting destination. Rollback must surface the + // failed rename instead of claiming the external data was restored. + if err := os.WriteFile(hostDir, []byte("conflict"), 0o644); err != nil { + t.Fatal(err) + } + if err := stage.Rollback(); err == nil || !strings.Contains(err.Error(), "restore") { + t.Fatalf("rollback err=%v, want restore failure", err) + } + if _, err := os.Stat(stage.stageDir); err != nil { + t.Fatalf("staging was removed after failed restore: %v", err) + } +} + // TestDeleteHostGCBlobs verifies blob garbage collection: after a host's trees // are removed, blobs referenced by no remaining exchange are deleted, while // blobs still referenced (including shared ones) survive. diff --git a/web/src/app/(main)/_components/main-content.tsx b/web/src/app/(main)/_components/main-content.tsx index 3c5a39c..c320a19 100644 --- a/web/src/app/(main)/_components/main-content.tsx +++ b/web/src/app/(main)/_components/main-content.tsx @@ -6,9 +6,9 @@ import { usePathname } from "next/navigation"; import { Separator } from "@/components/ui/separator"; import { SidebarTrigger } from "@/components/ui/sidebar"; +import { useCurrentUser } from "@/hooks/use-current-user"; import { api } from "@/lib/api"; import { cn } from "@/lib/utils"; -import { useCurrentUser } from "@/hooks/use-current-user"; import { AccountSwitcher } from "./sidebar/account-switcher"; import { LayoutControls } from "./sidebar/layout-controls"; @@ -62,7 +62,7 @@ export function MainContent({ children }: { children: ReactNode }) {
{version && ( - 版本 · {version} + 版本 · {version} )} diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index 19f1c2a..d5590e0 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -2,29 +2,20 @@ import * as React from "react"; import Link from "next/link"; -import { useRouter, useSearchParams } from "next/navigation"; +import { useSearchParams } from "next/navigation"; import { toast } from "sonner"; -import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon, Trash2Icon } from "lucide-react"; +import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react"; import { SidebarTrigger } from "@/components/ui/sidebar"; import { Separator } from "@/components/ui/separator"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, - AlertDialogTrigger, -} from "@/components/ui/alert-dialog"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Select, SelectContent, + SelectGroup, SelectItem, SelectTrigger, } from "@/components/ui/select"; @@ -72,9 +63,9 @@ function FieldRow({ } function FindingDetailInner() { - const router = useRouter(); const searchParams = useSearchParams(); const id = searchParams.get("id") ?? ""; + const contextTaskId = searchParams.get("context_task") ?? ""; const [finding, setFinding] = React.useState(null); const [loaded, setLoaded] = React.useState(false); const [tab, setTab] = React.useState("overview"); @@ -85,18 +76,18 @@ function FindingDetailInner() { return; } api - .getFinding(id) + .getFinding(id, contextTaskId || undefined) .then((f) => setFinding(f)) .catch(() => setFinding(null)) .finally(() => setLoaded(true)); - }, [id]); + }, [contextTaskId, id]); React.useEffect(() => { load(); }, [load]); const changeSeverity = React.useCallback( async (next: Severity) => { - if (!finding || next === finding.severity) return; + if (!finding || finding.inherited || next === finding.severity) return; const prev = finding.severity; setFinding({ ...finding, severity: next }); try { @@ -113,7 +104,7 @@ function FindingDetailInner() { const changeStatus = React.useCallback( async (next: FindingStatus) => { - if (!finding || next === finding.status) return; + if (!finding || finding.inherited || next === finding.status) return; const prev = finding.status; setFinding({ ...finding, status: next }); try { @@ -128,18 +119,6 @@ function FindingDetailInner() { [finding, id], ); - // deleteFinding 删除当前漏洞(需二次确认),成功后返回发现列表。 - const deleteFinding = React.useCallback(async () => { - if (!id) return; - try { - await api.deleteFinding(id); - toast.success("已删除漏洞"); - router.push("/function/findings"); - } catch (e) { - toast.error("删除失败:" + (e as Error).message); - } - }, [id, router]); - if (!finding) { return (
@@ -184,25 +163,11 @@ function FindingDetailInner() { - - - - - - - 确认删除该漏洞? - - 「{title}」将被永久删除,同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 - - - - 取消 - 删除 - - - + {finding.inherited && finding.source_task_id && ( + + 来源任务 #{finding.source_task_id} · 只读 + + )}
概览 @@ -271,48 +236,60 @@ function FindingDetailInner() {
- {/* 严重等级(可改) */} + {/* 严重等级 */} - changeSeverity(v as Severity)} > - - - - {SEVERITIES.map((sv) => ( - - {statusMeta("severity", sv).label} - - ))} - - + + + + + + {SEVERITIES.map((sv) => ( + + {statusMeta("severity", sv).label} + + ))} + + + + )} - {/* 处理状态(可改) */} + {/* 处理状态 */} - changeStatus(v as FindingStatus)} > - - - - {FINDING_STATUSES.map((st) => ( - - {statusMeta("finding", st).label} - - ))} - - + + + + + + {FINDING_STATUSES.map((st) => ( + + {statusMeta("finding", st).label} + + ))} + + + + )} {/* 漏洞类型 */} diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx index 7d066c6..c890352 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx @@ -152,7 +152,6 @@ const TABS: { key: string; label: string; icon: LucideIcon }[] = [ ]; export function AssetsTab({ taskId }: { taskId: string }) { - // 服务端分页:只取当前分类的当前页,分类总数走 /assets/counts?task_id=,不做数量截断 const [rows, setRows] = React.useState([]); const [total, setTotal] = React.useState(0); const [counts, setCounts] = React.useState>({}); @@ -161,7 +160,6 @@ export function AssetsTab({ taskId }: { taskId: string }) { const [page, setPage] = React.useState(0); const [size, setSize] = React.useState(50); - // 换分类/换页大小:回到第一页,并清掉旧行(各分类列不同,不能串用) React.useEffect(() => { setPage(0); setRows([]); @@ -171,17 +169,16 @@ export function AssetsTab({ taskId }: { taskId: string }) { let alive = true; const load = async () => { try { - const [cur, c] = await Promise.all([ + const [cur, nextCounts] = await Promise.all([ api.taskAssets(taskId, tab, size, page * size), api.assetCounts(taskId), ]); if (!alive) return; setRows(cur.assets); setTotal(cur.total); - setCounts(c ?? {}); + setCounts(nextCounts ?? {}); + setLoaded(true); } catch { - // 保留上一次的数据 - } finally { if (alive) setLoaded(true); } }; @@ -190,7 +187,6 @@ export function AssetsTab({ taskId }: { taskId: string }) { return () => { alive = false; clearInterval(timer); }; }, [taskId, tab, page, size]); - // 轮询期间资产变少时,别停在越界的空页上 React.useEffect(() => { const maxPage = Math.max(0, Math.ceil(total / size) - 1); if (page > maxPage) setPage(maxPage); diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx index 14201fe..e5cfdda 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx @@ -14,9 +14,9 @@ import { Waypoints, } from "lucide-react"; -import { cn } from "@/lib/utils"; -import { Card, CardContent } from "@/components/ui/card"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { Card, CardContent } from "@/components/ui/card"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Sheet, @@ -26,6 +26,7 @@ import { } from "@/components/ui/sheet"; import { api } from "@/lib/api"; import type { CoverageAssetRef, CoverageAssetRefs, CoverageGraphEdge, CoverageGraphNode } from "@/lib/types"; +import { cn } from "@/lib/utils"; // 每个父节点下、同一类型的子节点默认展示的数量;超出折叠,"展示更多"每次再拉这么多。 const FOLD_LIMIT = 20; @@ -274,10 +275,18 @@ function RefList({ title, items }: { title: string; items: CoverageAssetRef[] })
{items.map((r) => ( -
- #{r.id} - {r.state && {r.state}} - {r.summary || "—"} +
+ #{r.id} + {r.state && {r.state}} + {r.summary || "—"} + {r.inherited && r.source_task_id && ( + + 来源 #{r.source_task_id} · 只读 + + )}
))}
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 d4faf24..07d9449 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 @@ -2,25 +2,15 @@ import * as React from "react"; import Link from "next/link"; -import { ArrowUpRightIcon, ChevronRightIcon, Trash2Icon } from "lucide-react"; +import { ArrowUpRightIcon, ChevronRightIcon } from "lucide-react"; import { cn } from "@/lib/utils"; import { Card, CardContent } from "@/components/ui/card"; +import { Badge } from "@/components/ui/badge"; import { StatusBadge } from "@/components/status-badge"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, - AlertDialogTrigger, -} from "@/components/ui/alert-dialog"; -import { Button } from "@/components/ui/button"; import { Select, SelectContent, + SelectGroup, SelectItem, SelectTrigger, } from "@/components/ui/select"; @@ -42,12 +32,12 @@ const FINDING_STATUSES: FindingStatus[] = [ function Row({ f, + contextTaskId, onStatus, - onDelete, }: { f: Finding; + contextTaskId: string; onStatus: (f: Finding, next: FindingStatus) => void; - onDelete: (f: Finding) => void; }) { const [open, setOpen] = React.useState(false); return ( @@ -65,10 +55,17 @@ function Row({ )} /> -
- - {f.name || f.vulnclass || "未分类"} - +
+
+ + {f.name || f.vulnclass || "未分类"} + + {f.inherited && f.source_task_id && ( + + 来源 #{f.source_task_id} · 只读 + + )} +
{f.summary} @@ -92,7 +89,7 @@ function Row({ )}
)} - {f.finding_id ? ( + {f.finding_id && !f.inherited ? ( ) : ( @@ -119,7 +118,11 @@ function Row({ {f.finding_id && ( @@ -130,30 +133,8 @@ function Row({
{open && (
-
- 证据 / PoC - {f.finding_id && ( - - - - - - - 确认删除该漏洞? - - 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除, - 同时从发现列表、任务发现 Tab 与探索图中移除,此操作不可撤销。 - - - - 取消 - onDelete(f)}>删除 - - - - )} +
+ 证据 / PoC
             {f.evidence}
@@ -187,7 +168,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
 
   const onStatus = React.useCallback(
     async (f: Finding, next: FindingStatus) => {
-      if (!f.finding_id || next === f.status) return;
+      if (f.inherited || !f.finding_id || next === f.status) return;
       const prev = f.status;
       setFindings((cur) =>
         cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)),
@@ -205,19 +186,8 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     [],
   );
 
-  const onDelete = React.useCallback(async (f: Finding) => {
-    if (!f.finding_id) return;
-    try {
-      await api.deleteFinding(f.finding_id);
-      setFindings((cur) => cur.filter((x) => x.id !== f.id));
-      toast.success("已删除漏洞");
-    } catch (e) {
-      toast.error("删除失败:" + (e as Error).message);
-    }
-  }, []);
-
   const items = findings
-    .filter((f) => f.task_id === taskId)
+    .filter((f) => f.task_id === taskId || f.inherited)
     .sort((a, b) => {
       const order = { critical: 0, high: 1, medium: 2, low: 3 };
       return order[a.severity] - order[b.severity];
@@ -227,11 +197,11 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     
       
         {items.map((f) => (
-          
+          
         ))}
         {items.length === 0 && (
           

- 本任务暂无确认发现。 + 本任务及直接关联任务暂无确认发现。

)}
diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx index 090adb3..056b786 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx @@ -1,8 +1,11 @@ "use client"; import * as React from "react"; -import { FileTextIcon, CopyIcon, CheckIcon } from "lucide-react"; + +import { CheckIcon, CopyIcon, FileTextIcon } from "lucide-react"; import { toast } from "sonner"; + +import { Markdown } from "@/components/markdown"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { api } from "@/lib/api"; @@ -39,6 +42,29 @@ export function ReportTab({ taskId }: { taskId: string }) { setTimeout(() => setCopied(false), 1500); } + let content: React.ReactNode; + if (loading) { + content = ( +
+ + 加载中… +
+ ); + } else if (report) { + content = ( +
+ +
+ ); + } else { + content = ( +
+ + 暂无报告 +
+ ); + } + return ( @@ -53,23 +79,7 @@ export function ReportTab({ taskId }: { taskId: string }) { )}
- - {loading ? ( -
- - 加载中… -
- ) : report ? ( -
-            {report}
-          
- ) : ( -
- - 暂无报告 -
- )} -
+ {content} ); } diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx index 1365089..1e18d1b 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx @@ -1,30 +1,46 @@ "use client"; import * as React from "react"; + import { - RadioIcon, BrainIcon, - UserIcon, - Loader2Icon, - ZapOffIcon, - SendIcon, - SquareIcon, - ClockIcon, CircleCheckIcon, - CircleXIcon, CircleSlashIcon, - ShieldAlertIcon, - WifiOffIcon, - RotateCwIcon, + CircleXIcon, + ClockIcon, + HistoryIcon, + Loader2Icon, PaperclipIcon, + PauseIcon, + PlayIcon, + RadioIcon, + RotateCwIcon, + SendIcon, + ShieldAlertIcon, + SquareIcon, + UserIcon, + WifiOffIcon, XIcon, + ZapOffIcon, + Trash2Icon, } from "lucide-react"; import { toast } from "sonner"; -import { cn } from "@/lib/utils"; -import { Transcript } from "@/components/transcript"; + import { TodoPopover } from "@/components/todo-popover"; +import { Transcript } from "@/components/transcript"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { Textarea } from "@/components/ui/textarea"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupTextarea } from "@/components/ui/input-group"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Tooltip, @@ -43,6 +59,7 @@ import type { TaskNode, TokenTotal, } from "@/lib/types"; +import { cn } from "@/lib/utils"; // fmtBytes renders a human file size for attachment chips (mirrors transcript.tsx). function fmtBytes(n: number): string { @@ -63,6 +80,7 @@ function fmtBytes(n: number): string { // so refresh / tab-switch / sleep / reconnect never drop the newest records. const PAGE = 200; // history page size +const SYSTEM_SCAN_PAGE = 500; // generic activity pages scanned to recover sparse system audit events const MAX_KEEP = 4000; // per-session in-memory cap; older pages re-fetched on scroll-up const STREAM_WINDOW_MS = 5000; // "live" = activity seen within this window @@ -98,6 +116,7 @@ function emptyState(): SessionState { // sessionKeyOf routes an activity to its stable session key. worker="planner" covers // BOTH the Goal Agent's round-0 decomposition and the Planner (single Plan session). function sessionKeyOf(a: Activity): string { + if (a.worker === "system" || a.kind === "llm_switch" || a.kind === "llm_failover") return "system"; if (a.worker === "mainagent") return "main"; if (a.worker === "planner") return "plan"; if (a.intent_id) return `intent:${a.intent_id}`; @@ -122,6 +141,8 @@ function statusIcon(status: SessionStatus) { switch (status) { case "running": // 执行中 return ; + case "paused": + return ; case "pending": // 待领取(open intent) return ; case "done": // 完成 @@ -156,6 +177,7 @@ const roleMeta = { mainagent: { label: "主 Agent", icon: UserIcon }, planner: { label: "规划 Planner", icon: BrainIcon }, worker: { label: "Workers", icon: RadioIcon }, + system: { label: "系统审计", icon: HistoryIcon }, } as const; // The main-agent session is the interactive entry point of this tab and has no @@ -185,10 +207,24 @@ const PLANNER_SESSION: Session = { last_activity: "", }; +// LLM provider switches are task-level audit events rather than agent output. The +// history endpoint has no "system" filter, so this fixed session is populated by +// scanning the generic incremental activity endpoint and then tailed by task SSE. +const SYSTEM_ID = "s-system"; +const SYSTEM_SESSION: Session = { + id: SYSTEM_ID, + role: "system", + title: "系统事件 · LLM 故障转移", + status: "done", + live: false, + last_activity: "", +}; + // keyForSession maps a UI Session → its stable store key (main | plan | intent:). function keyForSession(s: Session): string { if (s.role === "mainagent") return "main"; if (s.role === "planner") return "plan"; + if (s.role === "system") return "system"; return `intent:${s.intent_id}`; } @@ -203,6 +239,8 @@ function intentStatus(state: string): SessionStatus { return "exhausted"; case "stopped": return "stopped"; + case "paused": + return "paused"; case "open": // 待领取,区别于执行中 return "pending"; default: // running @@ -220,9 +258,11 @@ function intentToSession(n: TaskNode): Session { role: "worker", title: label || `Intent ${n.id}`, status: state, - live: state === "running", + live: !n.inherited && state === "running", last_activity: n.ts, intent_id: n.id, + source_task_id: n.source_task_id, + inherited: n.inherited, }; } @@ -233,6 +273,10 @@ function SessionItem({ hasPending, unread, onClick, + onPause, + onResume, + onCancel, + controlling, }: { s: Session; active: boolean; @@ -240,6 +284,10 @@ function SessionItem({ hasPending?: boolean; unread?: number; onClick: () => void; + onPause?: () => void; + onResume?: () => void; + onCancel?: () => void; + controlling?: boolean; }) { const icon = s.role === "worker" ? ( @@ -248,36 +296,57 @@ function SessionItem({ ) : null; + const controllable = s.role === "worker" && !s.inherited && (s.status === "running" || s.status === "paused"); return ( - + {controllable && ( +
+ {s.status === "running" ? ( + + ) : ( + + )} + +
)} - {displayTitle} - {hasPending && ( - - )} - {!active && unread ? ( - - {unread > 99 ? "99+" : unread} - - ) : null} - {s.live && ( - - - 实时 - - )} - +
); } @@ -295,6 +364,8 @@ export function SessionsTab({ taskId }: { taskId: string }) { const [input, setInput] = React.useState(""); const [sending, setSending] = React.useState(false); const [stopping, setStopping] = React.useState(false); + const [controllingIntent, setControllingIntent] = React.useState(null); + const [cancelIntent, setCancelIntent] = React.useState(null); // 方式1 文件上传:选好的附件(已落到任务工作目录 uploads/),随下条消息一起发。 const [attachments, setAttachments] = React.useState([]); const [uploading, setUploading] = React.useState(false); @@ -313,6 +384,42 @@ export function SessionsTab({ taskId }: { taskId: string }) { if (fileInputRef.current) fileInputRef.current.value = ""; } } + + const patchIntentState = React.useCallback((intentId: string, state?: string) => { + const patch = (rows: TaskNode[]) => + state ? rows.map((row) => (row.id === intentId ? { ...row, state } : row)) : rows.filter((row) => row.id !== intentId); + setIntents(patch); + setOlderIntents(patch); + }, []); + + const controlWorker = React.useCallback(async (session: Session, action: "pause" | "resume" | "cancel") => { + if (!session.intent_id || session.inherited || controllingIntent) return; + setControllingIntent(session.intent_id); + try { + const result = await api.controlIntent(taskId, session.intent_id, action); + if (action === "pause") { + patchIntentState(session.intent_id, "paused"); + toast.success(`Worker #${session.intent_id} 已暂停`); + } else if (action === "resume") { + patchIntentState(session.intent_id, "open"); + toast.success(`Worker #${session.intent_id} 已恢复,等待重新领取`); + } else { + patchIntentState(session.intent_id); + if (activeId === session.id) setActiveId(MAIN_ID); + const deleted = result.deleted; + toast.success( + deleted + ? `Worker #${session.intent_id} 已取消,清理 ${deleted.facts} 条事实、${deleted.findings} 个漏洞` + : `Worker #${session.intent_id} 已取消`, + ); + } + } catch (error) { + toast.error("Worker 操作失败:" + (error as Error).message); + } finally { + setControllingIntent(null); + setCancelIntent(null); + } + }, [activeId, controllingIntent, patchIntentState, taskId]); // SSE connection state — surfaced so a dropped realtime link is visible, never // silently shown as "no messages". const [sseLive, setSseLive] = React.useState(false); @@ -326,6 +433,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { const esRef = React.useRef(null); const activeKeyRef = React.useRef("main"); // current session key (for SSE dispatch/unread) const atBottomRef = React.useRef(true); // transcript pinned to bottom? + const llmToastSeqRef = React.useRef>(new Set()); // Per-key request token: a stale response for a key is ignored (guards fast // latest/older interleaving). Writes are ALWAYS keyed, so a late response can only // touch its own session cache — never the currently-viewed one (see §7.5). @@ -353,6 +461,48 @@ export function SessionsTab({ taskId }: { taskId: string }) { const token = (reqTokenRef.current[key] ?? 0) + 1; reqTokenRef.current[key] = token; patchStore(key, (s) => ({ ...s, loading: true, error: undefined })); + + if (key === "system") { + void (async () => { + let since = 0; + let systemItems: Activity[] = []; + for (;;) { + const page = await api.activity(taskId, { since, limit: SYSTEM_SCAN_PAGE }); + if (reqTokenRef.current[key] !== token) return; + systemItems = mergeBySeq( + systemItems, + page.items.filter((item) => item.worker === "system" || item.kind === "llm_failover"), + ); + if (page.items.length < SYSTEM_SCAN_PAGE || page.cursor <= since) break; + since = page.cursor; + } + patchStore(key, (s) => { + const items = mergeBySeq(systemItems, s.items); + return { + ...s, + items, + loaded: true, + loading: false, + hasMore: false, + earliestSeq: items.length ? items[0].seq : 0, + unread: 0, + lastTs: items.length ? items[items.length - 1].ts : s.lastTs, + error: undefined, + }; + }); + })() + .catch((error) => { + if (reqTokenRef.current[key] !== token) return; + patchStore(key, (s) => ({ + ...s, + loading: false, + error: (error as Error).message || "加载失败", + })); + }) + .finally(() => loadingKeysRef.current.delete(key)); + return; + } + api .activityHistory(taskId, key, 0, PAGE) .then((r) => { @@ -424,8 +574,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { .then((r) => { if (alive) setTaskTokens(r.total); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { alive = false; @@ -438,8 +590,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { const load = () => api.interceptTask(taskId) .then((rows) => { if (alive) setPendingIntercepts(rows.filter((r) => r.status === "pending")); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { alive = false; clearInterval(t); }; }, [taskId]); @@ -490,9 +644,15 @@ export function SessionsTab({ taskId }: { taskId: string }) { st.lastTs = st.items.length ? st.items[st.items.length - 1].ts : ""; } buckets.main ??= { ...emptyState(), loaded: true }; + buckets.system ??= { ...emptyState(), loaded: true }; setStore(buckets); }) - .catch(() => setStore({ main: { ...emptyState(), loaded: true } })); + .catch(() => + setStore({ + main: { ...emptyState(), loaded: true }, + system: { ...emptyState(), loaded: true }, + }), + ); return () => { alive = false; }; @@ -532,6 +692,17 @@ export function SessionsTab({ taskId }: { taskId: string }) { } catch { return; // ignore malformed frame } + if ((a.kind === "llm_switch" || a.kind === "llm_failover") && !llmToastSeqRef.current.has(a.seq)) { + llmToastSeqRef.current.add(a.seq); + const transition = a.metadata?.llm_transition; + if (transition?.mode === "exhausted" || a.is_error) { + toast.error(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } else if (transition?.mode === "automatic") { + toast.success(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } else { + toast.info(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } + } const k = sessionKeyOf(a); setStore((prev) => { const cur = prev[k] ?? emptyState(); @@ -569,8 +740,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { esRef.current?.close(); esRef.current = null; }; - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [taskId]); + }, [taskId, patchStore]); // ── worker (intent) session list — paged, poll first page lightly ─────────────── React.useEffect(() => { @@ -583,8 +753,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { setIntents(r.items); setIntentsHasMore(r.hasMore); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { active = false; @@ -607,7 +779,9 @@ export function SessionsTab({ taskId }: { taskId: string }) { }); setIntentsHasMore(r.hasMore); }) - .catch(() => {}) + .catch(() => { + // A later manual retry can fetch this page again. + }) .finally(() => setLoadingOlderIntents(false)); }, [taskId, intents, olderIntents, loadingOlderIntents]); @@ -650,6 +824,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { // Worker agent_name format: "work#N · #intentID". Main/planner match by role key. const hasPendingForSession = React.useCallback( (s: Session): boolean => { + if (s.inherited) return false; if (!pendingIntercepts.length) return false; if (s.role === "mainagent") return pendingIntercepts.some((r) => r.agent_name === "mainagent"); if (s.role === "planner") return pendingIntercepts.some((r) => r.agent_name === "planner"); @@ -680,6 +855,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { { ...MAIN_SESSION, live: mainLive }, { ...PLANNER_SESSION, live: plannerLive }, ...workerSessions, + SYSTEM_SESSION, ], [workerSessions, mainLive, plannerLive], ); @@ -688,16 +864,19 @@ export function SessionsTab({ taskId }: { taskId: string }) { mainagent: sessions.filter((s) => s.role === "mainagent"), planner: sessions.filter((s) => s.role === "planner"), worker: sessions.filter((s) => s.role === "worker"), + system: sessions.filter((s) => s.role === "system"), }; const active = sessions.find((s) => s.id === activeId) ?? MAIN_SESSION; const isMain = active.role === "mainagent"; const isPlanner = active.role === "planner"; + const isSystem = active.role === "system"; const activeKey = keyForSession(active); const activeState = store[activeKey]; // Keep the SSE dispatcher's notion of the active session current, and lazily load // + clear unread whenever the active session changes. + // biome-ignore lint/correctness/useExhaustiveDependencies: cache updates must not reactivate the current session. React.useEffect(() => { activeKeyRef.current = activeKey; const st = store[activeKey]; @@ -706,7 +885,6 @@ export function SessionsTab({ taskId }: { taskId: string }) { } else if (st.unread) { patchStore(activeKey, (s) => ({ ...s, unread: 0 })); } - // eslint-disable-next-line react-hooks/exhaustive-deps }, [activeKey]); // Main agent is the human↔orchestrator CONSOLE: only the conversation (user msgs + @@ -721,7 +899,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { const pending = chatExtra.filter((a) => !seen.has(`${a.kind} ${a.summary}`)); return [...items, ...pending].sort((a, b) => a.seq - b.seq); } - if (isPlanner) return items; + if (isPlanner || isSystem) return items; // Worker session: the intent leads the transcript as a right-aligned "user"-style // message (the task handed to this worker), followed by its execution steps. const intentTitle = sessionMeta.get(active.id)?.title ?? active.title; @@ -731,9 +909,23 @@ export function SessionsTab({ taskId }: { taskId: string }) { ts: active.last_activity || "", kind: "intent", // LLM-generated objective — rendered as a distinct (non-human) bubble summary: intentTitle, + source_task_id: active.source_task_id, + inherited: active.inherited, }; return [intentMsg, ...items]; - }, [isMain, isPlanner, activeState, chatExtra, active.id, active.title, active.last_activity, sessionMeta]); + }, [ + isMain, + isPlanner, + isSystem, + activeState, + chatExtra, + active.id, + active.title, + active.last_activity, + active.source_task_id, + active.inherited, + sessionMeta, + ]); // seq of this session's most-recent TodoWrite call — for the Todo popover. const latestTodoSeq = React.useMemo(() => { @@ -791,6 +983,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { null, [], ); + // biome-ignore lint/correctness/useExhaustiveDependencies: activeId intentionally rebinds the scroll listener. React.useEffect(() => { const vp = viewport(); if (!vp) return; @@ -802,6 +995,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { return () => vp.removeEventListener("scroll", onScroll); }, [viewport, activeId, loadEarlier]); // open/switch a session → jump to the latest (bottom) + // biome-ignore lint/correctness/useExhaustiveDependencies: activeId intentionally scrolls a newly selected session. React.useLayoutEffect(() => { const vp = viewport(); if (vp) { @@ -810,6 +1004,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { } }, [activeId, viewport]); // new activity → stick to bottom only if the user is already pinned there + // biome-ignore lint/correctness/useExhaustiveDependencies: activity growth intentionally drives live-edge scrolling. React.useLayoutEffect(() => { if (!atBottomRef.current) return; const vp = viewport(); @@ -819,7 +1014,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { function stop() { if (stopping) return; setStopping(true); - api.stopChat(taskId).finally(() => setStopping(false)); + void api.stopChat(taskId).finally(() => setStopping(false)); } function send() { @@ -906,7 +1101,7 @@ export function SessionsTab({ taskId }: { taskId: string }) {
- {(["mainagent", "planner", "worker"] as const).map((role) => { + {(["mainagent", "planner", "system", "worker"] as const).map((role) => { const items = grouped[role]; if (!items.length) return null; const Meta = roleMeta[role]; @@ -927,11 +1122,16 @@ export function SessionsTab({ taskId }: { taskId: string }) { hasPending={hasPendingForSession(s)} unread={store[keyForSession(s)]?.unread} onClick={() => setActiveId(s.id)} + controlling={controllingIntent === s.intent_id} + onPause={() => void controlWorker(s, "pause")} + onResume={() => void controlWorker(s, "resume")} + onCancel={() => setCancelIntent(s)} /> ); })} {role === "worker" && intentsHasMore && (
)} -
- void pickFiles(e.target.files)} - /> - -