From cf0989c67d4bdbd953ae7db21d73f9add1f6a246 Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 07:37:24 +0800 Subject: [PATCH 1/7] =?UTF-8?q?=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1、修复流量包里的复制按钮不工作问题 2、修复新建任务关联资产不同步到任务测试范围问题 --- agent/blackboard_inheritance_test.go | 105 ++++++++++++++++++- db/companies_test.go | 10 +- db/tasks.go | 37 ++++++- db/tasks_test.go | 79 +++++++++++++- server/server.go | 2 +- web/src/app/(main)/function/tasks/page.tsx | 4 +- web/src/app/(main)/function/traffic/page.tsx | 10 +- web/src/lib/mock/handler.ts | 79 +++++++++++--- web/src/lib/task-assets.ts | 1 + web/src/lib/types.ts | 2 +- 10 files changed, 297 insertions(+), 32 deletions(-) diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go index fcbcb6f..ad4fe4b 100644 --- a/agent/blackboard_inheritance_test.go +++ b/agent/blackboard_inheritance_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "testing" "time" @@ -51,6 +52,12 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) { if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 { t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors) } + assets := d.Assets() + assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: domain}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) }) task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{ CompanyIDs: []int64{companyID}, @@ -62,7 +69,14 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) { tools := NewToolSet(d.Exploration(task.ExplorationID), "planner") tools.SetTaskID(task.ID) - tools.SetAssetStore(d.Assets(), companies) + tools.SetAssetStore(assets, companies) + linkedAssets, err := assets.QueryByTask(task.ID, "root_domain", 10, 0) + if err != nil || len(linkedAssets) != 1 || linkedAssets[0].ID != assetID { + t.Fatalf("company asset was not linked to task: assets=%+v err=%v", linkedAssets, err) + } + if linkedAssets[0].TaskSource != "company" || !strings.Contains(linkedAssets[0].TaskSourceSummary, companiesName(t, companies, companyID)) { + t.Fatalf("company asset provenance missing: %+v", linkedAssets[0]) + } overview := tools.graphOverviewData() coverage, ok := overview["coverage"].(map[string]any) if !ok { @@ -89,6 +103,95 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) { if !ok || len(keywords) != 1 || keywords[0] != keyword { t.Fatalf("company keywords missing: %#v", scopeRows[0]["company_keywords"]) } + hosts, ok := coverage["hosts"].([]string) + if !ok || !containsString(hosts, domain) { + t.Fatalf("company asset host missing from agent context: %#v", coverage["hosts"]) + } + untested := callReadJSON(t, tools.listUntestedAssets(), `{"type":"root_domain","page":1,"page_size":10}`) + if !strings.Contains(fmt.Sprint(untested), domain) { + t.Fatalf("company asset missing from untested backlog: %#v", untested) + } + byTask := callReadJSON(t, tools.listAssets(), fmt.Sprintf( + `{"dsl":"task_id==%d","type":"root_domain","limit":10}`, task.ID, + )) + if !strings.Contains(fmt.Sprint(byTask), domain) { + t.Fatalf("company asset missing from task-scoped agent query: %#v", byTask) + } + byCompany := callReadJSON(t, tools.listAssets(), fmt.Sprintf( + `{"dsl":"company_id==%d","type":"root_domain","limit":10}`, companyID, + )) + if !strings.Contains(fmt.Sprint(byCompany), domain) { + t.Fatalf("company asset missing from company-scoped agent query: %#v", byCompany) + } + assetJSON, _ := json.Marshal(assetID) + intentID, err := tools.addOneIntent(intentItem{ + Summary: "test associated company asset", + AssetIDs: []json.RawMessage{assetJSON}, + }) + if err != nil { + t.Fatal(err) + } + workerAssets, err := assets.IntentAssets(task.ID) + if err != nil || len(workerAssets) != 1 || workerAssets[0].IntentID != intentID || workerAssets[0].AssetID != assetID { + t.Fatalf("worker target did not retain company asset: assets=%+v err=%v", workerAssets, err) + } + + coverageDisabled := false + disabledTask, err := d.CreateTaskWithOptions("company context without coverage", "still expose company assets", db.TaskCreateOptions{ + CompanyIDs: []int64{companyID}, CoverageEnabled: &coverageDisabled, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(disabledTask.ID) }) + disabledTools := NewToolSet(d.Exploration(disabledTask.ExplorationID), "planner") + disabledTools.SetTaskID(disabledTask.ID) + disabledTools.SetCoverageEnabled(false) + disabledTools.SetAssetStore(assets, companies) + disabledOverview := disabledTools.graphOverviewData() + disabledCoverage, ok := disabledOverview["coverage"].(map[string]any) + if !ok { + t.Fatalf("coverage-disabled task lost asset context: %#v", disabledOverview["coverage"]) + } + disabledHosts, ok := disabledCoverage["hosts"].([]string) + if !ok || !containsString(disabledHosts, domain) { + t.Fatalf("coverage-disabled task lost company asset host: %#v", disabledCoverage["hosts"]) + } + if _, exists := disabledCoverage["denominator"]; exists { + t.Fatalf("coverage-disabled task unexpectedly exposed metrics: %#v", disabledCoverage) + } + if linked, err := assets.QueryByTask(disabledTask.ID, "root_domain", 10, 0); err != nil || len(linked) != 1 || linked[0].ID != assetID { + t.Fatalf("coverage-disabled task asset link=%+v err=%v", linked, err) + } + disabledIntentID, err := disabledTools.addOneIntent(intentItem{ + Summary: "test associated company asset without coverage metrics", + AssetIDs: []json.RawMessage{assetJSON}, + }) + if err != nil { + t.Fatal(err) + } + disabledWorkerAssets, err := assets.IntentAssets(disabledTask.ID) + if err != nil || len(disabledWorkerAssets) != 1 || disabledWorkerAssets[0].IntentID != disabledIntentID || disabledWorkerAssets[0].AssetID != assetID { + t.Fatalf("coverage-disabled worker target=%+v err=%v", disabledWorkerAssets, err) + } +} + +func containsString(items []string, target string) bool { + for _, item := range items { + if item == target { + return true + } + } + return false +} + +func companiesName(t *testing.T, companies *db.CompanyStore, companyID int64) string { + t.Helper() + company, err := companies.GetCompany(companyID) + if err != nil || company == nil { + t.Fatalf("company %d: company=%+v err=%v", companyID, company, err) + } + return company.Name } func TestBlackboardToolsReadDirectSources(t *testing.T) { diff --git a/db/companies_test.go b/db/companies_test.go index f1592c8..465966d 100644 --- a/db/companies_test.go +++ b/db/companies_test.go @@ -177,9 +177,13 @@ func TestCompanyScopeInvalid(t *testing.T) { } defer cleanupCompany(d, id) - // TLD-only domains and overly broad CIDRs should be rejected - lines := []string{"com", "1.2.3.4/8"} // rejected - added, _, invalid, _ := cs.AddScope(id, lines, "test") + // Explicitly typed TLD-only domains and overly broad CIDRs should be + // rejected. Untyped plain text is intentionally classified as a keyword. + inputs := []ScopeInput{ + {Kind: "domain", Value: "com"}, + {Kind: "cidr", Value: "1.2.3.4/8"}, + } + added, _, invalid, _ := cs.AddScopeInputs(id, inputs, "test") if added != 0 { t.Errorf("want 0 added for invalid lines, got %d", added) } diff --git a/db/tasks.go b/db/tasks.go index 115ed0c..f658305 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -89,10 +89,12 @@ const MinPlanHeartbeatSeconds = 600 const MaxTaskSourceCount = 8 // MaxTaskCompanyCount bounds the number of company asset scopes attached to a -// task. Company scopes are prompt context, so an unbounded list would inflate -// every planner and worker call. +// task. Company scopes are prompt context, and their currently attributed +// assets are snapshotted into the task at creation time. const MaxTaskCompanyCount = 32 +const taskCompanyAssetSource = "company" + var ( ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids") ErrTaskCompanyNotFound = errors.New("task company not found") @@ -248,6 +250,11 @@ func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error { if len(companyIDs) == 0 { return nil } + // Company scope edits rebuild assets.company_id. Serialize the creation-time + // snapshot with those edits so the task sees one committed attribution state. + if err := lockCompanyScopeMutation(tx); err != nil { + return err + } var inserted int err := tx.QueryRow(` WITH requested(company_id, position) AS ( @@ -268,6 +275,32 @@ SELECT count(*) FROM inserted`, taskID, companyIDs).Scan(&inserted) if inserted != len(companyIDs) { return fmt.Errorf("%w: one or more companies do not exist", ErrTaskCompanyNotFound) } + + // Updating task_ids fires sync_task_asset_links, which first creates generic + // source rows. The provenance upsert must therefore run afterwards so the + // company name and creation-time reason remain visible to operators. + if _, err := tx.Exec(` +UPDATE assets +SET task_ids=CASE + WHEN $1=ANY(task_ids) THEN task_ids + ELSE array_append(task_ids, $1) +END +WHERE company_id=ANY($2::bigint[])`, taskID, companyIDs); err != nil { + return err + } + if _, err := tx.Exec(` +INSERT INTO task_asset_links(task_id, asset_id, source, source_summary) +SELECT $1, asset.id, $3, '任务创建时关联企业：' || company.name +FROM assets asset +JOIN companies company ON company.id=asset.company_id +WHERE asset.company_id=ANY($2::bigint[]) + AND $1=ANY(asset.task_ids) +ON CONFLICT (task_id, asset_id) DO UPDATE +SET source=EXCLUDED.source, + source_summary=EXCLUDED.source_summary, + source_node_id=NULL`, taskID, companyIDs, taskCompanyAssetSource); err != nil { + return err + } return nil } diff --git a/db/tasks_test.go b/db/tasks_test.go index ef5dc26..ae63e01 100644 --- a/db/tasks_test.go +++ b/db/tasks_test.go @@ -509,10 +509,36 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) { if err != nil { t.Fatal(err) } + emptyCompany, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Empty Company %d", suffix), "") + if err != nil { + t.Fatal(err) + } t.Cleanup(func() { - _, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB) + _, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2,$3)`, companyA, companyB, emptyCompany) }) + existingTask, err := d.CreateTask("existing company asset owner", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(existingTask.ID) }) + var companyAssetA, companyAssetB int64 + domainA := fmt.Sprintf("task-company-a-%d.example.test", suffix) + domainB := fmt.Sprintf("task-company-b-%d.example.test", suffix) + if err := d.QueryRow(` +INSERT INTO assets(type, domain, root_domain, company_id, company_source, task_ids) +VALUES ('root_domain',$1,$1,$2,'explicit',ARRAY[$3]::bigint[]) +RETURNING id`, domainA, companyA, existingTask.ID).Scan(&companyAssetA); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(` +INSERT INTO assets(type, domain, root_domain, company_id, company_source) +VALUES ('root_domain',$1,$1,$2,'explicit') +RETURNING id`, domainB, companyB).Scan(&companyAssetB); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{companyAssetA, companyAssetB}) }) + task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{ CompanyIDs: []int64{companyA, companyB}, }) @@ -537,6 +563,57 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) { if scopeCount != 2 { t.Fatalf("company task scope rows=%d want 2", scopeCount) } + assets, err := d.Assets().QueryByTask(task.ID, "root_domain", 10, 0) + if err != nil { + t.Fatal(err) + } + if len(assets) != 2 { + t.Fatalf("company task assets=%d want 2: %+v", len(assets), assets) + } + assetByID := make(map[int64]*Asset, len(assets)) + for _, asset := range assets { + assetByID[asset.ID] = asset + } + for assetID, companyName := range map[int64]string{ + companyAssetA: fmt.Sprintf("Task Company A %d", suffix), + companyAssetB: fmt.Sprintf("Task Company B %d", suffix), + } { + asset := assetByID[assetID] + if asset == nil { + t.Errorf("company asset %d missing from task", assetID) + continue + } + if asset.TaskSource != taskCompanyAssetSource || asset.TaskSourceSummary != "任务创建时关联企业："+companyName { + t.Errorf("asset %d provenance=%q/%q", assetID, asset.TaskSource, asset.TaskSourceSummary) + } + } + var existingAssociation bool + if err := d.QueryRow(`SELECT $1=ANY(task_ids) FROM assets WHERE id=$2`, existingTask.ID, companyAssetA).Scan(&existingAssociation); err != nil { + t.Fatal(err) + } + if !existingAssociation { + t.Fatal("company snapshot removed the asset's existing task association") + } + var existingSource string + if err := d.QueryRow(`SELECT source FROM task_asset_links WHERE task_id=$1 AND asset_id=$2`, existingTask.ID, companyAssetA).Scan(&existingSource); err != nil { + t.Fatal(err) + } + if existingSource == taskCompanyAssetSource { + t.Fatalf("company snapshot overwrote another task's provenance: %q", existingSource) + } + + emptyTask, err := d.CreateTaskWithOptions("empty company scope", "no current assets", TaskCreateOptions{ + CompanyIDs: []int64{emptyCompany}, + }) + if err != nil { + t.Fatal(err) + } + if assets, err := d.Assets().QueryByTask(emptyTask.ID, "", 10, 0); err != nil || len(assets) != 0 { + t.Fatalf("empty company task assets=%+v err=%v", assets, err) + } + if err := d.DeleteTask(emptyTask.ID); err != nil { + t.Fatal(err) + } duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{ CompanyIDs: []int64{companyA, companyA, companyB, companyA}, diff --git a/server/server.go b/server/server.go index 999fcb6..266c754 100644 --- a/server/server.go +++ b/server/server.go @@ -1363,7 +1363,7 @@ type createTaskReq struct { LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效 SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务 - CompanyIDs []int64 `json:"company_ids,omitempty"` // 关联企业资产范围;不复制资产或强制生成意图 + CompanyIDs []int64 `json:"company_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 轮)。 diff --git a/web/src/app/(main)/function/tasks/page.tsx b/web/src/app/(main)/function/tasks/page.tsx index f7050e2..1996f79 100644 --- a/web/src/app/(main)/function/tasks/page.tsx +++ b/web/src/app/(main)/function/tasks/page.tsx @@ -2613,8 +2613,8 @@ function CreateTaskSheet({ portalContainer={sheetContentRef} /> - 将所选企业的域名、IP、CIDR、ICP 和企业关键词作为 Agent - 可选范围上下文；不会自动生成意图或强制改变执行目标。 + 创建任务时会将所选企业当前已有资产加入“测试资产”，并将域名、IP、CIDR、ICP 和企业关键词提供给 Agent + 作为范围上下文；不会自动生成意图或强制改变执行目标。 diff --git a/web/src/app/(main)/function/traffic/page.tsx b/web/src/app/(main)/function/traffic/page.tsx index b7623ff..b565110 100644 --- a/web/src/app/(main)/function/traffic/page.tsx +++ b/web/src/app/(main)/function/traffic/page.tsx @@ -39,7 +39,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { api } from "@/lib/api"; import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types"; -import { cn } from "@/lib/utils"; +import { cn, copyText } from "@/lib/utils"; function fmtTime(ts: string) { return new Date(ts).toLocaleString("zh-CN", { @@ -199,12 +199,12 @@ function HttpCodeBlock({ raw }: { raw: string }) { }, [copied]); const copyPacket = async () => { - try { - await navigator.clipboard.writeText(value); + const ok = await copyText(value); + if (ok) { setCopied(true); - } catch { - toast.error("复制失败，请使用 Ctrl/Cmd+A 后复制"); + return; } + toast.error("复制失败，请使用 Ctrl/Cmd+A 后复制"); }; const renderLine = (line: string, index: number) => { diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 6e9b4fd..71ccc50 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -42,9 +42,35 @@ const mockCompanies = structuredClone(D.companies); const mockAssets = structuredClone(D.assets); const mockTaskAssetIDs = new Map(D.tasks.map((task, index) => [task.id, index + 1])); const mockTaskScopes = new Map(); +type MockTaskAssetSource = Pick; +const mockTaskAssetSources = new Map(); let nextMockTaskAssetID = D.tasks.length + 1; let mockActiveTask = D.ACTIVE_TASK; +function mockTaskAssetSourceKey(taskID: string, assetID: number): string { + return JSON.stringify([taskID, assetID]); +} + +function setMockTaskAssetSource(taskID: string, assetID: number, source: MockTaskAssetSource) { + mockTaskAssetSources.set(mockTaskAssetSourceKey(taskID, assetID), source); +} + +function mockAssetForTask(taskID: string, asset: Asset): Asset { + const source = mockTaskAssetSources.get(mockTaskAssetSourceKey(taskID, asset.id)); + return source ? { ...asset, ...source } : asset; +} + +function deleteMockTaskAssetSources(taskID: string, assetID?: number) { + if (assetID !== undefined) { + mockTaskAssetSources.delete(mockTaskAssetSourceKey(taskID, assetID)); + return; + } + for (const key of mockTaskAssetSources.keys()) { + const [linkedTaskID] = JSON.parse(key) as [string, number]; + if (linkedTaskID === taskID) mockTaskAssetSources.delete(key); + } +} + function mockTaskAssetID(taskID: string): number | undefined { const numeric = Number(taskID); if (Number.isInteger(numeric) && numeric > 0) return numeric; @@ -289,7 +315,10 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re const now = new Date(); const profileIDs = [...((b.llm_profile_ids as number[] | undefined) ?? [])]; const sourceTaskIDs = [...((b.source_task_ids as string[] | undefined) ?? [])]; - const companyIDs = [...((b.company_ids as number[] | undefined) ?? [])]; + const companyIDs = [...new Set((b.company_ids as number[] | undefined) ?? [])]; + if (companyIDs.some((companyID) => !mockCompanies.some((company) => company.id === companyID))) { + throw new Error("关联企业不存在或无效"); + } const categoryID = typeof b.category_id === "number" ? b.category_id : undefined; const category = categoryID === undefined ? undefined : mockTaskCategories.find((item) => item.id === categoryID); if (categoryID !== undefined && !category) throw new Error("任务分类不存在"); @@ -318,7 +347,19 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re source_task_ids: sourceTaskIDs, company_ids: companyIDs, }; - mockTaskAssetIDs.set(id, nextMockTaskAssetID++); + const numericTaskID = nextMockTaskAssetID++; + mockTaskAssetIDs.set(id, numericTaskID); + const selectedCompanies = new Set(companyIDs); + for (const asset of mockAssets) { + if (asset.company_id === undefined || !selectedCompanies.has(asset.company_id)) continue; + if (!asset.task_ids.includes(numericTaskID)) asset.task_ids.push(numericTaskID); + const company = mockCompanies.find((candidate) => candidate.id === asset.company_id); + setMockTaskAssetSource(id, asset.id, { + task_source: "company", + task_source_summary: `任务创建时关联企业：${company?.name ?? `#${asset.company_id}`}`, + task_source_node_id: undefined, + }); + } for (const item of mockTasks) item.active = false; mockTasks.unshift(created); mockActiveTask = id; @@ -462,9 +503,14 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re } if (seg[0] === "tasks" && seg.length === 2 && m === "DELETE") { const id = seg[1]; + const numericTaskID = mockTaskAssetID(id); const index = mockTasks.findIndex((item) => item.id === id); if (index >= 0) mockTasks.splice(index, 1); mockTaskAssetIDs.delete(id); + deleteMockTaskAssetSources(id); + if (numericTaskID !== undefined) { + for (const asset of mockAssets) asset.task_ids = asset.task_ids.filter((taskID) => taskID !== numericTaskID); + } let findingsDeleted = 0; if (b.delete_findings) { @@ -701,7 +747,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re }); const limit = Number(q.get("limit") ?? 50); const offset = Number(q.get("offset") ?? 0); - const page = list.slice(offset, offset + limit); + const page = list.slice(offset, offset + limit).map((asset) => (taskID ? mockAssetForTask(taskID, asset) : asset)); return { count: page.length, total: list.length, assets: page }; } if (path === "/assets" && m === "DELETE") { @@ -792,9 +838,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re asset.task_ids.push(numericTaskID); mutation.assets_linked++; } - asset.task_source = "manual"; - asset.task_source_summary = "用户在测试资产页手工新增"; - asset.task_source_node_id = undefined; + setMockTaskAssetSource(seg[1], asset.id, { + task_source: "manual", + task_source_summary: "用户在测试资产页手工新增", + task_source_node_id: undefined, + }); } mockTaskScopes.set(seg[1], currentScopes); return mutation; @@ -812,9 +860,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re asset.task_ids.push(numericTaskID); mutation.attached++; } - asset.task_source = "manual"; - asset.task_source_summary = sourceSummary; - asset.task_source_node_id = undefined; + setMockTaskAssetSource(seg[1], asset.id, { + task_source: "manual", + task_source_summary: sourceSummary, + task_source_node_id: undefined, + }); } return mutation; } @@ -824,11 +874,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (numericTaskID === undefined || !asset) throw new Error("任务或资产不存在"); if (!asset.task_ids.includes(numericTaskID)) throw new Error("资产未关联当前任务"); asset.task_ids = asset.task_ids.filter((id) => id !== numericTaskID); - if (asset.task_ids.length === 0) { - asset.task_source = undefined; - asset.task_source_summary = undefined; - asset.task_source_node_id = undefined; - } + deleteMockTaskAssetSources(seg[1], asset.id); return { detached: asset.id }; } if (seg[0] === "tasks" && seg[2] === "intent-assets" && seg.length === 3 && m === "GET") { @@ -843,8 +889,9 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re } const sourceTaskID = mockTaskAssetID(seg[1]) ?? 0; const assets: IntentAsset[] = mappings.flatMap((mapping) => { - const asset = mockAssets.find((item) => item.id === mapping.assetID); - if (!asset) return []; + const storedAsset = mockAssets.find((item) => item.id === mapping.assetID); + if (!storedAsset) return []; + const asset = mockAssetForTask(seg[1], storedAsset); return [ { intent_id: mapping.intentID, diff --git a/web/src/lib/task-assets.ts b/web/src/lib/task-assets.ts index 1bfccc7..1981c15 100644 --- a/web/src/lib/task-assets.ts +++ b/web/src/lib/task-assets.ts @@ -13,6 +13,7 @@ const TASK_ASSET_SOURCE_LABELS: Record = { agent: "Agent 发现", anchor: "黑板锚点", api: "资产 API", + company: "企业关联", legacy: "历史关联", manual: "人工加入", system: "系统关联", diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 59c4b3c..8b0208d 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -35,7 +35,7 @@ export interface Task { llm_failover_state?: "default" | "ready" | "chain_exhausted" | string; llm_failover_reason?: string; source_task_ids?: string[]; // directly related tasks inherited as read-only context - company_ids?: number[]; // associated company asset scopes available as task context + company_ids?: number[]; // associated company scopes; current company assets join the task at creation coverage_enabled?: boolean; // 资产覆盖度功能开关(创建时定,默认开)；false=不计算/不展示覆盖度 } From 594436fddd72e87da93fcdd1a9939f45c40dc45c Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 08:00:36 +0800 Subject: [PATCH 2/7] =?UTF-8?q?=E4=BF=AE=E5=A4=8DBUG=EF=BC=9A=201=E3=80=81?= =?UTF-8?q?=E5=BD=93=E5=AD=98=E5=9C=A8=E8=BE=83=E5=A4=9A=E4=BB=BB=E5=8A=A1?= =?UTF-8?q?=E6=97=B6=EF=BC=8C=E7=B3=BB=E7=BB=9F=E4=BC=9A=E5=8F=91=E7=94=9F?= =?UTF-8?q?=E5=8D=A1=E9=A1=BF=EF=BC=8C=E5=8A=A0=E8=BD=BD=E5=AF=B9=E8=AF=9D?= =?UTF-8?q?=E6=88=96=E8=80=85=E5=8A=A0=E8=BD=BD=E7=95=8C=E9=9D=A2=E6=97=B6?= =?UTF-8?q?=E5=80=99=E5=AD=98=E5=9C=A8=E6=98=8E=E6=98=BE=E5=8D=A1=E9=A1=BF?= =?UTF-8?q?=EF=BC=8C=E4=BC=98=E5=8C=96=E6=80=A7=E8=83=BD=E6=89=BE=E5=87=BA?= =?UTF-8?q?=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- db/exploration.go | 68 +++++++ db/schema.sql | 7 + db/task_context.go | 166 +++++++++++++++++- db/task_list_performance_test.go | 117 ++++++++++++ db/tasks.go | 6 +- server/server.go | 14 +- web/src/app/(main)/chat/page.tsx | 161 ++++++++++------- web/src/app/(main)/dashboard/page.tsx | 62 +++++-- .../tasks/detail/_tabs/overview-tab.tsx | 27 ++- .../app/(main)/function/tasks/detail/page.tsx | 74 ++++---- web/src/lib/api.ts | 9 +- web/src/lib/mock/handler.ts | 5 + 12 files changed, 576 insertions(+), 140 deletions(-) create mode 100644 db/task_list_performance_test.go diff --git a/db/exploration.go b/db/exploration.go index 3a8c563..f62e280 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -991,6 +991,74 @@ func (d *DB) LastActivityAll() (map[int64]int64, error) { // GoalCounts is the goal summary for one exploration. type GoalCounts struct{ Total, Met int } +// TaskListMetrics contains the aggregates rendered in task lists. +type TaskListMetrics struct { + Tokens TokenUsage + LastActivity int64 + Goals GoalCounts +} + +// TaskListMetricsAll returns list aggregates for every live task in one query. +// The lateral lookups use the per-exploration indexes instead of grouping the +// complete activity history on every poll. +func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) { + rows, err := d.Query(` + SELECT task.exploration_id, + COALESCE(token_metrics.input_tokens,0), + COALESCE(token_metrics.output_tokens,0), + COALESCE(token_metrics.cache_read_tokens,0), + COALESCE(token_metrics.cache_write_tokens,0), + COALESCE(latest_activity.created_at,0), + COALESCE(goal_metrics.total,0), + COALESCE(goal_metrics.met,0) + FROM tasks task + LEFT JOIN LATERAL ( + SELECT SUM(input_tokens) AS input_tokens, + SUM(output_tokens) AS output_tokens, + SUM(cache_read_tokens) AS cache_read_tokens, + SUM(cache_write_tokens) AS cache_write_tokens + FROM activity + WHERE exploration_id=task.exploration_id AND kind='result' + ) token_metrics ON true + LEFT JOIN LATERAL ( + SELECT EXTRACT(EPOCH FROM created_at)::bigint AS created_at + FROM activity + WHERE exploration_id=task.exploration_id + ORDER BY created_at DESC + LIMIT 1 + ) latest_activity ON true + LEFT JOIN LATERAL ( + SELECT COUNT(*) AS total, + COUNT(*) FILTER (WHERE state='met') AS met + FROM exploration_nodes + WHERE exploration_id=task.exploration_id AND kind='goal' + ) goal_metrics ON true + WHERE task.deleted_at IS NULL`) + if err != nil { + return nil, err + } + defer rows.Close() + out := map[int64]TaskListMetrics{} + for rows.Next() { + var explorationID int64 + var metrics TaskListMetrics + if err := rows.Scan( + &explorationID, + &metrics.Tokens.InputTokens, + &metrics.Tokens.OutputTokens, + &metrics.Tokens.CacheReadTokens, + &metrics.Tokens.CacheWriteTokens, + &metrics.LastActivity, + &metrics.Goals.Total, + &metrics.Goals.Met, + ); err != nil { + return nil, err + } + out[explorationID] = metrics + } + return out, rows.Err() +} + // GoalCountsAll returns goal totals (total / met) per exploration id, one query // for all tasks — used by listTasks so the task list shows progress for every task. func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) { diff --git a/db/schema.sql b/db/schema.sql index b1a53fe..6d962d9 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -271,6 +271,13 @@ 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 -- can't distinguish them); this covers reverse pagination of those sessions. CREATE INDEX IF NOT EXISTS idx_act_worker ON activity(exploration_id, worker, id); +-- Task-list polls aggregate result usage and find the latest event repeatedly. +-- Cover the token columns for index-only aggregation and the timestamp order for +-- per-exploration latest-activity lookups. +CREATE INDEX IF NOT EXISTS idx_act_result_usage ON activity(exploration_id) + INCLUDE (input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) + WHERE kind='result'; +CREATE INDEX IF NOT EXISTS idx_act_latest ON activity(exploration_id, created_at DESC); -- ===================================================================== -- C. LLM profiles diff --git a/db/task_context.go b/db/task_context.go index 14053d4..cf23c7c 100644 --- a/db/task_context.go +++ b/db/task_context.go @@ -52,9 +52,6 @@ func (d *DB) hydrateTaskContext(t *Task) error { 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 @@ -65,8 +62,22 @@ func (d *DB) hydrateTaskContext(t *Task) error { return err } t.CompanyIDs = companies + applyTaskLLMContext(t, legacyProfileID, activeProfileID, revision, chain) + return nil +} + +func applyTaskLLMContext( + t *Task, + legacyProfileID, activeProfileID *int64, + revision int64, + chain []TaskLLMProfile, +) { + t.LLMProfileID = legacyProfileID + t.ActiveLLMProfileID = activeProfileID + t.LLMChainRevision = revision t.LLMProfileIDs = make([]int64, 0, len(chain)) t.LLMFailoverState = "default" + t.LLMFailoverReason = "" var latest *TaskLLMProfile activeReady := false for i := range chain { @@ -90,6 +101,155 @@ func (d *DB) hydrateTaskContext(t *Task) error { if latest != nil { t.LLMFailoverReason = latest.LastError } +} + +type taskBatchContext struct { + legacyProfileID *int64 + activeProfileID *int64 + revision int64 + chain []TaskLLMProfile + sourceIDs []int64 + companyIDs []int64 +} + +// hydrateTasksContext loads every task's LLM chain, source tasks, and company +// scopes with three bulk queries. ListTasks used to issue these queries once per +// task, making startup and task-list hydration grow as 3N+1 database round trips. +func (d *DB) hydrateTasksContext(tasks []*Task) error { + if len(tasks) == 0 { + return nil + } + + ids := make([]int64, 0, len(tasks)) + contexts := make(map[int64]*taskBatchContext, len(tasks)) + byID := make(map[int64]*Task, len(tasks)) + for _, task := range tasks { + if task == nil { + continue + } + ids = append(ids, task.ID) + contexts[task.ID] = &taskBatchContext{} + byID[task.ID] = task + } + if len(ids) == 0 { + return nil + } + + rows, err := d.Query(` +SELECT t.id, 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=ANY($1::bigint[]) +ORDER BY t.id, p.position NULLS LAST`, ids) + if err != nil { + return err + } + for rows.Next() { + var taskID, revision int64 + var legacy, active, profileID, position sql.NullInt64 + var status, lastError sql.NullString + var exhaustedAt sql.NullTime + if err := rows.Scan( + &taskID, &legacy, &active, &revision, + &profileID, &position, &status, &lastError, &exhaustedAt, + ); err != nil { + rows.Close() + return err + } + context := contexts[taskID] + if context == nil { + continue + } + context.revision = revision + if legacy.Valid { + id := legacy.Int64 + context.legacyProfileID = &id + } + if active.Valid { + id := active.Int64 + context.activeProfileID = &id + } + if profileID.Valid { + entry := TaskLLMProfile{ + ProfileID: profileID.Int64, + Position: int(position.Int64), + Status: status.String, + LastError: lastError.String, + } + if exhaustedAt.Valid { + ts := exhaustedAt.Time + entry.ExhaustedAt = &ts + } + context.chain = append(context.chain, entry) + } + } + if err := rows.Err(); err != nil { + rows.Close() + return err + } + rows.Close() + + rows, err = d.Query(` +SELECT task_id, source_task_id +FROM task_relations +WHERE task_id=ANY($1::bigint[]) +ORDER BY task_id, created_at, source_task_id`, ids) + if err != nil { + return err + } + for rows.Next() { + var taskID, sourceID int64 + if err := rows.Scan(&taskID, &sourceID); err != nil { + rows.Close() + return err + } + if context := contexts[taskID]; context != nil { + context.sourceIDs = append(context.sourceIDs, sourceID) + } + } + if err := rows.Err(); err != nil { + rows.Close() + return err + } + rows.Close() + + rows, err = d.Query(` +SELECT task_id, company_id +FROM task_scope +WHERE task_id=ANY($1::bigint[]) AND kind='company' AND company_id IS NOT NULL +ORDER BY task_id, id`, ids) + if err != nil { + return err + } + for rows.Next() { + var taskID, companyID int64 + if err := rows.Scan(&taskID, &companyID); err != nil { + rows.Close() + return err + } + if context := contexts[taskID]; context != nil { + context.companyIDs = append(context.companyIDs, companyID) + } + } + if err := rows.Err(); err != nil { + rows.Close() + return err + } + rows.Close() + + for id, context := range contexts { + task := byID[id] + task.SourceTaskIDs = context.sourceIDs + task.CompanyIDs = context.companyIDs + applyTaskLLMContext( + task, + context.legacyProfileID, + context.activeProfileID, + context.revision, + context.chain, + ) + } return nil } diff --git a/db/task_list_performance_test.go b/db/task_list_performance_test.go new file mode 100644 index 0000000..104bedf --- /dev/null +++ b/db/task_list_performance_test.go @@ -0,0 +1,117 @@ +package db + +import ( + "fmt" + "testing" + "time" +) + +func TestListTasksBulkHydratesTaskContext(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + stamp := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("bulk-list-profile-%d", stamp), Format: "openai", Model: "test", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteProfile(profileID) }) + source, err := d.CreateTask(fmt.Sprintf("bulk-list-source-%d", stamp), "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(source.ID) }) + companyID, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Bulk List Company %d", stamp), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.Companies().DeleteCompany(companyID) }) + child, err := d.CreateTaskWithOptions("bulk-list-child", "goal", TaskCreateOptions{ + LLMProfileIDs: []int64{profileID}, + SourceTaskIDs: []int64{source.ID}, + CompanyIDs: []int64{companyID}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(child.ID) }) + + tasks, err := d.ListTasks() + if err != nil { + t.Fatal(err) + } + var got *Task + for _, task := range tasks { + if task.ID == child.ID { + got = task + break + } + } + if got == nil { + t.Fatalf("task %d missing from list", child.ID) + } + if len(got.LLMProfileIDs) != 1 || got.LLMProfileIDs[0] != profileID || + got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileID || got.LLMFailoverState != "ready" { + t.Fatalf("LLM context was not bulk hydrated: %+v", got) + } + if len(got.SourceTaskIDs) != 1 || got.SourceTaskIDs[0] != source.ID { + t.Fatalf("source task context was not bulk hydrated: %v", got.SourceTaskIDs) + } + if len(got.CompanyIDs) != 1 || got.CompanyIDs[0] != companyID { + t.Fatalf("company context was not bulk hydrated: %v", got.CompanyIDs) + } +} + +func TestTaskListMetricsAll(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + task, err := d.CreateTask("task-list-metrics", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + + resultAt := time.Now().Add(-time.Second).Truncate(time.Second) + latestAt := resultAt.Add(time.Second) + if _, err := d.Exec(` +INSERT INTO activity( + exploration_id, kind, input_tokens, output_tokens, + cache_read_tokens, cache_write_tokens, created_at +) VALUES ($1,'result',11,7,3,2,$2), ($1,'tool_result',999,999,999,999,$3)`, + task.ExplorationID, resultAt, latestAt); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(` +INSERT INTO exploration_nodes(exploration_id, kind, payload, state) +VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExplorationID); err != nil { + t.Fatal(err) + } + + all, err := d.TaskListMetricsAll() + if err != nil { + t.Fatal(err) + } + metrics, ok := all[task.ExplorationID] + if !ok { + t.Fatalf("metrics for exploration %d missing", task.ExplorationID) + } + if metrics.Tokens.InputTokens != 11 || metrics.Tokens.OutputTokens != 7 || + metrics.Tokens.CacheReadTokens != 3 || metrics.Tokens.CacheWriteTokens != 2 { + t.Fatalf("unexpected token metrics: %+v", metrics.Tokens) + } + if metrics.LastActivity != latestAt.Unix() { + t.Fatalf("last activity=%d, want %d", metrics.LastActivity, latestAt.Unix()) + } + if metrics.Goals.Total != 2 || metrics.Goals.Met != 1 { + t.Fatalf("unexpected goal metrics: %+v", metrics.Goals) + } +} diff --git a/db/tasks.go b/db/tasks.go index f658305..69a1a34 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -377,10 +377,8 @@ ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, id DESC`) return nil, err } rows.Close() - for _, task := range out { - if err := d.hydrateTaskContext(task); err != nil { - return nil, err - } + if err := d.hydrateTasksContext(out); err != nil { + return nil, err } return out, nil } diff --git a/server/server.go b/server/server.go index 266c754..e46755a 100644 --- a/server/server.go +++ b/server/server.go @@ -995,22 +995,20 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { active = t.ID } list := s.m.List() - toks, _ := s.m.PG().TokenTotalsAll() // whole-task token totals, one query for all tasks - lastAct, _ := s.m.PG().LastActivityAll() // persisted last-activity per task, one query - goalCounts, _ := s.m.PG().GoalCountsAll() // goal progress per exploration, one query + metrics, _ := s.m.PG().TaskListMetricsAll() dtos := make([]TaskDTO, 0, len(list)) for _, t := range list { dto := taskDTO(t, s.resolvedTaskStatus(t)) - dto.Tokens = tokenTotalDTO(toks[t.ExpID]) + metric := metrics[t.ExpID] + dto.Tokens = tokenTotalDTO(metric.Tokens) // prefer the live in-memory heartbeat (fresher) and fall back to the // persisted max activity time (survives restarts) for run-duration display. - dto.LastActivity = lastAct[t.ExpID] + dto.LastActivity = metric.LastActivity if live := s.engine.LastActivity(t.ID); live > dto.LastActivity { dto.LastActivity = live } - gc := goalCounts[t.ExpID] - dto.GoalsTotal = gc.Total - dto.GoalsMet = gc.Met + dto.GoalsTotal = metric.Goals.Total + dto.GoalsMet = metric.Goals.Met dtos = append(dtos, dto) } writeJSON(w, 200, map[string]any{"tasks": dtos, "active": active}) diff --git a/web/src/app/(main)/chat/page.tsx b/web/src/app/(main)/chat/page.tsx index 04a4616..57f6f0d 100644 --- a/web/src/app/(main)/chat/page.tsx +++ b/web/src/app/(main)/chat/page.tsx @@ -80,6 +80,7 @@ function fmtDuration(ms: number): string { // then one more page each time the user scrolls to the top. Kept modest so a long // thread stays snappy (only ~a page of rows is in the DOM until you scroll up). const HISTORY_PAGE = 200; +const CONVERSATION_LIST_PAGE = 100; function conversationIsPinned(conversation: Conversation): boolean { return conversation.pinned ?? Boolean(conversation.pinned_at); @@ -749,7 +750,7 @@ function ChatView({ // ConversationItem is one row in the left rail: title, agent subtitle, inline // rename, pin marker, and a compact action menu. -function ConversationItem({ +const ConversationItem = React.memo(function ConversationItem({ conv, agent, active, @@ -770,15 +771,15 @@ function ConversationItem({ active: boolean; renaming: boolean; renameText: string; - onSelect: () => void; - onStartRename: () => void; + onSelect: (id: number) => void; + onStartRename: (conversation: Conversation) => void; onRenameText: (v: string) => void; - onCommitRename: () => void; + onCommitRename: (id: number, title: string) => void; onCancelRename: () => void; - onTogglePinned: () => void; - onDelete: () => void; + onTogglePinned: (conversation: Conversation) => void; + onDelete: (id: number) => void; selectedForDelete: boolean; - onSelectedForDeleteChange: (checked: boolean) => void; + onSelectedForDeleteChange: (id: number, checked: boolean) => void; }) { const [deleteOpen, setDeleteOpen] = React.useState(false); const renameInputRef = React.useRef(null); @@ -801,7 +802,7 @@ function ConversationItem({ > onSelectedForDeleteChange(checked === true)} + onCheckedChange={(checked) => onSelectedForDeleteChange(conv.id, checked === true)} aria-label={`选择对话「${conv.title || "新对话"}」`} className="ml-1 shrink-0" /> @@ -811,10 +812,10 @@ function ConversationItem({ value={renameText} onChange={(e) => onRenameText(e.target.value)} onBlur={() => { - if (!cancelRenameRef.current) onCommitRename(); + if (!cancelRenameRef.current) onCommitRename(conv.id, renameText); }} onKeyDown={(e) => { - if (e.key === "Enter") onCommitRename(); + if (e.key === "Enter") onCommitRename(conv.id, renameText); if (e.key === "Escape") { e.preventDefault(); cancelRenameRef.current = true; @@ -826,8 +827,8 @@ function ConversationItem({ ) : ( + )} diff --git a/web/src/app/(main)/dashboard/page.tsx b/web/src/app/(main)/dashboard/page.tsx index df1b5b3..ec79ed8 100644 --- a/web/src/app/(main)/dashboard/page.tsx +++ b/web/src/app/(main)/dashboard/page.tsx @@ -148,13 +148,46 @@ export default function DashboardPage() { const [tools, setTools] = React.useState([]); const [llmProfiles, setLLMProfiles] = React.useState([]); - // fast poll: tasks, findings, stats, pending, activity (every 5s) + // The task list is the most expensive dashboard source. Poll it independently + // so a large history cannot hold back every other dashboard panel. React.useEffect(() => { let alive = true; + let loading = false; + let signature = ""; const load = async () => { + if (loading) return; + loading = true; try { - const [tr, fr, sr, sets, pr, act, tok, ctok] = await Promise.all([ - api.tasks(), + const tr = await api.tasks(); + if (!alive) return; + const nextSignature = JSON.stringify(tr.tasks); + if (nextSignature !== signature) { + signature = nextSignature; + setTasks(tr.tasks); + } + } catch { + /* transient errors — next poll retries */ + } finally { + loading = false; + } + }; + void load(); + const t = setInterval(load, 5000); + return () => { + alive = false; + clearInterval(t); + }; + }, []); + + // The remaining fast sources stay batched so one poll causes a single render. + React.useEffect(() => { + let alive = true; + let loading = false; + const load = async () => { + if (loading) return; + loading = true; + try { + const [findings, stats, settings, pending, activity, tokens, conversationTokens] = await Promise.all([ api.findings(), api.stats(), api.settings(), @@ -164,23 +197,24 @@ export default function DashboardPage() { api.conversationTokens(), ]); if (!alive) return; - setTasks(tr.tasks); - setFindings(fr); - setStats(sr); - setSettings(sets); - setPending(pr); - setActivity(act.items); - setTokens(tok.total ?? null); - setConvTokens(ctok); + setFindings(findings); + setStats(stats); + setSettings(settings); + setPending(pending); + setActivity(activity.items); + setTokens(tokens.total ?? null); + setConvTokens(conversationTokens); } catch { - /* transient errors — next poll retries */ + // Preserve stale data and retry on the next interval. + } finally { + loading = false; } }; void load(); - const t = setInterval(load, 5000); + const timer = setInterval(load, 5000); return () => { alive = false; - clearInterval(t); + clearInterval(timer); }; }, []); diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx index 37810d9..d7babef 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx @@ -56,7 +56,7 @@ function scopeValue(row: TaskScopeRow): string { if (row.value) return row.value; if (row.domain) return row.domain; if (row.net) return row.net; - if (row.company_id) return row.company_name || `企业 #${row.company_id}`; + if (row.company_id) return row.company_name?.trim() ? row.company_name : `企业 #${row.company_id}`; return "—"; } @@ -346,17 +346,32 @@ export function OverviewTab({ taskId }: { taskId: string }) { React.useEffect(() => { let cancelled = false; + let loading = false; const load = async () => { + if (loading) return; + loading = true; try { - const [tasksResp, statsResp, intentsResp, findingsResp] = await Promise.all([ - api.tasks(), + const [taskResp, statsResp, intentsResp, findingsResp] = await Promise.all([ + api.task(taskId), api.stats(taskId), api.intents(taskId), api.findings(taskId), ]); if (cancelled) return; - setTask(tasksResp.tasks.find((t) => t.id === taskId) ?? null); + const activeTask = statsResp.active_task; + setTask( + activeTask + ? { + ...taskResp, + in_flight: activeTask.in_flight, + goals_total: activeTask.goals_total, + goals_met: activeTask.goals_met, + engine_mode: statsResp.engine_mode ?? activeTask.engine_mode, + paused: activeTask.paused, + } + : taskResp, + ); setStats(statsResp); setIntents(intentsResp); setFindings(findingsResp); @@ -372,10 +387,12 @@ export function OverviewTab({ taskId }: { taskId: string }) { }); } catch { // transient errors are ignored; the next poll will retry + } finally { + loading = false; } }; - load(); + void load(); void loadScope(); void loadTokens(); void loadGoals(); diff --git a/web/src/app/(main)/function/tasks/detail/page.tsx b/web/src/app/(main)/function/tasks/detail/page.tsx index b07434a..f28a12d 100644 --- a/web/src/app/(main)/function/tasks/detail/page.tsx +++ b/web/src/app/(main)/function/tasks/detail/page.tsx @@ -1,22 +1,16 @@ "use client"; import * as React from "react"; + import Link from "next/link"; import { useSearchParams } from "next/navigation"; -import { toast } from "sonner"; -import { - ArrowLeftIcon, - PauseIcon, - PlayIcon, - BrainIcon, - CheckIcon, - CircleAlertIcon, -} from "lucide-react"; +import { ArrowLeftIcon, BrainIcon, CheckIcon, CircleAlertIcon, PauseIcon, PlayIcon } from "lucide-react"; +import { toast } from "sonner"; + +import { StatusBadge } from "@/components/status-badge"; import { TaskLLMProfileChain } from "@/components/task-llm-profile-chain"; import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; -import { SidebarTrigger } from "@/components/ui/sidebar"; -import { Separator } from "@/components/ui/separator"; import { Button } from "@/components/ui/button"; import { Popover, @@ -26,30 +20,31 @@ import { PopoverTitle, PopoverTrigger, } from "@/components/ui/popover"; +import { Separator } from "@/components/ui/separator"; +import { SidebarTrigger } from "@/components/ui/sidebar"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { StatusBadge } from "@/components/status-badge"; import { api } from "@/lib/api"; import type { LLMProfile, Task } from "@/lib/types"; -import { SessionsTab } from "./_tabs/sessions-tab"; -import { OverviewTab } from "./_tabs/overview-tab"; -import { GraphTab } from "./_tabs/graph-tab"; -import { FindingsTab } from "./_tabs/findings-tab"; -import { ReportTab } from "./_tabs/report-tab"; -import { InterceptTab } from "./_tabs/intercept-tab"; import { AssetsTab } from "./_tabs/assets-tab"; import { CoverageGraphTab } from "./_tabs/coverage-graph-tab"; +import { FindingsTab } from "./_tabs/findings-tab"; +import { GraphTab } from "./_tabs/graph-tab"; +import { InterceptTab } from "./_tabs/intercept-tab"; +import { OverviewTab } from "./_tabs/overview-tab"; +import { ReportTab } from "./_tabs/report-tab"; +import { SessionsTab } from "./_tabs/sessions-tab"; const TABS = [ - { value: "sessions", label: "会话" }, - { value: "overview", label: "总览" }, - { value: "graph", label: "探索链路" }, - { value: "findings", label: "发现" }, - { value: "assets", label: "测试资产" }, - { value: "coverage", label: "资产覆盖图" }, + { value: "sessions", label: "会话" }, + { value: "overview", label: "总览" }, + { value: "graph", label: "探索链路" }, + { value: "findings", label: "发现" }, + { value: "assets", label: "测试资产" }, + { value: "coverage", label: "资产覆盖图" }, { value: "intercept", label: "拦截审批" }, - { value: "report", label: "报告" }, + { value: "report", label: "报告" }, ]; function taskProfileIDs(task: Task): string[] { @@ -81,7 +76,7 @@ function TaskLLMControl({ task, profiles, onUpdated }: { task: Task; profiles: L const activeProfile = profiles.find((profile) => profile.id === activeID); const currentLabel = exhausted ? "配置链已耗尽" - : activeProfile?.name ?? (activeID ? `配置 #${activeID}` : "跟随默认配置"); + : (activeProfile?.name ?? (activeID ? `配置 #${activeID}` : "跟随默认配置")); const activeIndex = chain.indexOf(activeID); const backupCount = activeIndex >= 0 ? Math.max(0, chain.length - activeIndex - 1) : 0; const currentTitle = [currentLabel, activeProfile?.model, backupCount > 0 ? `${backupCount} 个备用` : ""] @@ -151,11 +146,7 @@ function TaskLLMControl({ task, profiles, onUpdated }: { task: Task; profiles: L {backupCount > 0 && +{backupCount}} - + 任务 LLM 配置链 {editorDescription} @@ -234,12 +225,16 @@ function TaskDetailInner() { }; }, [id]); + const taskLoadInFlight = React.useRef(null); const load = React.useCallback(() => { - Promise.all([api.tasks(), api.stats(id).catch(() => null)]) - .then(([r, s]) => { - const base = r.tasks.find((t) => t.id === id) ?? null; + if (taskLoadInFlight.current === id) return; + taskLoadInFlight.current = id; + Promise.all([api.task(id), api.stats(id).catch(() => null)]) + .then(([task, s]) => { + if (taskLoadInFlight.current !== id) return; + const base = { ...task }; const at = s?.active_task; - if (base && at) { + if (at) { base.in_flight = at.in_flight; base.goals_total = at.goals_total; base.goals_met = at.goals_met; @@ -247,12 +242,17 @@ function TaskDetailInner() { base.paused = at.paused; } setTask(base); - setPaused(at?.paused ?? base?.paused ?? false); + setPaused(at?.paused ?? base.paused ?? false); }) .catch(() => { // Keep the last rendered task state during a transient poll failure. }) - .finally(() => setLoaded(true)); + .finally(() => { + if (taskLoadInFlight.current === id) { + taskLoadInFlight.current = null; + setLoaded(true); + } + }); }, [id]); React.useEffect(() => { load(); diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index 9bb02e9..c28c328 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -188,6 +188,7 @@ export const api = { // ---- tasks ---- tasks: () => get<{ tasks: Task[]; active: string }>("/tasks").then((r) => ({ tasks: arr(r.tasks), active: r.active ?? "" })), + task: (id: string) => get(`/tasks/${encodeURIComponent(id)}`), createTask: (input: { name?: string; categoryId?: number; @@ -283,7 +284,8 @@ export const api = { usageStats: (days = 365) => get(`/tokens/usage?days=${days}`), // ---- 目标管理（总览）---- // 本任务全部目标（text/vulnclass/state）。 - taskGoals: (id: string) => get<{ goals: TaskGoal[] }>(`/tasks/${id}/goals`), + taskGoals: (id: string) => + get<{ goals: TaskGoal[] | null }>(`/tasks/${id}/goals`).then((response) => ({ goals: arr(response.goals) })), // 人工新增目标：写入图谱并通知 planner、复活任务。 addGoal: (id: string, text: string, vulnclass?: string) => post(`/tasks/${id}/goals`, { text, vulnclass: vulnclass ?? "" }), @@ -295,7 +297,10 @@ export const api = { // ---- 操作约束管理（总览）---- // 本任务全部操作约束（allow/deny）。 - taskConstraints: (id: string) => get<{ constraints: TaskConstraint[] }>(`/tasks/${id}/constraints`), + taskConstraints: (id: string) => + get<{ constraints: TaskConstraint[] | null }>(`/tasks/${id}/constraints`).then((response) => ({ + constraints: arr(response.constraints), + })), // 新增约束（不通知 planner，下一轮规划自然读到）。 addConstraint: (id: string, text: string, kind: TaskConstraint["kind"]) => post(`/tasks/${id}/constraints`, { text, kind }), diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 71ccc50..2b6f536 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -490,6 +490,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (index >= 0) mockTaskTemplates.splice(index, 1); return { deleted: id }; } + if (seg[0] === "tasks" && seg.length === 2 && m === "GET") { + const task = mockTasks.find((item) => item.id === seg[1]); + if (!task) throw new Error("任务不存在"); + return structuredClone(task); + } if (seg[0] === "tasks" && seg.length === 2 && m === "PATCH") { const task = mockTasks.find((item) => item.id === seg[1]); if (!task) throw new Error("任务不存在"); From 9da521ee4e0708f8d2add9b35701c52f7a1e9d4f Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 12:56:29 +0800 Subject: [PATCH 3/7] =?UTF-8?q?=E5=BD=92=E6=A1=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- db/findings.go | 49 +- db/findings_test.go | 36 + db/llm_usage.go | 83 +- db/schema.sql | 46 + db/skill_usage.go | 25 +- db/task_archive_aggregate_stats.go | 91 ++ db/task_archives.go | 667 ++++++++++ db/task_archives_restore.go | 648 ++++++++++ db/task_archives_test.go | 297 +++++ db/tool_usage.go | 14 +- server/dto.go | 14 + server/findings_groups.go | 1 + server/findings_groups_test.go | 22 + server/manager.go | 19 + server/server.go | 27 +- server/task_archive_package.go | 655 ++++++++++ server/task_archive_package_test.go | 209 ++++ server/task_archives.go | 656 ++++++++++ server/task_archives_test.go | 88 ++ server/task_llm_test.go | 17 + traffic/archive.go | 237 ++++ traffic/archive_test.go | 158 +++ traffic/traffic.go | 234 +++- web/src/app/(main)/function/findings/page.tsx | 85 +- .../app/(main)/function/tasks/detail/page.tsx | 72 +- web/src/app/(main)/function/tasks/page.tsx | 1089 +++++++++++++---- web/src/app/(main)/function/traffic/page.tsx | 57 +- web/src/components/transcript.tsx | 25 +- web/src/lib/api.ts | 27 + web/src/lib/mock/handler.ts | 317 ++++- web/src/lib/types.ts | 70 ++ 31 files changed, 5768 insertions(+), 267 deletions(-) create mode 100644 db/task_archive_aggregate_stats.go create mode 100644 db/task_archives.go create mode 100644 db/task_archives_restore.go create mode 100644 db/task_archives_test.go create mode 100644 server/task_archive_package.go create mode 100644 server/task_archive_package_test.go create mode 100644 server/task_archives.go create mode 100644 server/task_archives_test.go create mode 100644 traffic/archive.go create mode 100644 traffic/archive_test.go diff --git a/db/findings.go b/db/findings.go index b79f00c..bae1304 100644 --- a/db/findings.go +++ b/db/findings.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "sort" "strconv" "strings" "time" @@ -145,6 +146,7 @@ type FindingFilter struct { Status string // pending | false_positive | ignored | resolved VulnClass string TaskID string // 任务 id(字符串形式;空/非法 = 不按任务筛选) + Query string // 名称/类型/摘要/证据/报告正文的模糊检索关键词 Sort string // "severity" | "time" } @@ -154,7 +156,8 @@ type FindingFilter struct { const FindingUnassignedTask = "__unassigned__" // where builds the WHERE clause (shared by the page and count queries) plus its -// positional args. Only equality filters, all parameterized. +// positional args. All values are parameterized; Query also escapes ILIKE +// wildcards so user input is always matched literally. func (f FindingFilter) where() (string, []any) { var conds []string var args []any @@ -175,6 +178,18 @@ func (f FindingFilter) where() (string, []any) { args = append(args, tid) conds = append(conds, fmt.Sprintf("f.task_id = $%d", len(args))) } + if query := strings.TrimSpace(f.Query); query != "" { + escaped := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(query) + args = append(args, "%"+escaped+"%") + placeholder := fmt.Sprintf("$%d", len(args)) + conds = append(conds, fmt.Sprintf(`( + COALESCE(f.name, '') ILIKE %s ESCAPE '\' OR + f.vulnclass ILIKE %s ESCAPE '\' OR + f.summary ILIKE %s ESCAPE '\' OR + f.evidence ILIKE %s ESCAPE '\' OR + COALESCE(f.report, '') ILIKE %s ESCAPE '\' + )`, placeholder, placeholder, placeholder, placeholder, placeholder)) + } if len(conds) == 0 { return "", args } @@ -532,7 +547,37 @@ func (d *DB) FindingStats() (*FindingStats, error) { } st.Tasks = append(st.Tasks, opt) } - return st, trows.Err() + if err := trows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + vulnclasses := make(map[string]bool, len(st.VulnClasses)) + for _, vulnclass := range st.VulnClasses { + vulnclasses[vulnclass] = true + } + for _, aggregate := range archived { + cold := aggregate.FindingStats + st.Total += cold.Total + st.Pending += cold.Pending + st.Critical += cold.Critical + st.High += cold.High + st.Medium += cold.Medium + st.Low += cold.Low + for _, vulnclass := range cold.VulnClasses { + if vulnclass != "" { + vulnclasses[vulnclass] = true + } + } + } + st.VulnClasses = st.VulnClasses[:0] + for vulnclass := range vulnclasses { + st.VulnClasses = append(st.VulnClasses, vulnclass) + } + sort.Strings(st.VulnClasses) + return st, nil } // GetFinding returns a single finding row (with task_description joined and the diff --git a/db/findings_test.go b/db/findings_test.go index 4ed3ce2..8ec9347 100644 --- a/db/findings_test.go +++ b/db/findings_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "slices" + "strings" "testing" "time" ) @@ -171,6 +172,33 @@ func TestFindingsPageAndStats(t *testing.T) { if one, _ = d.GetFinding(ids[0]); one.Report != "# 报告\n正文" { t.Fatalf("report not read back: %q", one.Report) } + for _, test := range []struct { + name string + query string + want int + }{ + {name: "name", query: "严重漏洞标题", want: 1}, + {name: "summary", query: "summary", want: 6}, + {name: "evidence", query: "poc", want: 6}, + {name: "report", query: "正文", want: 1}, + {name: "case insensitive vulnclass", query: strings.ToUpper(vc), want: 6}, + } { + t.Run("query_"+test.name, func(t *testing.T) { + matches, searchTotal, searchErr := d.ListFindingsPage( + FindingFilter{VulnClass: vc, Query: test.query}, 1, 20, + ) + if searchErr != nil || searchTotal != test.want || len(matches) != test.want { + t.Fatalf("query %q: len=%d total=%d err=%v, want %d", test.query, len(matches), searchTotal, searchErr, test.want) + } + }) + } + if _, err := d.Exec(`UPDATE findings SET name=$1 WHERE id=$2`, "literal %_ marker", ids[1]); err != nil { + t.Fatal(err) + } + matches, searchTotal, err := d.ListFindingsPage(FindingFilter{VulnClass: vc, Query: "%_"}, 1, 20) + if err != nil || searchTotal != 1 || len(matches) != 1 || matches[0].ID != ids[1] { + t.Fatalf("query wildcards must be literal: matches=%+v total=%d err=%v", matches, searchTotal, err) + } if miss, err := d.GetFinding(-1); err != nil || miss != nil { t.Fatalf("GetFinding(-1): want nil,nil got %+v,%v", miss, err) } @@ -277,6 +305,14 @@ func TestFindingGroupsAndUnassignedPaging(t *testing.T) { t.Fatalf("SetFindingStatus[%d]: %v", i, setErr) } } + matchedGroups, matchedGroupTotal, matchedFindingTotal, err := d.ListFindingGroups( + FindingFilter{VulnClass: vc, Query: "summary 0"}, 1, 10, + ) + if err != nil || matchedGroupTotal != 1 || matchedFindingTotal != 1 || len(matchedGroups) != 1 || + matchedGroups[0].TaskID == nil || *matchedGroups[0].TaskID != taskA.ID { + t.Fatalf("query-filtered groups: %+v groups=%d findings=%d err=%v", + matchedGroups, matchedGroupTotal, matchedFindingTotal, err) + } // Retained findings from deleted tasks join the same bucket as findings that // were created without any task. if err := d.DeleteTask(taskB.ID); err != nil { diff --git a/db/llm_usage.go b/db/llm_usage.go index ecbf83a..21b3b8c 100644 --- a/db/llm_usage.go +++ b/db/llm_usage.go @@ -1,6 +1,10 @@ package db -import "strconv" +import ( + "sort" + "strconv" + "time" +) // LLMUsage is one lightweight LLM-call metering row — the always-on usage ledger, // distinct from llm_records (which stores full request/response bodies and is a @@ -129,7 +133,43 @@ ORDER BY SUM(input_tokens) + SUM(output_tokens) DESC`) } out = append(out, p) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + byName := make(map[string]ProfileUsage, len(out)) + for _, current := range out { + byName[current.ProfileName] = current + } + for _, aggregate := range archived { + for _, cold := range aggregate.TokenProfiles { + current := byName[cold.ProfileName] + current.ProfileName = cold.ProfileName + current.Calls += cold.Calls + current.Tasks += cold.Tasks + current.InputTokens += cold.InputTokens + current.OutputTokens += cold.OutputTokens + current.CacheReadTokens += cold.CacheReadTokens + current.CacheWriteTokens += cold.CacheWriteTokens + byName[cold.ProfileName] = current + } + } + out = out[:0] + for _, current := range byName { + out = append(out, current) + } + sort.Slice(out, func(i, j int) bool { + left := out[i].InputTokens + out[i].OutputTokens + right := out[j].InputTokens + out[j].OutputTokens + if left != right { + return left > right + } + return out[i].ProfileName < out[j].ProfileName + }) + return out, nil } // ProfileDayUsage is one (profile, UTC calendar day) token bucket for the daily @@ -169,7 +209,44 @@ ORDER BY day`, days) } out = append(out, p) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + cutoff := time.Now().UTC().AddDate(0, 0, -days).Format("2006-01-02") + byKey := make(map[string]ProfileDayUsage, len(out)) + for _, current := range out { + byKey[current.ProfileName+"\x00"+current.Date] = current + } + for _, aggregate := range archived { + for _, cold := range aggregate.TokenDaily { + if cold.Date < cutoff { + continue + } + key := cold.ProfileName + "\x00" + cold.Date + current := byKey[key] + current.ProfileName = cold.ProfileName + current.Date = cold.Date + current.InputTokens += cold.InputTokens + current.OutputTokens += cold.OutputTokens + current.CacheReadTokens += cold.CacheReadTokens + byKey[key] = current + } + } + out = out[:0] + for _, current := range byKey { + out = append(out, current) + } + sort.Slice(out, func(i, j int) bool { + if out[i].Date != out[j].Date { + return out[i].Date < out[j].Date + } + return out[i].ProfileName < out[j].ProfileName + }) + return out, nil } // ParseExpID turns the exploration-id segment parsed from a session string into an diff --git a/db/schema.sql b/db/schema.sql index 6d962d9..de0b7cd 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -395,6 +395,7 @@ CREATE TABLE IF NOT EXISTS tasks ( pinned_at TIMESTAMPTZ, first_run_at TIMESTAMPTZ, deadline_at TIMESTAMPTZ, + archived_at TIMESTAMPTZ, deleted_at TIMESTAMPTZ, completed_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), @@ -424,10 +425,55 @@ ALTER TABLE tasks ADD COLUMN IF NOT EXISTS category_id BIGINT REFERENCES task_ca ALTER TABLE tasks ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ; 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; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS archived_at TIMESTAMPTZ; CREATE INDEX IF NOT EXISTS idx_tasks_category ON tasks(category_id, created_at DESC) WHERE deleted_at IS NULL AND category_id IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_tasks_pinned ON tasks(pinned_at DESC) WHERE deleted_at IS NULL AND pinned_at IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_tasks_archived ON tasks(archived_at DESC) + WHERE archived_at IS NOT NULL; + +-- Cold task archives retain only compact metadata in PostgreSQL. The complete +-- task payload lives in a versioned .tar.zst package under data/archives/tasks. +-- task_id stays unique so an operation can be retried safely after a restart. +CREATE TABLE IF NOT EXISTS task_archives ( + id BIGSERIAL PRIMARY KEY, + task_id BIGINT NOT NULL UNIQUE REFERENCES tasks(id) ON DELETE CASCADE, + state TEXT NOT NULL DEFAULT 'archive_queued' CHECK (state IN ( + 'archive_queued','archiving','archive_failed','ready', + 'restore_queued','restoring','restore_failed', + 'delete_queued','deleting','delete_failed' + )), + phase TEXT NOT NULL DEFAULT 'queued', + progress INTEGER NOT NULL DEFAULT 0 CHECK (progress BETWEEN 0 AND 100), + error TEXT NOT NULL DEFAULT '', + warnings JSONB NOT NULL DEFAULT '[]', + format_version INTEGER NOT NULL DEFAULT 1, + archive_path TEXT NOT NULL DEFAULT '', + sha256 TEXT NOT NULL DEFAULT '', + original_size BIGINT NOT NULL DEFAULT 0, + compressed_size BIGINT NOT NULL DEFAULT 0, + task_name TEXT NOT NULL DEFAULT '', + task_description TEXT NOT NULL DEFAULT '', + task_goal TEXT NOT NULL DEFAULT '', + original_status TEXT NOT NULL DEFAULT '', + category_id_snapshot BIGINT, + category_name_snapshot TEXT NOT NULL DEFAULT '', + source_task_ids BIGINT[] NOT NULL DEFAULT '{}', + remaining_timeout_seconds BIGINT NOT NULL DEFAULT 0, + data_counts JSONB NOT NULL DEFAULT '{}', + aggregate_stats JSONB NOT NULL DEFAULT '{}', + archived_at TIMESTAMPTZ, + requested_at TIMESTAMPTZ NOT NULL DEFAULT now(), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_task_archives_state ON task_archives(state, requested_at, id); +CREATE INDEX IF NOT EXISTS idx_task_archives_archived ON task_archives(archived_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_task_archives_sources ON task_archives USING GIN(source_task_ids); +DROP TRIGGER IF EXISTS trg_task_archives_upd ON task_archives; +CREATE TRIGGER trg_task_archives_upd BEFORE UPDATE ON task_archives + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); -- Reusable task description/goal presets. nkey is the normalized, case-insensitive -- identity used to reject visually equivalent duplicate names. diff --git a/db/skill_usage.go b/db/skill_usage.go index 61d3a63..cdabee1 100644 --- a/db/skill_usage.go +++ b/db/skill_usage.go @@ -90,7 +90,14 @@ ORDER BY COUNT(*) DESC, skill`) } out = append(out, s) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + return mergeArchivedSkillStats(out, archived, false), nil } // MissingSkillStats returns the skill names agents asked for that do not exist, @@ -107,8 +114,7 @@ SELECT skill, COUNT(*) AS calls, FROM skill_usage WHERE NOT found GROUP BY skill -ORDER BY COUNT(*) DESC, skill -LIMIT $1`, limit) +ORDER BY COUNT(*) DESC, skill`) if err != nil { return nil, err } @@ -131,7 +137,18 @@ LIMIT $1`, limit) } out = append(out, s) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + out = mergeArchivedSkillStats(out, archived, true) + if len(out) > limit { + out = out[:limit] + } + return out, nil } // SkillCall is one row of a skill's recent-call list (detail panel). diff --git a/db/task_archive_aggregate_stats.go b/db/task_archive_aggregate_stats.go new file mode 100644 index 0000000..c7e149b --- /dev/null +++ b/db/task_archive_aggregate_stats.go @@ -0,0 +1,91 @@ +package db + +import ( + "encoding/json" + "sort" +) + +// taskArchiveAggregate is the compact, non-sensitive statistics section of a +// cold archive manifest. Fields are additive so older archive formats remain +// readable when new dashboard dimensions are introduced. +type taskArchiveAggregate struct { + TokenProfiles []ProfileUsage `json:"token_profiles"` + TokenDaily []ProfileDayUsage `json:"token_daily"` + Skills map[string]int `json:"skills"` + SkillStats []SkillStat `json:"skill_stats"` + MissingSkillStats []SkillStat `json:"missing_skill_stats"` + Tools map[string]int `json:"tools"` + FindingStats FindingStats `json:"finding_stats"` +} + +func (d *DB) archivedTaskAggregates() ([]taskArchiveAggregate, error) { + rawItems, err := d.ArchivedAggregateStats() + if err != nil { + return nil, err + } + out := make([]taskArchiveAggregate, 0, len(rawItems)) + for _, raw := range rawItems { + var item taskArchiveAggregate + if err := json.Unmarshal(raw, &item); err != nil { + return nil, err + } + out = append(out, item) + } + return out, nil +} + +func mergeAgentKeys(current, incoming []string) []string { + seen := make(map[string]bool, len(current)+len(incoming)) + out := make([]string, 0, len(current)+len(incoming)) + for _, group := range [][]string{current, incoming} { + for _, key := range group { + if key == "" || seen[key] { + continue + } + seen[key] = true + out = append(out, key) + } + } + return out +} + +func mergeArchivedSkillStats(live []SkillStat, archived []taskArchiveAggregate, missing bool) []SkillStat { + byName := make(map[string]SkillStat, len(live)) + for _, current := range live { + byName[current.Skill] = current + } + for _, aggregate := range archived { + coldStats := aggregate.SkillStats + if missing { + coldStats = aggregate.MissingSkillStats + } else if len(coldStats) == 0 { + // Format v1 initially retained only the call-count map. Preserve those + // summaries even though agent and timestamp dimensions are unavailable. + for skill, calls := range aggregate.Skills { + coldStats = append(coldStats, SkillStat{Skill: skill, Calls: calls, Tasks: 1}) + } + } + for _, cold := range coldStats { + current := byName[cold.Skill] + current.Skill = cold.Skill + current.Calls += cold.Calls + current.Tasks += cold.Tasks + current.Agents = mergeAgentKeys(current.Agents, cold.Agents) + if cold.LastUsed != nil && (current.LastUsed == nil || cold.LastUsed.After(*current.LastUsed)) { + current.LastUsed = cold.LastUsed + } + byName[cold.Skill] = current + } + } + out := make([]SkillStat, 0, len(byName)) + for _, current := range byName { + out = append(out, current) + } + sort.Slice(out, func(i, j int) bool { + if out[i].Calls != out[j].Calls { + return out[i].Calls > out[j].Calls + } + return out[i].Skill < out[j].Skill + }) + return out +} diff --git a/db/task_archives.go b/db/task_archives.go new file mode 100644 index 0000000..bef9f3a --- /dev/null +++ b/db/task_archives.go @@ -0,0 +1,667 @@ +package db + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "net/url" + "sort" + "strconv" + "strings" + "time" +) + +const TaskArchiveFormatVersion = 1 + +const ( + ArchiveQueued = "archive_queued" + Archiving = "archiving" + ArchiveFailed = "archive_failed" + ArchiveReady = "ready" + RestoreQueued = "restore_queued" + Restoring = "restoring" + RestoreFailed = "restore_failed" + DeleteQueued = "delete_queued" + Deleting = "deleting" + DeleteFailed = "delete_failed" +) + +var ( + ErrTaskArchiveNotFound = errors.New("task archive not found") + ErrTaskArchiveIneligible = errors.New("task must be paused or terminal before archiving") + ErrTaskArchiveQueued = errors.New("queued task must be paused before archiving") + ErrTaskArchiveDependent = errors.New("task is inherited by a live task") + ErrTaskArchiveState = errors.New("task archive state does not allow this operation") + ErrTaskArchiveDeleteBlocked = errors.New("task archive is required by another archive") + ErrTaskArchiveFormatMismatch = errors.New("task archive format is not supported") +) + +// TaskArchive is the compact PostgreSQL record retained while a task is cold. +// Sensitive profile configuration and API keys are intentionally absent. +type TaskArchive struct { + ID int64 `json:"id"` + TaskID int64 `json:"task_id"` + State string `json:"state"` + Phase string `json:"phase"` + Progress int `json:"progress"` + Error string `json:"error,omitempty"` + Warnings json.RawMessage `json:"warnings"` + FormatVersion int `json:"format_version"` + ArchivePath string `json:"-"` + SHA256 string `json:"sha256,omitempty"` + OriginalSize int64 `json:"original_size"` + CompressedSize int64 `json:"compressed_size"` + TaskName string `json:"task_name"` + TaskDescription string `json:"task_description"` + TaskGoal string `json:"task_goal"` + OriginalStatus string `json:"original_status"` + CategoryIDSnapshot *int64 `json:"category_id,omitempty"` + CategoryNameSnapshot string `json:"category_name,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids"` + RemainingTimeoutSeconds int64 `json:"remaining_timeout_seconds"` + DataCounts json.RawMessage `json:"data_counts"` + AggregateStats json.RawMessage `json:"aggregate_stats"` + ArchivedAt *time.Time `json:"archived_at,omitempty"` + RequestedAt time.Time `json:"requested_at"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// TaskArchiveBlockers returns one live direct dependent for every source task +// that cannot currently be archived. Dependents already queued for archiving do +// not block their source because the FIFO worker will compact them first. +func (d *DB) TaskArchiveBlockers() (map[int64]int64, error) { + rows, err := d.Query(`SELECT relation.source_task_id, MIN(child.id) +FROM task_relations relation +JOIN tasks child ON child.id=relation.task_id AND child.deleted_at IS NULL +LEFT JOIN task_archives pending ON pending.task_id=child.id +WHERE pending.id IS NULL OR pending.state NOT IN ('archive_queued','archiving') +GROUP BY relation.source_task_id`) + if err != nil { + return nil, err + } + defer rows.Close() + blockers := map[int64]int64{} + for rows.Next() { + var sourceID, dependentID int64 + if err := rows.Scan(&sourceID, &dependentID); err != nil { + return nil, err + } + blockers[sourceID] = dependentID + } + return blockers, rows.Err() +} + +type TaskArchivePage struct { + Items []TaskArchive `json:"items"` + Total int `json:"total"` + Page int `json:"page"` + Size int `json:"size"` +} + +// TaskArchiveSnapshot is serialized into manifest.json inside the cold package. +// Each value in Tables is a JSON array generated by PostgreSQL row_to_json, which +// preserves exact IDs, timestamps, arrays and JSON payloads for restoration. +type TaskArchiveSnapshot struct { + FormatVersion int `json:"format_version"` + CreatedAt time.Time `json:"created_at"` + TaskID int64 `json:"task_id"` + ExplorationID int64 `json:"exploration_id"` + SourceTaskIDs []int64 `json:"source_task_ids"` + Hosts []string `json:"hosts"` + ExclusiveHosts []string `json:"exclusive_hosts"` + ExclusiveAssetIDs []int64 `json:"exclusive_asset_ids"` + Tables map[string]json.RawMessage `json:"tables"` + DataCounts map[string]int64 `json:"data_counts"` + AggregateStats map[string]any `json:"aggregate_stats"` +} + +func scanTaskArchive(sc interface{ Scan(...any) error }) (*TaskArchive, error) { + var item TaskArchive + var sources string + err := sc.Scan( + &item.ID, &item.TaskID, &item.State, &item.Phase, &item.Progress, &item.Error, + &item.Warnings, &item.FormatVersion, &item.ArchivePath, &item.SHA256, + &item.OriginalSize, &item.CompressedSize, &item.TaskName, &item.TaskDescription, + &item.TaskGoal, &item.OriginalStatus, &item.CategoryIDSnapshot, + &item.CategoryNameSnapshot, &sources, &item.RemainingTimeoutSeconds, + &item.DataCounts, &item.AggregateStats, &item.ArchivedAt, &item.RequestedAt, + &item.CreatedAt, &item.UpdatedAt, + ) + if err != nil { + return nil, err + } + if err := json.Unmarshal([]byte(sources), &item.SourceTaskIDs); err != nil { + return nil, err + } + if len(item.Warnings) == 0 { + item.Warnings = json.RawMessage("[]") + } + if len(item.DataCounts) == 0 { + item.DataCounts = json.RawMessage("{}") + } + if len(item.AggregateStats) == 0 { + item.AggregateStats = json.RawMessage("{}") + } + return &item, nil +} + +const taskArchiveCols = `id, task_id, state, phase, progress, COALESCE(error,''), warnings, +format_version, COALESCE(archive_path,''), COALESCE(sha256,''), original_size, +compressed_size, COALESCE(task_name,''), COALESCE(task_description,''), +COALESCE(task_goal,''), COALESCE(original_status,''), category_id_snapshot, +COALESCE(category_name_snapshot,''), array_to_json(source_task_ids)::text, +remaining_timeout_seconds, data_counts, aggregate_stats, archived_at, +requested_at, created_at, updated_at` + +func (d *DB) GetTaskArchive(id int64) (*TaskArchive, error) { + item, err := scanTaskArchive(d.QueryRow(`SELECT `+taskArchiveCols+` FROM task_archives WHERE id=$1`, id)) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + return item, err +} + +func (d *DB) GetTaskArchiveByTask(taskID int64) (*TaskArchive, error) { + item, err := scanTaskArchive(d.QueryRow(`SELECT `+taskArchiveCols+` FROM task_archives WHERE task_id=$1`, taskID)) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + return item, err +} + +func (d *DB) ListTaskArchives(search, state string, page, size int) (TaskArchivePage, error) { + if page < 1 { + page = 1 + } + if size <= 0 || size > 100 { + size = 20 + } + search = strings.TrimSpace(search) + state = strings.TrimSpace(state) + where := `WHERE ($1='' OR task_id::text ILIKE '%'||$1||'%' OR task_name ILIKE '%'||$1||'%' OR task_description ILIKE '%'||$1||'%') +AND ($2='' OR state=$2)` + var out TaskArchivePage + out.Page, out.Size = page, size + if err := d.QueryRow(`SELECT count(*) FROM task_archives `+where, search, state).Scan(&out.Total); err != nil { + return out, err + } + rows, err := d.Query(`SELECT `+taskArchiveCols+` FROM task_archives `+where+` +ORDER BY COALESCE(archived_at, requested_at) DESC, id DESC LIMIT $3 OFFSET $4`, search, state, size, (page-1)*size) + if err != nil { + return out, err + } + defer rows.Close() + for rows.Next() { + item, err := scanTaskArchive(rows) + if err != nil { + return out, err + } + out.Items = append(out.Items, *item) + } + return out, rows.Err() +} + +// QueueTaskArchive validates lifecycle and direct inheritance while holding the +// task row. A failed archive can be explicitly retried through the same API. +func (d *DB) QueueTaskArchive(taskID int64) (*TaskArchive, error) { + tx, err := d.Begin() + if err != nil { + return nil, err + } + defer tx.Rollback() //nolint:errcheck + var name, description, goal, status, categoryName string + var categoryID *int64 + var paused, queued bool + var deadline *time.Time + err = tx.QueryRow(`SELECT COALESCE(t.name,''), t.description, t.goal, t.status, + t.category_id, COALESCE(c.name,''), t.paused, t.queued, t.deadline_at +FROM tasks t LEFT JOIN task_categories c ON c.id=t.category_id +WHERE t.id=$1 AND t.deleted_at IS NULL FOR UPDATE OF t`, taskID).Scan( + &name, &description, &goal, &status, &categoryID, &categoryName, &paused, &queued, &deadline, + ) + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrTaskArchiveNotFound + } + if err != nil { + return nil, err + } + if queued { + return nil, ErrTaskArchiveQueued + } + if !paused && !IsTerminal(status) { + return nil, ErrTaskArchiveIneligible + } + var dependent int64 + err = tx.QueryRow(`SELECT child.id FROM task_relations relation +JOIN tasks child ON child.id=relation.task_id AND child.deleted_at IS NULL +LEFT JOIN task_archives pending ON pending.task_id=child.id +WHERE relation.source_task_id=$1 + AND (pending.id IS NULL OR pending.state NOT IN ('archive_queued','archiving')) +LIMIT 1`, taskID).Scan(&dependent) + if err == nil { + return nil, fmt.Errorf("%w: task %d", ErrTaskArchiveDependent, dependent) + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + var sources []int64 + rows, err := tx.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 + } + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + rows.Close() + return nil, err + } + sources = append(sources, id) + } + if err := rows.Close(); err != nil { + return nil, err + } + if sources == nil { + sources = []int64{} + } + remaining := int64(0) + if paused && deadline != nil { + remaining = int64(time.Until(*deadline).Seconds()) + if remaining < 0 { + remaining = 0 + } + } + _, err = tx.Exec(`INSERT INTO task_archives( + task_id,state,phase,progress,error,warnings,format_version,task_name, + task_description,task_goal,original_status,category_id_snapshot, + category_name_snapshot,source_task_ids,remaining_timeout_seconds,requested_at) +VALUES ($1,$2,'queued',0,'','[]',$3,$4,$5,$6,$7,$8,$9,$10,$11,now()) +ON CONFLICT (task_id) DO UPDATE SET + state=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.state ELSE task_archives.state END, + phase=CASE WHEN task_archives.state IN ('archive_failed') THEN 'queued' ELSE task_archives.phase END, + progress=CASE WHEN task_archives.state IN ('archive_failed') THEN 0 ELSE task_archives.progress END, + error=CASE WHEN task_archives.state IN ('archive_failed') THEN '' ELSE task_archives.error END, + requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`, + taskID, ArchiveQueued, TaskArchiveFormatVersion, name, description, goal, status, + categoryID, categoryName, sources, remaining) + if err != nil { + return nil, err + } + item, err := scanTaskArchive(tx.QueryRow(`SELECT `+taskArchiveCols+` FROM task_archives WHERE task_id=$1`, taskID)) + if err != nil { + return nil, err + } + if item.State != ArchiveQueued && item.State != ArchiveFailed { + return nil, fmt.Errorf("%w: current state %s", ErrTaskArchiveState, item.State) + } + return item, tx.Commit() +} + +func (d *DB) QueueTaskArchiveRestore(id int64) (*TaskArchive, error) { + item, err := scanTaskArchive(d.QueryRow(`UPDATE task_archives +SET state=$2, phase='queued', progress=0, error='', requested_at=now() +WHERE id=$1 AND state IN ('ready','restore_failed') RETURNING `+taskArchiveCols, id, RestoreQueued)) + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrTaskArchiveState + } + return item, err +} + +func (d *DB) QueueTaskArchiveDelete(id int64) (*TaskArchive, error) { + tx, err := d.Begin() + if err != nil { + return nil, err + } + defer tx.Rollback() //nolint:errcheck + var taskID int64 + if err := tx.QueryRow(`SELECT task_id FROM task_archives WHERE id=$1 AND state IN ('ready','delete_failed') FOR UPDATE`, id).Scan(&taskID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrTaskArchiveState + } + return nil, err + } + var dependent int64 + err = tx.QueryRow(`SELECT task_id FROM task_archives +WHERE id<>$1 AND $2=ANY(source_task_ids) AND state NOT IN ('delete_queued','deleting') LIMIT 1`, id, taskID).Scan(&dependent) + if err == nil { + return nil, fmt.Errorf("%w: task %d", ErrTaskArchiveDeleteBlocked, dependent) + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + item, err := scanTaskArchive(tx.QueryRow(`UPDATE task_archives +SET state=$2, phase='queued', progress=0, error='', requested_at=now() +WHERE id=$1 RETURNING `+taskArchiveCols, id, DeleteQueued)) + if err != nil { + return nil, err + } + return item, tx.Commit() +} + +// RecoverTaskArchiveJobs makes active operations retryable after an unclean +// shutdown. Idempotent restore/import semantics keep already-completed pieces safe. +func (d *DB) RecoverTaskArchiveJobs() error { + _, err := d.Exec(`UPDATE task_archives SET + state=CASE state WHEN 'archiving' THEN 'archive_queued' + WHEN 'restoring' THEN 'restore_queued' + WHEN 'deleting' THEN 'delete_queued' ELSE state END, + phase='recovered', error='' +WHERE state IN ('archiving','restoring','deleting')`) + return err +} + +// ClaimTaskArchiveJob claims one persistent FIFO item for the single archive +// worker. It returns nil when the queue is empty. +func (d *DB) ClaimTaskArchiveJob(ctx context.Context) (*TaskArchive, error) { + tx, err := d.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer tx.Rollback() //nolint:errcheck + var id int64 + var queuedState string + err = tx.QueryRowContext(ctx, `SELECT id,state FROM task_archives +WHERE state IN ('archive_queued','restore_queued','delete_queued') +ORDER BY requested_at,id FOR UPDATE SKIP LOCKED LIMIT 1`).Scan(&id, &queuedState) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + active := map[string]string{ArchiveQueued: Archiving, RestoreQueued: Restoring, DeleteQueued: Deleting}[queuedState] + item, err := scanTaskArchive(tx.QueryRowContext(ctx, `UPDATE task_archives +SET state=$2,phase='starting',progress=1,error='' WHERE id=$1 RETURNING `+taskArchiveCols, id, active)) + if err != nil { + return nil, err + } + return item, tx.Commit() +} + +func (d *DB) UpdateTaskArchiveProgress(id int64, phase string, progress int) error { + if progress < 0 { + progress = 0 + } + if progress > 100 { + progress = 100 + } + _, err := d.Exec(`UPDATE task_archives SET phase=$2,progress=$3 WHERE id=$1`, id, phase, progress) + return err +} + +func (d *DB) AppendTaskArchiveWarning(id int64, warning string) error { + if strings.TrimSpace(warning) == "" { + return nil + } + _, err := d.Exec(`UPDATE task_archives SET warnings=warnings || jsonb_build_array($2::text) WHERE id=$1`, id, warning) + return err +} + +func (d *DB) IsTaskArchiveRestored(id int64) (bool, error) { + var restored bool + err := d.QueryRow(`SELECT task.deleted_at IS NULL AND task.archived_at IS NULL +FROM task_archives archive JOIN tasks task ON task.id=archive.task_id WHERE archive.id=$1`, id).Scan(&restored) + if errors.Is(err, sql.ErrNoRows) { + return false, ErrTaskArchiveNotFound + } + return restored, err +} + +func (d *DB) FailTaskArchiveJob(id int64, activeState string, cause error) error { + failed := map[string]string{Archiving: ArchiveFailed, Restoring: RestoreFailed, Deleting: DeleteFailed}[activeState] + if failed == "" { + return fmt.Errorf("unknown active archive state %q", activeState) + } + message := "unknown archive failure" + if cause != nil { + message = cause.Error() + } + _, err := d.Exec(`UPDATE task_archives SET state=$2,phase='failed',error=$3 WHERE id=$1`, id, failed, message) + return err +} + +func queryArchiveRows(q interface { + QueryRow(query string, args ...any) *sql.Row +}, inner string, args ...any) (json.RawMessage, error) { + var raw []byte + err := q.QueryRow(`SELECT COALESCE(jsonb_agg(to_jsonb(row_data)), '[]'::jsonb) FROM (`+inner+`) row_data`, args...).Scan(&raw) + return json.RawMessage(raw), err +} + +func rawRowCount(raw json.RawMessage) int64 { + var rows []json.RawMessage + if json.Unmarshal(raw, &rows) != nil { + return 0 + } + return int64(len(rows)) +} + +func archiveAssetIDsQuery() string { + return `SELECT id FROM assets WHERE $1=ANY(task_ids) +UNION SELECT link.asset_id FROM task_asset_links link WHERE link.task_id=$1 +UNION SELECT anchor.asset_id FROM exploration_anchors anchor + JOIN exploration_nodes node ON node.id=anchor.node_id WHERE node.exploration_id=$2 +UNION SELECT value::bigint FROM findings finding + CROSS JOIN LATERAL jsonb_array_elements_text( + CASE WHEN jsonb_typeof(finding.asset_ids)='array' THEN finding.asset_ids ELSE '[]'::jsonb END + ) value WHERE finding.task_id=$1 AND value ~ '^[0-9]+$'` +} + +// SnapshotTaskArchive reads one repeatable PostgreSQL snapshot. Task-owned Agent +// writes are already quiescent at the server barrier; repeatable-read also keeps +// the asset and accounting views mutually consistent during serialization. +func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { + tx, err := d.BeginTx(context.Background(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true}) + if err != nil { + return nil, err + } + defer tx.Rollback() //nolint:errcheck + var expID int64 + if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 AND deleted_at IS NULL`, taskID).Scan(&expID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrTaskArchiveNotFound + } + return nil, err + } + tables := make(map[string]json.RawMessage) + queries := []struct { + name string + query string + args []any + }{ + {"tasks", `SELECT * FROM tasks WHERE id=$1`, []any{taskID}}, + {"explorations", `SELECT * FROM explorations WHERE id=$1`, []any{expID}}, + {"exploration_nodes", `SELECT * FROM exploration_nodes WHERE exploration_id=$1 ORDER BY id`, []any{expID}}, + {"exploration_edges", `SELECT * FROM exploration_edges WHERE exploration_id=$1 ORDER BY src_id,dst_id`, []any{expID}}, + {"exploration_anchors", `SELECT anchor.* FROM exploration_anchors anchor JOIN exploration_nodes node ON node.id=anchor.node_id WHERE node.exploration_id=$1 ORDER BY node_id,asset_id`, []any{expID}}, + {"task_constraints", `SELECT * FROM task_constraints WHERE exploration_id=$1 ORDER BY id`, []any{expID}}, + {"activity", `SELECT * FROM activity WHERE exploration_id=$1 ORDER BY id`, []any{expID}}, + {"task_relations", `SELECT * FROM task_relations WHERE task_id=$1 ORDER BY created_at,source_task_id`, []any{taskID}}, + {"task_asset_links", `SELECT * FROM task_asset_links WHERE task_id=$1 ORDER BY asset_id`, []any{taskID}}, + {"task_llm_profiles", `SELECT * FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, []any{taskID}}, + {"task_scope", `SELECT * FROM task_scope WHERE task_id=$1 ORDER BY id`, []any{taskID}}, + {"findings", `SELECT * FROM findings WHERE task_id=$1 ORDER BY id`, []any{taskID}}, + {"llm_records", `SELECT * FROM llm_records WHERE COALESCE(task_id,'')=$1 ORDER BY id`, []any{strconv.FormatInt(taskID, 10)}}, + {"llm_usage", `SELECT * FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 ORDER BY id`, []any{strconv.FormatInt(taskID, 10), expID}}, + {"skill_usage", `SELECT * FROM skill_usage WHERE task_id=$1 OR exploration_id=$2 ORDER BY id`, []any{taskID, expID}}, + {"tool_usage", `SELECT * FROM tool_usage WHERE task_id=$1 OR exploration_id=$2 ORDER BY id`, []any{taskID, expID}}, + {"intercept_pending", `SELECT * FROM intercept_pending WHERE COALESCE(task_id,'')=$1 ORDER BY id`, []any{strconv.FormatInt(taskID, 10)}}, + {"assets", `SELECT asset.* FROM assets asset WHERE asset.id IN (` + archiveAssetIDsQuery() + `) ORDER BY asset.id`, []any{taskID, expID}}, + } + counts := make(map[string]int64, len(queries)) + for _, query := range queries { + raw, err := queryArchiveRows(tx, query.query, query.args...) + if err != nil { + return nil, fmt.Errorf("snapshot %s: %w", query.name, err) + } + tables[query.name] = raw + counts[query.name] = rawRowCount(raw) + } + + assetIDs, exclusiveAssetIDs, hosts, exclusiveHosts, err := archiveAssetMetadata(tx, taskID, expID, tables["assets"]) + if err != nil { + return nil, err + } + _ = assetIDs // retained in the assets table payload; only exclusive ids need a side channel. + var sources []int64 + rows, err := tx.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 + } + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + rows.Close() + return nil, err + } + sources = append(sources, id) + } + if err := rows.Close(); err != nil { + return nil, err + } + stats, err := taskArchiveAggregates(tx, taskID, expID) + if err != nil { + return nil, err + } + snapshot := &TaskArchiveSnapshot{ + FormatVersion: TaskArchiveFormatVersion, CreatedAt: time.Now().UTC(), TaskID: taskID, + ExplorationID: expID, SourceTaskIDs: sources, Hosts: hosts, ExclusiveHosts: exclusiveHosts, + ExclusiveAssetIDs: exclusiveAssetIDs, Tables: tables, DataCounts: counts, AggregateStats: stats, + } + return snapshot, tx.Commit() +} + +func archiveAssetMetadata(tx *sql.Tx, taskID, expID int64, assetRows json.RawMessage) ([]int64, []int64, []string, []string, error) { + var rows []map[string]any + if err := json.Unmarshal(assetRows, &rows); err != nil { + return nil, nil, nil, nil, err + } + allHosts := map[string]struct{}{} + assetIDs := make([]int64, 0, len(rows)) + for _, row := range rows { + if id, ok := jsonInt64(row["id"]); ok { + assetIDs = append(assetIDs, id) + } + for _, key := range []string{"domain", "ip"} { + if value, _ := row[key].(string); strings.TrimSpace(value) != "" { + allHosts[strings.ToLower(strings.TrimSpace(value))] = struct{}{} + } + } + if rawURL, _ := row["url"].(string); rawURL != "" { + if parsed, err := url.Parse(rawURL); err == nil && parsed.Hostname() != "" { + allHosts[strings.ToLower(parsed.Hostname())] = struct{}{} + } + } + } + exclusiveHosts, err := hostsForTaskDeletion(tx, taskID, expID) + if err != nil { + return nil, nil, nil, nil, err + } + rowsID, err := tx.Query(`WITH candidate AS (`+archiveAssetIDsQuery()+`) +SELECT asset.id FROM assets asset JOIN candidate ON candidate.id=asset.id +WHERE asset.company_id IS NULL +AND NOT EXISTS (SELECT 1 FROM tasks task WHERE task.id<>$1 AND task.deleted_at IS NULL AND task.id=ANY(asset.task_ids)) +AND NOT EXISTS ( + SELECT 1 FROM exploration_anchors anchor JOIN exploration_nodes node ON node.id=anchor.node_id + JOIN tasks task ON task.exploration_id=node.exploration_id + WHERE anchor.asset_id=asset.id AND task.id<>$1 AND task.deleted_at IS NULL +) ORDER BY asset.id`, taskID, expID) + if err != nil { + return nil, nil, nil, nil, err + } + var exclusiveIDs []int64 + for rowsID.Next() { + var id int64 + if err := rowsID.Scan(&id); err != nil { + rowsID.Close() + return nil, nil, nil, nil, err + } + exclusiveIDs = append(exclusiveIDs, id) + } + if err := rowsID.Close(); err != nil { + return nil, nil, nil, nil, err + } + hosts := make([]string, 0, len(allHosts)) + for host := range allHosts { + hosts = append(hosts, host) + } + sort.Strings(hosts) + return assetIDs, exclusiveIDs, hosts, exclusiveHosts, nil +} + +func taskArchiveAggregates(tx *sql.Tx, taskID, expID int64) (map[string]any, error) { + stats := map[string]any{} + var calls, input, output, cacheRead, cacheWrite int64 + if err := tx.QueryRow(`SELECT count(*),COALESCE(sum(input_tokens),0),COALESCE(sum(output_tokens),0), +COALESCE(sum(cache_read),0),COALESCE(sum(cache_write),0) +FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2`, strconv.FormatInt(taskID, 10), expID). + Scan(&calls, &input, &output, &cacheRead, &cacheWrite); err != nil { + return nil, err + } + stats["tokens"] = map[string]int64{"calls": calls, "input_tokens": input, "output_tokens": output, "cache_read_tokens": cacheRead, "cache_write_tokens": cacheWrite} + for _, item := range []struct { + name string + query string + args []any + }{ + {"token_profiles", `SELECT COALESCE(jsonb_agg(to_jsonb(x)),'[]'::jsonb) FROM ( +SELECT COALESCE(profile_name,'') profile_name,count(*) calls,1 tasks, + COALESCE(sum(input_tokens),0) input_tokens,COALESCE(sum(output_tokens),0) output_tokens, + COALESCE(sum(cache_read),0) cache_read_tokens,COALESCE(sum(cache_write),0) cache_write_tokens +FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 +GROUP BY profile_name ORDER BY sum(input_tokens)+sum(output_tokens) DESC) x`, []any{strconv.FormatInt(taskID, 10), expID}}, + {"token_daily", `SELECT COALESCE(jsonb_agg(to_jsonb(x)),'[]'::jsonb) FROM ( +SELECT COALESCE(profile_name,'') profile_name,to_char(ts AT TIME ZONE 'UTC','YYYY-MM-DD') date, + COALESCE(sum(input_tokens),0) input_tokens,COALESCE(sum(output_tokens),0) output_tokens, + COALESCE(sum(cache_read),0) cache_read_tokens +FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 +GROUP BY profile_name,date ORDER BY date) x`, []any{strconv.FormatInt(taskID, 10), expID}}, + {"skills", `SELECT COALESCE(jsonb_object_agg(name,n),'{}'::jsonb) FROM (SELECT skill name,count(*) n FROM skill_usage WHERE (task_id=$1 OR exploration_id=$2) AND found GROUP BY skill) x`, []any{taskID, expID}}, + {"skill_stats", `SELECT COALESCE(jsonb_agg(to_jsonb(x)),'[]'::jsonb) FROM ( +SELECT skill,count(*) calls,1 tasks, + COALESCE(array_agg(DISTINCT agent_key) FILTER (WHERE agent_key IS NOT NULL),ARRAY[]::text[]) agents, + max(ts) last_used +FROM skill_usage WHERE (task_id=$1 OR exploration_id=$2) AND found GROUP BY skill) x`, []any{taskID, expID}}, + {"missing_skill_stats", `SELECT COALESCE(jsonb_agg(to_jsonb(x)),'[]'::jsonb) FROM ( +SELECT skill,count(*) calls,0 tasks, + COALESCE(array_agg(DISTINCT agent_key) FILTER (WHERE agent_key IS NOT NULL),ARRAY[]::text[]) agents, + max(ts) last_used +FROM skill_usage WHERE (task_id=$1 OR exploration_id=$2) AND NOT found GROUP BY skill) x`, []any{taskID, expID}}, + {"tools", `SELECT COALESCE(jsonb_object_agg(name,n),'{}'::jsonb) FROM (SELECT tool_key name,count(*) n FROM tool_usage WHERE task_id=$1 OR exploration_id=$2 GROUP BY tool_key) x`, []any{taskID, expID}}, + {"findings", `SELECT COALESCE(jsonb_object_agg(name,n),'{}'::jsonb) FROM (SELECT COALESCE(NULLIF(severity,''),'unknown') name,count(*) n FROM findings WHERE task_id=$1 GROUP BY severity) x`, []any{taskID}}, + {"finding_stats", `SELECT jsonb_build_object( + 'total',count(*),'pending',count(*) FILTER (WHERE status='pending'), + 'critical',count(*) FILTER (WHERE severity='critical'),'high',count(*) FILTER (WHERE severity='high'), + 'medium',count(*) FILTER (WHERE severity='medium'),'low',count(*) FILTER (WHERE severity='low'), + 'vulnclasses',COALESCE(jsonb_agg(DISTINCT vulnclass) FILTER (WHERE vulnclass<>''),'[]'::jsonb)) +FROM findings WHERE task_id=$1`, []any{taskID}}, + } { + var raw []byte + if err := tx.QueryRow(item.query, item.args...).Scan(&raw); err != nil { + return nil, err + } + var value any + if err := json.Unmarshal(raw, &value); err != nil { + return nil, err + } + stats[item.name] = value + } + return stats, nil +} + +func jsonInt64(value any) (int64, bool) { + switch value := value.(type) { + case float64: + return int64(value), value == float64(int64(value)) + case json.Number: + id, err := value.Int64() + return id, err == nil + case string: + id, err := strconv.ParseInt(value, 10, 64) + return id, err == nil + default: + return 0, false + } +} diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go new file mode 100644 index 0000000..d3f15b0 --- /dev/null +++ b/db/task_archives_restore.go @@ -0,0 +1,648 @@ +package db + +import ( + "bytes" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "time" +) + +// CompleteTaskArchive performs the hot-store compaction only after the external +// package has been fully written and checksummed. The task/exploration rows remain +// as minimal ID stubs; all heavyweight task-owned rows move into the package. +func (d *DB) CompleteTaskArchive( + archiveID int64, + snapshot *TaskArchiveSnapshot, + archivePath, sha256 string, + originalSize, compressedSize int64, +) error { + if snapshot == nil || snapshot.FormatVersion != TaskArchiveFormatVersion { + return ErrTaskArchiveFormatMismatch + } + countsRaw, err := json.Marshal(snapshot.DataCounts) + if err != nil { + return err + } + statsRaw, err := json.Marshal(snapshot.AggregateStats) + if err != nil { + return err + } + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck + var taskID, expID int64 + var state string + if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state +FROM task_archives archive JOIN tasks task ON task.id=archive.task_id +WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID, &state); err != nil { + return err + } + if state != Archiving || taskID != snapshot.TaskID || expID != snapshot.ExplorationID { + return fmt.Errorf("%w: archive snapshot identity/state mismatch", ErrTaskArchiveState) + } + var dependent int64 + err = tx.QueryRow(`SELECT child.id FROM task_relations relation +JOIN tasks child ON child.id=relation.task_id AND child.deleted_at IS NULL +WHERE relation.source_task_id=$1 LIMIT 1`, taskID).Scan(&dependent) + if err == nil { + return fmt.Errorf("%w: task %d", ErrTaskArchiveDependent, dependent) + } + if !errors.Is(err, sql.ErrNoRows) { + return err + } + // Prevent task/asset ownership from changing while exclusivity is rechecked. + if _, err := tx.Exec(`LOCK TABLE assets, exploration_anchors IN SHARE ROW EXCLUSIVE MODE`); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM intercept_pending WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(taskID, 10)); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM llm_records WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(taskID, 10)); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2`, strconv.FormatInt(taskID, 10), expID); err != nil { + return err + } + for _, table := range []string{"skill_usage", "tool_usage"} { + if _, err := tx.Exec(`DELETE FROM `+table+` WHERE task_id=$1 OR exploration_id=$2`, taskID, expID); err != nil { + return err + } + } + if _, err := tx.Exec(`DELETE FROM findings WHERE task_id=$1`, taskID); err != nil { + return err + } + for _, statement := range []struct { + query string + args []any + }{ + {`DELETE FROM task_relations WHERE task_id=$1 OR source_task_id=$1`, []any{taskID}}, + {`DELETE FROM task_asset_links WHERE task_id=$1`, []any{taskID}}, + {`DELETE FROM task_llm_profiles WHERE task_id=$1`, []any{taskID}}, + {`DELETE FROM task_scope WHERE task_id=$1`, []any{taskID}}, + {`DELETE FROM task_constraints WHERE exploration_id=$1`, []any{expID}}, + {`DELETE FROM activity WHERE exploration_id=$1`, []any{expID}}, + {`DELETE FROM exploration_nodes WHERE exploration_id=$1`, []any{expID}}, + } { + if _, err := tx.Exec(statement.query, statement.args...); err != nil { + return err + } + } + if _, err := tx.Exec(`UPDATE assets SET task_ids=array_remove(task_ids,$1) WHERE $1=ANY(task_ids)`, taskID); err != nil { + return err + } + if len(snapshot.ExclusiveAssetIDs) > 0 { + if _, err := tx.Exec(`DELETE FROM assets asset +WHERE asset.id=ANY($2::bigint[]) AND asset.company_id IS NULL +AND NOT EXISTS (SELECT 1 FROM tasks task WHERE task.id<>$1 AND task.deleted_at IS NULL AND task.id=ANY(asset.task_ids)) +AND NOT EXISTS ( + SELECT 1 FROM exploration_anchors anchor JOIN exploration_nodes node ON node.id=anchor.node_id + JOIN tasks task ON task.exploration_id=node.exploration_id + WHERE anchor.asset_id=asset.id AND task.id<>$1 AND task.deleted_at IS NULL +)`, taskID, snapshot.ExclusiveAssetIDs); err != nil { + return err + } + } + if _, err := tx.Exec(`UPDATE explorations SET description='',goal='',status='open' WHERE id=$1`, expID); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE tasks SET + name='',category_id=NULL,description='',goal='',paused=true,queued=false,queued_at=NULL,queue_mode='', + llm_profile_id=NULL,active_llm_profile_id=NULL,llm_chain_revision=llm_chain_revision+1, + company_id=NULL,parent_ref=NULL,timeout_seconds=0,coverage_enabled=true,pinned_at=NULL, + first_run_at=NULL,deadline_at=NULL,archived_at=now(),deleted_at=now() +WHERE id=$1`, taskID); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE task_archives SET + state=$2,phase='ready',progress=100,error='',archive_path=$3,sha256=$4, + original_size=$5,compressed_size=$6,data_counts=$7,aggregate_stats=$8, + archived_at=now(),warnings='[]' +WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize, string(countsRaw), string(statsRaw)); err != nil { + return err + } + return tx.Commit() +} + +// RestoreTaskArchive restores PostgreSQL rows from a verified manifest. It is +// idempotent for accounting/traffic retry scenarios and returns non-fatal +// warnings for global objects that intentionally are not recreated. +func (d *DB) RestoreTaskArchive(archiveID int64, snapshot *TaskArchiveSnapshot, remainingTimeoutSeconds int64) ([]string, error) { + if snapshot == nil || snapshot.FormatVersion != TaskArchiveFormatVersion { + return nil, ErrTaskArchiveFormatMismatch + } + tx, err := d.Begin() + if err != nil { + return nil, err + } + defer tx.Rollback() //nolint:errcheck + var taskID, expID int64 + var state string + if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state +FROM task_archives archive JOIN tasks task ON task.id=archive.task_id +WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID, &state); err != nil { + return nil, err + } + if state != Restoring || taskID != snapshot.TaskID || expID != snapshot.ExplorationID { + return nil, fmt.Errorf("%w: restore snapshot identity/state mismatch", ErrTaskArchiveState) + } + warnings := []string{} + assetMap, assetWarnings, err := restoreArchiveAssets(tx, taskID, snapshot.Tables["assets"]) + if err != nil { + return nil, err + } + warnings = append(warnings, assetWarnings...) + + // The task row exists as an archived stub. Restore its global references only + // when the current instance still owns them; never recreate categories/profiles. + taskRows, err := decodeArchiveRows(snapshot.Tables["tasks"]) + if err != nil || len(taskRows) != 1 { + return nil, fmt.Errorf("restore task row: expected one row: %w", err) + } + taskRow := taskRows[0] + categoryID, _ := jsonInt64(taskRow["category_id"]) + if categoryID > 0 && !rowExists(tx, "task_categories", categoryID) { + taskRow["category_id"] = nil + warnings = append(warnings, fmt.Sprintf("任务分类 %d 已删除，已恢复为未分类", categoryID)) + } + companyID, _ := jsonInt64(taskRow["company_id"]) + if companyID > 0 && !rowExists(tx, "companies", companyID) { + taskRow["company_id"] = nil + warnings = append(warnings, fmt.Sprintf("任务企业 %d 已删除，企业关联已跳过", companyID)) + } + for _, key := range []string{"llm_profile_id", "active_llm_profile_id"} { + profileID, _ := jsonInt64(taskRow[key]) + if profileID > 0 && !rowExists(tx, "llm_profiles", profileID) { + taskRow[key] = nil + warnings = append(warnings, fmt.Sprintf("LLM 配置 %d 已删除，已从任务配置中移除", profileID)) + } + } + if err := restoreExplorationStub(tx, snapshot.Tables["explorations"], expID); err != nil { + return nil, err + } + if err := restoreTaskStub(tx, taskRow, taskID, remainingTimeoutSeconds); err != nil { + return nil, err + } + + remappedTables, err := remapArchiveAssetReferences(snapshot.Tables, assetMap) + if err != nil { + return nil, err + } + // Insert graph rows in foreign-key order. The archived stub has no graph rows, + // so an ID conflict signals external corruption and must stop the restore. + for _, table := range []string{"exploration_nodes", "exploration_edges", "exploration_anchors", "task_constraints", "activity"} { + if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil { + return nil, fmt.Errorf("restore %s: %w", table, err) + } + } + if warning, err := restoreTaskRelations(tx, taskID, remappedTables["task_relations"]); err != nil { + return nil, err + } else { + warnings = append(warnings, warning...) + } + if warning, err := restoreTaskScopes(tx, remappedTables["task_scope"]); err != nil { + return nil, err + } else { + warnings = append(warnings, warning...) + } + if warning, err := restoreTaskLLMProfiles(tx, remappedTables["task_llm_profiles"]); err != nil { + return nil, err + } else { + warnings = append(warnings, warning...) + } + for _, table := range []string{"task_asset_links", "findings", "llm_records", "llm_usage", "skill_usage", "tool_usage"} { + if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil { + return nil, fmt.Errorf("restore %s: %w", table, err) + } + } + if err := restoreInterceptRows(tx, remappedTables["intercept_pending"]); err != nil { + return nil, err + } + warningsRaw, _ := json.Marshal(warnings) + if _, err := tx.Exec(`UPDATE task_archives SET warnings=$2,phase='database_restored',progress=85,error='' WHERE id=$1`, archiveID, string(warningsRaw)); err != nil { + return nil, err + } + return warnings, tx.Commit() +} + +func restoreExplorationStub(tx *sql.Tx, raw json.RawMessage, expID int64) error { + _, err := tx.Exec(`UPDATE explorations current SET + description=archived.description,goal=archived.goal,status=archived.status, + created_at=archived.created_at,updated_at=archived.updated_at +FROM json_populate_record(NULL::explorations,$2::json) archived +WHERE current.id=$1 AND archived.id=$1`, expID, string(firstArchiveRow(raw))) + return err +} + +func restoreTaskStub(tx *sql.Tx, row map[string]any, taskID, remaining int64) error { + raw, err := json.Marshal(row) + if err != nil { + return err + } + _, err = tx.Exec(`UPDATE tasks current SET + name=archived.name,category_id=archived.category_id,description=archived.description,goal=archived.goal, + status=archived.status,paused=archived.paused,queued=false,queued_at=NULL,queue_mode='', + llm_profile_id=archived.llm_profile_id,active_llm_profile_id=archived.active_llm_profile_id, + llm_chain_revision=archived.llm_chain_revision,company_id=archived.company_id,parent_ref=archived.parent_ref, + timeout_seconds=archived.timeout_seconds,plan_heartbeat_seconds=archived.plan_heartbeat_seconds, + coverage_enabled=archived.coverage_enabled,pinned_at=archived.pinned_at,first_run_at=archived.first_run_at, + deadline_at=CASE WHEN archived.paused AND $3>0 THEN now()+make_interval(secs=>$3::double precision) + ELSE archived.deadline_at END, + deleted_at=NULL,archived_at=NULL,completed_at=archived.completed_at, + created_at=archived.created_at,updated_at=archived.updated_at +FROM json_populate_record(NULL::tasks,$2::json) archived +WHERE current.id=$1 AND archived.id=$1`, taskID, string(raw), remaining) + return err +} + +func firstArchiveRow(raw json.RawMessage) json.RawMessage { + var rows []json.RawMessage + if json.Unmarshal(raw, &rows) != nil || len(rows) == 0 { + return json.RawMessage("{}") + } + return rows[0] +} + +func decodeArchiveRows(raw json.RawMessage) ([]map[string]any, error) { + if len(raw) == 0 { + return nil, nil + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var rows []map[string]any + if err := decoder.Decode(&rows); err != nil { + return nil, err + } + return rows, nil +} + +func rowExists(tx *sql.Tx, table string, id int64) bool { + if table != "task_categories" && table != "llm_profiles" && table != "companies" && table != "tasks" && table != "assets" { + return false + } + var exists bool + _ = tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM `+table+` WHERE id=$1)`, id).Scan(&exists) + return exists +} + +func insertArchiveRows(tx *sql.Tx, table string, raw json.RawMessage) error { + allowed := map[string]bool{ + "exploration_nodes": true, "exploration_edges": true, "exploration_anchors": true, + "task_constraints": true, "activity": true, "task_asset_links": true, "findings": true, + "llm_records": true, "llm_usage": true, "skill_usage": true, "tool_usage": true, + } + if !allowed[table] { + return fmt.Errorf("archive restore table %q is not allowed", table) + } + if rawRowCount(raw) == 0 { + return nil + } + _, err := tx.Exec(`INSERT INTO `+table+` SELECT * FROM json_populate_recordset(NULL::`+table+`,$1::json)`, string(raw)) + return err +} + +func restoreArchiveAssets(tx *sql.Tx, taskID int64, raw json.RawMessage) (map[int64]int64, []string, error) { + rows, err := decodeArchiveRows(raw) + if err != nil { + return nil, nil, err + } + mapping := make(map[int64]int64, len(rows)) + warnings := []string{} + for _, row := range rows { + oldID, ok := jsonInt64(row["id"]) + if !ok || oldID <= 0 { + return nil, nil, errors.New("archived asset has invalid id") + } + if companyID, ok := jsonInt64(row["company_id"]); ok && companyID > 0 && !rowExists(tx, "companies", companyID) { + row["company_id"] = nil + row["company_source"] = "explicit" + warnings = append(warnings, fmt.Sprintf("资产 %d 的企业 %d 已删除，已恢复为未归属", oldID, companyID)) + } + if existing, found, err := findArchiveAssetNaturalID(tx, row); err != nil { + return nil, nil, err + } else if found { + mapping[oldID] = existing + if _, err := tx.Exec(`UPDATE assets SET task_ids=CASE WHEN $1=ANY(task_ids) THEN task_ids ELSE array_append(task_ids,$1) END WHERE id=$2`, taskID, existing); err != nil { + return nil, nil, err + } + continue + } + candidate := oldID + if rowExists(tx, "assets", oldID) { + if err := tx.QueryRow(`SELECT nextval(pg_get_serial_sequence('assets','id'))`).Scan(&candidate); err != nil { + return nil, nil, err + } + row["id"] = candidate + warnings = append(warnings, fmt.Sprintf("资产 ID %d 已被占用，恢复为 %d", oldID, candidate)) + } + row["task_ids"] = mergeJSONTaskID(row["task_ids"], taskID) + assetRaw, _ := json.Marshal(row) + res, err := tx.Exec(`INSERT INTO assets SELECT * FROM json_populate_record(NULL::assets,$1::json) ON CONFLICT DO NOTHING`, string(assetRaw)) + if err != nil { + return nil, nil, err + } + inserted, _ := res.RowsAffected() + if inserted == 0 { + existing, found, findErr := findArchiveAssetNaturalID(tx, row) + if findErr != nil || !found { + return nil, nil, errors.Join(findErr, fmt.Errorf("could not restore asset %d", oldID)) + } + candidate = existing + if _, err := tx.Exec(`UPDATE assets SET task_ids=CASE WHEN $1=ANY(task_ids) THEN task_ids ELSE array_append(task_ids,$1) END WHERE id=$2`, taskID, existing); err != nil { + return nil, nil, err + } + } + mapping[oldID] = candidate + } + return mapping, warnings, nil +} + +func mergeJSONTaskID(value any, taskID int64) []int64 { + out := []int64{} + seen := map[int64]bool{} + if values, ok := value.([]any); ok { + for _, item := range values { + if id, ok := jsonInt64(item); ok && id > 0 && !seen[id] { + seen[id] = true + out = append(out, id) + } + } + } + if !seen[taskID] { + out = append(out, taskID) + } + return out +} + +func findArchiveAssetNaturalID(tx *sql.Tx, row map[string]any) (int64, bool, error) { + typeName, _ := row["type"].(string) + stringValue := func(key string) string { + value, _ := row[key].(string) + return value + } + var query string + var args []any + switch typeName { + case "root_domain": + query, args = `SELECT id FROM assets WHERE type='root_domain' AND domain=$1`, []any{stringValue("domain")} + case "ip": + query, args = `SELECT id FROM assets WHERE type='ip' AND ip=$1`, []any{stringValue("ip")} + case "subdomain": + query, args = `SELECT id FROM assets WHERE type='subdomain' AND domain=$1 AND COALESCE(record_type,'')=$2`, []any{stringValue("domain"), stringValue("record_type")} + case "app": + if bundle := stringValue("bundle_id"); bundle != "" { + query, args = `SELECT id FROM assets WHERE type='app' AND bundle_id=$1`, []any{bundle} + } else { + query, args = `SELECT id FROM assets WHERE type='app' AND bundle_id IS NULL AND app_name=$1`, []any{stringValue("app_name")} + } + case "service": + if stringValue("service_type") == "http" { + query, args = `SELECT id FROM assets WHERE type='service' AND service_type='http' AND url=$1`, []any{stringValue("url")} + } else { + port, _ := jsonInt64(row["port"]) + query, args = `SELECT id FROM assets WHERE type='service' AND service_type='other' AND COALESCE(domain,'')=$1 AND COALESCE(ip,'')=$2 AND port=$3 AND service_name=$4`, []any{stringValue("domain"), stringValue("ip"), port, stringValue("service_name")} + } + case "endpoint": + query, args = `SELECT id FROM assets WHERE type='endpoint' AND url=$1 AND method=$2`, []any{stringValue("url"), stringValue("method")} + default: + return 0, false, fmt.Errorf("unsupported archived asset type %q", typeName) + } + var id int64 + err := tx.QueryRow(query, args...).Scan(&id) + if errors.Is(err, sql.ErrNoRows) { + return 0, false, nil + } + return id, err == nil, err +} + +func remapArchiveAssetReferences(tables map[string]json.RawMessage, mapping map[int64]int64) (map[string]json.RawMessage, error) { + out := make(map[string]json.RawMessage, len(tables)) + for name, raw := range tables { + out[name] = raw + } + for _, table := range []string{"exploration_anchors", "task_asset_links"} { + rows, err := decodeArchiveRows(tables[table]) + if err != nil { + return nil, err + } + for _, row := range rows { + if old, ok := jsonInt64(row["asset_id"]); ok { + if replacement, exists := mapping[old]; exists { + row["asset_id"] = replacement + } + } + } + out[table], _ = json.Marshal(rows) + } + for _, table := range []string{"exploration_nodes", "findings"} { + rows, err := decodeArchiveRows(tables[table]) + if err != nil { + return nil, err + } + for _, row := range rows { + key := "asset_ids" + container := row + if table == "exploration_nodes" { + payload, ok := row["payload"].(map[string]any) + if !ok { + continue + } + container = payload + } + if values, ok := container[key].([]any); ok { + for i, value := range values { + if old, ok := jsonInt64(value); ok { + if replacement, exists := mapping[old]; exists { + values[i] = replacement + } + } + } + container[key] = values + } + } + out[table], _ = json.Marshal(rows) + } + return out, nil +} + +func restoreTaskRelations(tx *sql.Tx, taskID int64, raw json.RawMessage) ([]string, error) { + rows, err := decodeArchiveRows(raw) + if err != nil { + return nil, err + } + warnings := []string{} + for _, row := range rows { + sourceID, ok := jsonInt64(row["source_task_id"]) + if !ok || !liveTaskExists(tx, sourceID) { + warnings = append(warnings, fmt.Sprintf("来源任务 %d 不可用，继承关系已跳过", sourceID)) + continue + } + created, _ := row["created_at"].(string) + if _, err := tx.Exec(`INSERT INTO task_relations(task_id,source_task_id,created_at) VALUES($1,$2,COALESCE($3::timestamptz,now())) ON CONFLICT DO NOTHING`, taskID, sourceID, nilIfEmptyString(created)); err != nil { + return nil, err + } + } + return warnings, nil +} + +func restoreTaskScopes(tx *sql.Tx, raw json.RawMessage) ([]string, error) { + rows, err := decodeArchiveRows(raw) + if err != nil { + return nil, err + } + warnings := []string{} + kept := make([]map[string]any, 0, len(rows)) + for _, row := range rows { + if companyID, ok := jsonInt64(row["company_id"]); ok && companyID > 0 && !rowExists(tx, "companies", companyID) { + warnings = append(warnings, fmt.Sprintf("企业 %d 已删除，关联范围已跳过", companyID)) + continue + } + kept = append(kept, row) + } + if len(kept) == 0 { + return warnings, nil + } + encoded, _ := json.Marshal(kept) + _, err = tx.Exec(`INSERT INTO task_scope SELECT * FROM json_populate_recordset(NULL::task_scope,$1::json) ON CONFLICT DO NOTHING`, string(encoded)) + return warnings, err +} + +func restoreTaskLLMProfiles(tx *sql.Tx, raw json.RawMessage) ([]string, error) { + rows, err := decodeArchiveRows(raw) + if err != nil { + return nil, err + } + warnings := []string{} + position := 0 + for _, row := range rows { + profileID, ok := jsonInt64(row["profile_id"]) + if !ok || !rowExists(tx, "llm_profiles", profileID) { + warnings = append(warnings, fmt.Sprintf("LLM 配置 %d 已删除，配置链项已跳过", profileID)) + continue + } + taskID, _ := jsonInt64(row["task_id"]) + status, _ := row["status"].(string) + lastError, _ := row["last_error"].(string) + exhaustedAt, _ := row["exhausted_at"].(string) + createdAt, _ := row["created_at"].(string) + updatedAt, _ := row["updated_at"].(string) + if _, err := tx.Exec(`INSERT INTO task_llm_profiles(task_id,profile_id,position,status,last_error,exhausted_at,created_at,updated_at) + VALUES($1,$2,$3,$4,NULLIF($5,''),$6::timestamptz,COALESCE($7::timestamptz,now()),COALESCE($8::timestamptz,now())) + ON CONFLICT DO NOTHING`, taskID, profileID, position, status, lastError, + nilIfEmptyString(exhaustedAt), nilIfEmptyString(createdAt), nilIfEmptyString(updatedAt)); err != nil { + return nil, err + } + position++ + } + return warnings, nil +} + +func restoreInterceptRows(tx *sql.Tx, raw json.RawMessage) error { + rows, err := decodeArchiveRows(raw) + if err != nil { + return err + } + for _, row := range rows { + if status, _ := row["status"].(string); status == "pending" { + row["status"] = "timeout" + row["reason"] = "任务归档期间审批已超时" + row["decided_at"] = time.Now().UTC() + } + } + if len(rows) == 0 { + return nil + } + encoded, _ := json.Marshal(rows) + _, err = tx.Exec(`INSERT INTO intercept_pending SELECT * FROM json_populate_recordset(NULL::intercept_pending,$1::json)`, string(encoded)) + return err +} + +func liveTaskExists(tx *sql.Tx, id int64) bool { + var exists bool + _ = tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, id).Scan(&exists) + return exists +} + +func nilIfEmptyString(value string) any { + if strings.TrimSpace(value) == "" { + return nil + } + return value +} + +// CompleteTaskArchiveRestore removes compact metadata after every external +// component has been verified and the package has been consumed. +func (d *DB) CompleteTaskArchiveRestore(archiveID int64) error { + res, err := d.Exec(`DELETE FROM task_archives archive USING tasks task +WHERE archive.id=$1 AND archive.task_id=task.id AND task.deleted_at IS NULL AND task.archived_at IS NULL`, archiveID) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n != 1 { + return ErrTaskArchiveState + } + return nil +} + +// DeleteTaskArchiveStub permanently removes the cold task after its package has +// been staged for deletion. Dependency protection is rechecked transactionally. +func (d *DB) DeleteTaskArchiveStub(archiveID int64) error { + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck + var taskID, expID int64 + var state string + if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state +FROM task_archives archive JOIN tasks task ON task.id=archive.task_id +WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID, &state); err != nil { + return err + } + if state != Deleting { + return ErrTaskArchiveState + } + var dependent int64 + err = tx.QueryRow(`SELECT task_id FROM task_archives WHERE id<>$1 AND $2=ANY(source_task_ids) LIMIT 1`, archiveID, taskID).Scan(&dependent) + if err == nil { + return fmt.Errorf("%w: task %d", ErrTaskArchiveDeleteBlocked, dependent) + } + if !errors.Is(err, sql.ErrNoRows) { + return err + } + if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1 AND archived_at IS NOT NULL`, taskID); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil { + return err + } + return tx.Commit() +} + +// ArchivedAggregateStats returns compact summaries used by global dashboards so +// cold data does not disappear from historical totals. +func (d *DB) ArchivedAggregateStats() ([]json.RawMessage, error) { + rows, err := d.Query(`SELECT archive.aggregate_stats +FROM task_archives archive +JOIN tasks task ON task.id=archive.task_id +WHERE task.archived_at IS NOT NULL`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []json.RawMessage + for rows.Next() { + var raw json.RawMessage + if err := rows.Scan(&raw); err != nil { + return nil, err + } + out = append(out, raw) + } + return out, rows.Err() +} diff --git a/db/task_archives_test.go b/db/task_archives_test.go new file mode 100644 index 0000000..dbe282f --- /dev/null +++ b/db/task_archives_test.go @@ -0,0 +1,297 @@ +package db + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + "time" +) + +func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + if err := d.EnsureLLMRecordsTable(); err != nil { + t.Fatal(err) + } + if err := d.EnsureLLMUsageTable(); err != nil { + t.Fatal(err) + } + task, err := d.CreateTaskWithOptions("archive database roundtrip", "restore exact graph", TaskCreateOptions{Name: "cold task"}) + if err != nil { + t.Fatal(err) + } + var companyID, llmProfileID int64 + if err := d.QueryRow(`INSERT INTO companies(name,nkey) VALUES($1,$2) RETURNING id`, + fmt.Sprintf("archive-company-%d", task.ID), fmt.Sprintf("archive-company-%d", task.ID)).Scan(&companyID); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`INSERT INTO llm_profiles(name,format,model) VALUES($1,'openai','archive-model') RETURNING id`, + fmt.Sprintf("archive-chain-%d", task.ID)).Scan(&llmProfileID); err != nil { + t.Fatal(err) + } + exhaustedAt := time.Now().UTC().Add(-time.Hour).Truncate(time.Microsecond) + chainCreatedAt := exhaustedAt.Add(-time.Hour) + if _, err := d.Exec(`UPDATE tasks SET company_id=$2,llm_profile_id=$3,active_llm_profile_id=$3 WHERE id=$1`, task.ID, companyID, llmProfileID); err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`INSERT INTO task_llm_profiles(task_id,profile_id,position,status,last_error,exhausted_at,created_at,updated_at) +VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmProfileID, exhaustedAt, chainCreatedAt); err != nil { + t.Fatal(err) + } + defer func() { + _, _ = d.Exec(`DELETE FROM llm_usage WHERE task_id=$1 OR exploration_id=$2`, fmt.Sprint(task.ID), task.ExplorationID) + _, _ = d.Exec(`DELETE FROM skill_usage WHERE task_id=$1 OR exploration_id=$2`, task.ID, task.ExplorationID) + _, _ = d.Exec(`DELETE FROM tool_usage WHERE task_id=$1 OR exploration_id=$2`, task.ID, task.ExplorationID) + _, _ = d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID) + _ = d.DeleteTask(task.ID) + _, _ = d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, llmProfileID) + _, _ = d.Exec(`DELETE FROM companies WHERE id=$1`, companyID) + }() + if err := d.SetPaused(task.ID, true); err != nil { + t.Fatal(err) + } + assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: fmt.Sprintf("archive-%d.example", task.ID), TaskID: task.ID}) + if err != nil { + t.Fatal(err) + } + store := d.Exploration(task.ExplorationID) + nodeID, err := store.AddNode(KindFact, map[string]any{"summary": "archived fact", "asset_ids": []int64{assetID}}, 1, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + if err := store.Anchor(nodeID, assetID); err != nil { + t.Fatal(err) + } + profileName := fmt.Sprintf("archive-profile-%d", task.ID) + skillName := fmt.Sprintf("archive-skill-%d", task.ID) + toolName := fmt.Sprintf("archive-tool-%d", task.ID) + vulnclass := fmt.Sprintf("archive-vuln-%d", task.ID) + if err := d.InsertLLMUsage(&LLMUsage{TaskID: fmt.Sprint(task.ID), ExplorationID: task.ExplorationID, Worker: "worker", Model: "test", ProfileName: profileName, InputTokens: 11, OutputTokens: 7}); err != nil { + t.Fatal(err) + } + if err := d.InsertSkillUsage(&SkillUsage{Skill: skillName, AgentKey: "worker", TaskID: task.ID, ExplorationID: task.ExplorationID, Found: true}); err != nil { + t.Fatal(err) + } + if err := d.InsertToolUsage(&ToolUsage{ToolKey: toolName, AgentKey: "worker", TaskID: task.ID, ExplorationID: task.ExplorationID}); err != nil { + t.Fatal(err) + } + if _, err := d.AddFinding(task.ID, nodeID, vulnclass, "archive finding", SeverityCritical, "summary", "evidence", "worker", []int64{assetID}); err != nil { + t.Fatal(err) + } + archive, err := d.QueueTaskArchive(task.ID) + if err != nil { + t.Fatal(err) + } + claimed, err := d.ClaimTaskArchiveJob(t.Context()) + if err != nil || claimed == nil || claimed.ID != archive.ID || claimed.State != Archiving { + t.Fatalf("claim = %+v, %v", claimed, err) + } + snapshot, err := d.SnapshotTaskArchive(task.ID) + if err != nil { + t.Fatal(err) + } + if snapshot.DataCounts["assets"] != 1 || snapshot.DataCounts["exploration_nodes"] < 2 { + t.Fatalf("unexpected snapshot counts: %#v", snapshot.DataCounts) + } + if err := d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/test-task.tar.zst", "abc", 100, 50); err != nil { + t.Fatal(err) + } + if live, err := d.GetTask(task.ID); err != nil || live != nil { + t.Fatalf("archived task must be hidden, got %+v, %v", live, err) + } + var nodes, assets, usage int + if err := d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodes); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE id=$1`, assetID).Scan(&assets); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT count(*) FROM llm_usage WHERE task_id=$1`, fmt.Sprint(task.ID)).Scan(&usage); err != nil { + t.Fatal(err) + } + if nodes != 0 || assets != 0 || usage != 0 { + t.Fatalf("hot compaction left nodes=%d assets=%d usage=%d", nodes, assets, usage) + } + ready, err := d.GetTaskArchive(archive.ID) + if err != nil || ready == nil || ready.State != ArchiveReady { + t.Fatalf("ready archive = %+v, %v", ready, err) + } + var stats map[string]any + if err := json.Unmarshal(ready.AggregateStats, &stats); err != nil { + t.Fatal(err) + } + if _, ok := stats["tokens"]; !ok { + t.Fatalf("archive token summary missing: %#v", stats) + } + assertArchiveGlobalStats(t, d, profileName, skillName, toolName, vulnclass) + if _, err := d.Exec(`DELETE FROM companies WHERE id=$1`, companyID); err != nil { + t.Fatal(err) + } + if _, err := d.QueueTaskArchiveRestore(archive.ID); err != nil { + t.Fatal(err) + } + claimed, err = d.ClaimTaskArchiveJob(t.Context()) + if err != nil || claimed == nil || claimed.State != Restoring { + t.Fatalf("restore claim = %+v, %v", claimed, err) + } + warnings, err := d.RestoreTaskArchive(archive.ID, snapshot, int64((10 * time.Minute).Seconds())) + if err != nil { + t.Fatal(err) + } + warningFound := false + for _, warning := range warnings { + if strings.Contains(warning, fmt.Sprintf("任务企业 %d 已删除", companyID)) { + warningFound = true + } + } + if !warningFound { + t.Fatalf("missing deleted company warning: %v", warnings) + } + live, err := d.GetTask(task.ID) + if err != nil || live == nil { + t.Fatalf("restored task = %+v, %v", live, err) + } + if live.Name != "cold task" || !live.Paused { + t.Fatalf("restored task metadata mismatch: %+v", live) + } + if err := d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodes); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE id=$1 AND $2=ANY(task_ids)`, assetID, task.ID).Scan(&assets); err != nil { + t.Fatal(err) + } + if err := d.QueryRow(`SELECT count(*) FROM llm_usage WHERE task_id=$1`, fmt.Sprint(task.ID)).Scan(&usage); err != nil { + t.Fatal(err) + } + if nodes < 2 || assets != 1 || usage != 1 { + t.Fatalf("restore incomplete nodes=%d assets=%d usage=%d", nodes, assets, usage) + } + var restoredCompanyID *int64 + var restoredStatus, restoredError string + var restoredExhaustedAt, restoredCreatedAt time.Time + if err := d.QueryRow(`SELECT company_id FROM tasks WHERE id=$1`, task.ID).Scan(&restoredCompanyID); err != nil { + t.Fatal(err) + } + if restoredCompanyID != nil { + t.Fatalf("deleted legacy company restored as %v", *restoredCompanyID) + } + if err := d.QueryRow(`SELECT status,COALESCE(last_error,''),exhausted_at,created_at FROM task_llm_profiles WHERE task_id=$1 AND profile_id=$2`, + task.ID, llmProfileID).Scan(&restoredStatus, &restoredError, &restoredExhaustedAt, &restoredCreatedAt); err != nil { + t.Fatal(err) + } + if restoredStatus != "quota_exhausted" || restoredError != "balance exhausted" || + !restoredExhaustedAt.Equal(exhaustedAt) || !restoredCreatedAt.Equal(chainCreatedAt) { + t.Fatalf("LLM chain state/time mismatch: status=%s error=%q exhausted=%s created=%s", + restoredStatus, restoredError, restoredExhaustedAt, restoredCreatedAt) + } + assertArchiveGlobalStats(t, d, profileName, skillName, toolName, vulnclass) + if err := d.CompleteTaskArchiveRestore(archive.ID); err != nil { + t.Fatal(err) + } + if item, err := d.GetTaskArchive(archive.ID); err != nil || item != nil { + t.Fatalf("archive metadata must be consumed: %+v, %v", item, err) + } +} + +func TestTaskArchiveBlockersIgnoreQueuedDependents(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + source, err := d.CreateTaskWithOptions("archive blocker source", "source", TaskCreateOptions{}) + if err != nil { + t.Fatal(err) + } + dependent, err := d.CreateTaskWithOptions("archive blocker dependent", "dependent", TaskCreateOptions{ + SourceTaskIDs: []int64{source.ID}, + }) + if err != nil { + _ = d.DeleteTask(source.ID) + t.Fatal(err) + } + defer func() { + _, _ = d.Exec(`DELETE FROM task_archives WHERE task_id IN ($1,$2)`, source.ID, dependent.ID) + _ = d.DeleteTask(dependent.ID) + _ = d.DeleteTask(source.ID) + }() + + blockers, err := d.TaskArchiveBlockers() + if err != nil { + t.Fatal(err) + } + if blockers[source.ID] != dependent.ID { + t.Fatalf("source blocker=%d, want dependent %d", blockers[source.ID], dependent.ID) + } + if err := d.SetPaused(dependent.ID, true); err != nil { + t.Fatal(err) + } + if _, err := d.QueueTaskArchive(dependent.ID); err != nil { + t.Fatal(err) + } + blockers, err = d.TaskArchiveBlockers() + if err != nil { + t.Fatal(err) + } + if blocker := blockers[source.ID]; blocker != 0 { + t.Fatalf("queued dependent must not block source, got %d", blocker) + } +} + +func assertArchiveGlobalStats(t *testing.T, d *DB, profileName, skillName, toolName, vulnclass string) { + t.Helper() + profiles, err := d.UsageByProfile() + if err != nil { + t.Fatal(err) + } + matchedProfile := false + for _, profile := range profiles { + if profile.ProfileName != profileName { + continue + } + matchedProfile = true + if profile.Calls != 1 || profile.Tasks != 1 || profile.InputTokens != 11 || profile.OutputTokens != 7 { + t.Fatalf("archive profile aggregate double-counted or missing: %+v", profile) + } + } + if !matchedProfile { + t.Fatalf("archive profile aggregate %q missing", profileName) + } + skills, err := d.SkillStats() + if err != nil { + t.Fatal(err) + } + matchedSkill := false + for _, skill := range skills { + if skill.Skill == skillName { + matchedSkill = skill.Calls == 1 && skill.Tasks == 1 + } + } + if !matchedSkill { + t.Fatalf("archive skill aggregate %q missing or double-counted: %+v", skillName, skills) + } + tools, err := d.ToolUsageCounts() + if err != nil { + t.Fatal(err) + } + if tools[toolName] != 1 { + t.Fatalf("archive tool aggregate %q=%d, want 1", toolName, tools[toolName]) + } + findings, err := d.FindingStats() + if err != nil { + t.Fatal(err) + } + foundClass := false + for _, item := range findings.VulnClasses { + if item == vulnclass { + foundClass = true + } + } + if !foundClass { + t.Fatalf("archive finding class %q missing", vulnclass) + } +} diff --git a/db/tool_usage.go b/db/tool_usage.go index 8c2632b..53c2280 100644 --- a/db/tool_usage.go +++ b/db/tool_usage.go @@ -42,5 +42,17 @@ GROUP BY tool_key`) } out[key] = calls } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + archived, err := d.archivedTaskAggregates() + if err != nil { + return nil, err + } + for _, aggregate := range archived { + for key, calls := range aggregate.Tools { + out[key] += calls + } + } + return out, nil } diff --git a/server/dto.go b/server/dto.go index 8db7ba7..cf7133d 100644 --- a/server/dto.go +++ b/server/dto.go @@ -56,10 +56,24 @@ type TaskDTO struct { LLMFailoverState string `json:"llm_failover_state"` LLMFailoverReason string `json:"llm_failover_reason,omitempty"` SourceTaskIDs []string `json:"source_task_ids"` + ArchiveBlockedBy string `json:"archive_blocked_by_task_id,omitempty"` CompanyIDs []int64 `json:"company_ids"` CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定) } +func applyTaskArchiveBlocker(dto *TaskDTO, blockers map[int64]int64) { + if dto == nil || len(blockers) == 0 { + return + } + taskID, err := strconv.ParseInt(dto.ID, 10, 64) + if err != nil { + return + } + if dependentID := blockers[taskID]; dependentID > 0 { + dto.ArchiveBlockedBy = strconv.FormatInt(dependentID, 10) + } +} + // TokenTotalDTO is a whole-task (all agents) token aggregate. type TokenTotalDTO struct { InputTokens int `json:"input_tokens"` diff --git a/server/findings_groups.go b/server/findings_groups.go index 86f0972..e370069 100644 --- a/server/findings_groups.go +++ b/server/findings_groups.go @@ -37,6 +37,7 @@ func (s *Server) findingGroups(w http.ResponseWriter, r *http.Request) { Status: normFilter(q.Get("status")), VulnClass: normFilter(q.Get("vulnclass")), TaskID: normFilter(q.Get("task_id")), + Query: q.Get("q"), Sort: q.Get("sort"), } groups, total, findingTotal, err := s.m.pg.ListFindingGroups(filter, page, limit) diff --git a/server/findings_groups_test.go b/server/findings_groups_test.go index 5cb44ed..61c05b6 100644 --- a/server/findings_groups_test.go +++ b/server/findings_groups_test.go @@ -102,6 +102,28 @@ func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) { t.Fatalf("unexpected first group: %+v", page.Items[0]) } + request = httptest.NewRequest(http.MethodGet, + "/api/exploration/findings/groups?vulnclass="+marker+"&q=finding+2&page=1&limit=10", nil) + recorder = httptest.NewRecorder() + s.findingGroups(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("searched groups status=%d body=%s", recorder.Code, recorder.Body.String()) + } + page = struct { + Items []db.FindingGroup `json:"items"` + Total int `json:"total"` + FindingTotal int `json:"finding_total"` + Page int `json:"page"` + PageSize int `json:"page_size"` + }{} + if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil { + t.Fatal(err) + } + if len(page.Items) != 1 || page.Total != 1 || page.FindingTotal != 1 || page.Items[0].TaskID == nil || + *page.Items[0].TaskID != mustTaskID(t, taskB.ID) { + t.Fatalf("unexpected searched group page: %+v", page) + } + request = httptest.NewRequest(http.MethodGet, "/api/exploration/findings/groups?vulnclass="+marker+"&page=-4&limit=0", nil) recorder = httptest.NewRecorder() diff --git a/server/manager.go b/server/manager.go index 7902ce5..cb198c2 100644 --- a/server/manager.go +++ b/server/manager.go @@ -358,6 +358,25 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { if err != nil { log.Printf("[traffic] disabled: %v", err) } else { + err = tr.RecoverHostDeleteStages(func(_ int64, taskID int64) (bool, error) { + if taskID <= 0 { + return false, errors.New("归档流量暂存日志缺少任务 ID") + } + task, taskErr := pg.GetTask(taskID) + if taskErr != nil { + return false, taskErr + } + // Archived/permanently deleted tasks are hidden from GetTask. A + // restored task is visible and needs the staged traffic put back. + return task == nil, nil + }) + if err != nil { + _ = tr.Close() + _ = pg.Close() + return nil, fmt.Errorf("recover traffic delete staging: %w", err) + } + } + if tr != nil { m.traffic = tr go func() { log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr()) diff --git a/server/server.go b/server/server.go index e46755a..942c8ca 100644 --- a/server/server.go +++ b/server/server.go @@ -116,6 +116,11 @@ type Server struct { // the persisted current profile at every new LLM call. taskAgentMu sync.Mutex taskAgents map[string]*taskAgentBundle + + // Cold task archives run through one persistent FIFO worker. The buffered wake + // channel coalesces enqueue bursts; the database remains the source of truth. + archiveWake chan struct{} + archiveWG sync.WaitGroup } // profBundle is a planner/worker pair built from one LLM profile. @@ -151,7 +156,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi 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), - taskAgents: map[string]*taskAgentBundle{}} + taskAgents: map[string]*taskAgentBundle{}, archiveWake: make(chan struct{}, 1)} // 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. @@ -264,6 +269,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi } } go s.reconcileConcurrency() + s.startTaskArchiveWorker() s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定 return s } @@ -690,6 +696,14 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("PATCH /api/tasks/{id}", s.updateTaskMetadata) mux.HandleFunc("PATCH /api/tasks/{id}/category", s.updateTaskCategory) mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch) + mux.HandleFunc("GET /api/task-archives", s.listTaskArchives) + mux.HandleFunc("GET /api/task-archives/{id}", s.getTaskArchive) + mux.HandleFunc("POST /api/tasks/{id}/archive", s.queueTaskArchive) + mux.HandleFunc("POST /api/tasks/archive/batch", s.queueTaskArchivesBatch) + mux.HandleFunc("POST /api/task-archives/{id}/restore", s.queueTaskArchiveRestore) + mux.HandleFunc("POST /api/task-archives/restore/batch", s.restoreTaskArchivesBatch) + mux.HandleFunc("DELETE /api/task-archives/{id}", s.queueTaskArchiveDelete) + mux.HandleFunc("POST /api/task-archives/delete/batch", s.deleteTaskArchivesBatch) mux.HandleFunc("GET /api/tasks/{id}/coverage", s.taskCoverage) mux.HandleFunc("GET /api/tasks/{id}/coverage-graph", s.taskCoverageGraph) mux.HandleFunc("GET /api/tasks/{id}/asset-refs", s.taskAssetRefs) @@ -996,9 +1010,11 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { } list := s.m.List() metrics, _ := s.m.PG().TaskListMetricsAll() + archiveBlockers, _ := s.m.PG().TaskArchiveBlockers() dtos := make([]TaskDTO, 0, len(list)) for _, t := range list { dto := taskDTO(t, s.resolvedTaskStatus(t)) + applyTaskArchiveBlocker(&dto, archiveBlockers) metric := metrics[t.ExpID] dto.Tokens = tokenTotalDTO(metric.Tokens) // prefer the live in-memory heartbeat (fresher) and fall back to the @@ -1640,7 +1656,10 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "task not found") return } - writeJSON(w, 200, taskDTO(t, s.resolvedTaskStatus(t))) + dto := taskDTO(t, s.resolvedTaskStatus(t)) + archiveBlockers, _ := s.m.PG().TaskArchiveBlockers() + applyTaskArchiveBlocker(&dto, archiveBlockers) + writeJSON(w, 200, dto) } // taskCoverage returns a task's rough asset test coverage (denominator/tested/backlog). @@ -1842,6 +1861,7 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) { VulnClass: normFilter(q.Get("vulnclass")), // task_id(独立于会切到「按任务节点」分支的 task 参数):全局表按任务筛选。 TaskID: normFilter(q.Get("task_id")), + Query: q.Get("q"), Sort: q.Get("sort"), } fs, total, err := s.m.pg.ListFindingsPage(filter, page, limit) @@ -2011,7 +2031,7 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) { // format = md-single（整合一份 .md）| md-zip（一漏洞一 .md,打包 zip） // | csv | json // ids = 逗号分隔的 finding id（scope=selected 时必填） -// 筛选参数 severity/status/vulnclass/task_id/sort 与列表接口一致（scope=filtered 用）。 +// 筛选参数 severity/status/vulnclass/task_id/q/sort 与列表接口一致（scope=filtered 用）。 func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { q := r.URL.Query() scope := q.Get("scope") @@ -2045,6 +2065,7 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { Status: normFilter(q.Get("status")), VulnClass: normFilter(q.Get("vulnclass")), TaskID: normFilter(q.Get("task_id")), + Query: q.Get("q"), Sort: q.Get("sort"), } default: diff --git a/server/task_archive_package.go b/server/task_archive_package.go new file mode 100644 index 0000000..ce4494d --- /dev/null +++ b/server/task_archive_package.go @@ -0,0 +1,655 @@ +package server + +import ( + "archive/tar" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + + pgdb "github.com/Autumn-27/artex/db" + "github.com/klauspost/compress/zstd" +) + +const ( + archiveDirMode = 0o700 + archiveFileMode = 0o600 + maxArchiveFiles = 1_000_000 + maxArchiveBytes = int64(1 << 47) // 128 TiB safety ceiling for corrupt headers. +) + +type archiveFileMove struct { + Source string `json:"source"` + Relative string `json:"relative"` +} + +type archiveStageJournal struct { + ArchiveID int64 `json:"archive_id"` + TaskID string `json:"task_id"` + Moves []archiveFileMove `json:"moves"` +} + +type archiveRestoreJournal struct { + ArchiveID int64 `json:"archive_id"` + TaskID int64 `json:"task_id"` + Moves []restoredArchivePath `json:"moves"` +} + +type taskArchiveFileStage struct { + root string + payload string + journal archiveStageJournal + done bool +} + +type restoredArchivePath struct { + Source string + Destination string +} + +type taskArchiveRestoreFiles struct { + extracted string + moves []restoredArchivePath + done bool +} + +func writeArchiveJournal(path string, value any) error { + raw, err := json.Marshal(value) + if err != nil { + return err + } + file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, archiveFileMode) + if err != nil { + return err + } + if _, err := file.Write(raw); err != nil { + _ = file.Close() + return err + } + if err := file.Sync(); err != nil { + _ = file.Close() + return err + } + return file.Close() +} + +func taskArchiveRoot(dataDir string) string { + return filepath.Join(dataDir, "archives", "tasks") +} + +func taskArchivePath(dataDir string, archiveID int64, taskID string) string { + return filepath.Join(taskArchiveRoot(dataDir), fmt.Sprintf("task-%s-%d.tar.zst", taskID, archiveID)) +} + +func stageTaskArchiveFiles(dataDir string, archiveID int64, taskID string, explorationID int64) (*taskArchiveFileStage, error) { + root := filepath.Join(taskArchiveRoot(dataDir), ".staging", strconv.FormatInt(archiveID, 10)) + stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload")} + journalPath := filepath.Join(root, "journal.json") + if raw, err := os.ReadFile(journalPath); err == nil { + if err := json.Unmarshal(raw, &stage.journal); err != nil { + return nil, fmt.Errorf("read task archive staging journal: %w", err) + } + return stage, nil + } else if !os.IsNotExist(err) { + return nil, err + } + if err := os.MkdirAll(filepath.Join(stage.payload, "files"), archiveDirMode); err != nil { + return nil, err + } + stage.journal = archiveStageJournal{ArchiveID: archiveID, TaskID: taskID} + targets := []archiveFileMove{} + taskDir := filepath.Join(dataDir, "tasks", taskID) + if _, err := os.Lstat(taskDir); err == nil { + targets = append(targets, archiveFileMove{Source: taskDir, Relative: filepath.Join("files", "tasks", taskID)}) + } else if !os.IsNotExist(err) { + return nil, err + } + transcriptDir := filepath.Join(dataDir, "transcripts") + entries, err := os.ReadDir(transcriptDir) + if err != nil && !os.IsNotExist(err) { + return nil, err + } + 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, archiveFileMove{ + Source: filepath.Join(transcriptDir, name), Relative: filepath.Join("files", "transcripts", name), + }) + } + // Persist the complete plan before the first rename. Recovery can therefore + // roll back any prefix of the moves after an abrupt process termination. + stage.journal.Moves = targets + if err := writeArchiveJournal(journalPath, stage.journal); err != nil { + _ = stage.rollback() + return nil, err + } + for _, move := range targets { + destination := filepath.Join(stage.payload, move.Relative) + if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil { + _ = stage.rollback() + return nil, err + } + if err := os.Rename(move.Source, destination); err != nil { + _ = stage.rollback() + return nil, fmt.Errorf("stage task archive path %s: %w", move.Source, err) + } + } + return stage, nil +} + +func (s *taskArchiveFileStage) rollback() error { + if s == nil || s.done { + return nil + } + var errs []error + for i := len(s.journal.Moves) - 1; i >= 0; i-- { + move := s.journal.Moves[i] + staged := filepath.Join(s.payload, move.Relative) + if _, err := os.Lstat(staged); os.IsNotExist(err) { + continue + } else if err != nil { + errs = append(errs, err) + continue + } + if _, err := os.Lstat(move.Source); err == nil { + errs = append(errs, fmt.Errorf("archive rollback destination exists: %s", move.Source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, err) + continue + } + if err := os.MkdirAll(filepath.Dir(move.Source), 0o755); err != nil { + errs = append(errs, err) + continue + } + if err := os.Rename(staged, move.Source); err != nil { + errs = append(errs, err) + } + } + if len(errs) == 0 { + if err := os.RemoveAll(s.root); err != nil { + errs = append(errs, err) + } + } + s.done = true + return errors.Join(errs...) +} + +func (s *taskArchiveFileStage) commit() error { + if s == nil || s.done { + return nil + } + s.done = true + return os.RemoveAll(s.root) +} + +func installTaskArchiveFiles(dataDir, extractedDir, taskID string, archiveID int64) (*taskArchiveRestoreFiles, error) { + stage := &taskArchiveRestoreFiles{extracted: extractedDir} + numericTaskID, err := strconv.ParseInt(taskID, 10, 64) + if err != nil || numericTaskID <= 0 { + return nil, fmt.Errorf("invalid restore task id %q", taskID) + } + sources := []restoredArchivePath{} + workspace := filepath.Join(extractedDir, "files", "tasks", taskID) + if _, err := os.Lstat(workspace); err == nil { + sources = append(sources, restoredArchivePath{Source: workspace, Destination: filepath.Join(dataDir, "tasks", taskID)}) + } else if !os.IsNotExist(err) { + return nil, err + } + transcripts := filepath.Join(extractedDir, "files", "transcripts") + entries, err := os.ReadDir(transcripts) + if err != nil && !os.IsNotExist(err) { + return nil, err + } + for _, entry := range entries { + sources = append(sources, restoredArchivePath{ + Source: filepath.Join(transcripts, entry.Name()), Destination: filepath.Join(dataDir, "transcripts", entry.Name()), + }) + } + for _, move := range sources { + if _, err := os.Lstat(move.Destination); err == nil { + return nil, fmt.Errorf("restore destination already exists: %s", move.Destination) + } else if !os.IsNotExist(err) { + return nil, err + } + } + stage.moves = sources + journal := archiveRestoreJournal{ArchiveID: archiveID, TaskID: numericTaskID, Moves: sources} + if err := writeArchiveJournal(filepath.Join(extractedDir, "restore-journal.json"), journal); err != nil { + return nil, err + } + for _, move := range sources { + if err := os.MkdirAll(filepath.Dir(move.Destination), 0o755); err != nil { + _ = stage.rollback() + return nil, err + } + if err := os.Rename(move.Source, move.Destination); err != nil { + _ = stage.rollback() + return nil, err + } + } + return stage, nil +} + +func (s *taskArchiveRestoreFiles) 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.Destination); os.IsNotExist(err) { + continue + } else if err != nil { + errs = append(errs, err) + continue + } + if _, err := os.Lstat(move.Source); err == nil { + errs = append(errs, fmt.Errorf("restore rollback source exists: %s", move.Source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, err) + continue + } + if err := os.MkdirAll(filepath.Dir(move.Source), archiveDirMode); err != nil { + errs = append(errs, err) + continue + } + if err := os.Rename(move.Destination, move.Source); err != nil { + errs = append(errs, err) + } + } + s.done = true + return errors.Join(errs...) +} + +func (s *taskArchiveRestoreFiles) commit() error { + if s == nil || s.done { + return nil + } + s.done = true + return os.RemoveAll(s.extracted) +} + +// recoverTaskArchiveRestoreStages resolves file installs left by an interrupted +// restore. If PostgreSQL committed, installed files are authoritative and only +// the extraction directory is stale. Otherwise all completed renames are moved +// back so the persistent restore job can retry from a clean destination. +func recoverTaskArchiveRestoreStages(dataDir string, pg *pgdb.DB) error { + parent := filepath.Join(taskArchiveRoot(dataDir), ".restore") + entries, err := os.ReadDir(parent) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + var errs []error + for _, entry := range entries { + if !entry.IsDir() { + continue + } + root := filepath.Join(parent, entry.Name()) + raw, err := os.ReadFile(filepath.Join(root, "restore-journal.json")) + if os.IsNotExist(err) { + // Extraction was interrupted before any destination rename. + if err := os.RemoveAll(root); err != nil { + errs = append(errs, err) + } + continue + } + if err != nil { + errs = append(errs, err) + continue + } + var journal archiveRestoreJournal + if err := json.Unmarshal(raw, &journal); err != nil { + errs = append(errs, err) + continue + } + restored, err := pg.IsTaskArchiveRestored(journal.ArchiveID) + if errors.Is(err, pgdb.ErrTaskArchiveNotFound) && journal.TaskID > 0 { + task, taskErr := pg.GetTask(journal.TaskID) + if taskErr != nil { + errs = append(errs, taskErr) + continue + } + restored = task != nil + err = nil + } + if err != nil { + errs = append(errs, err) + continue + } + if restored { + if err := os.RemoveAll(root); err != nil { + errs = append(errs, err) + } + continue + } + stage := &taskArchiveRestoreFiles{extracted: root, moves: journal.Moves} + if err := stage.rollback(); err != nil { + errs = append(errs, err) + continue + } + if err := os.RemoveAll(root); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +// recoverTaskArchiveStages resolves file renames left by an interrupted archive. +// A committed cold task already has a verified package, so stale staging can be +// discarded; otherwise files are moved back before the persistent job retries. +func recoverTaskArchiveStages(dataDir string, pg *pgdb.DB) error { + parent := filepath.Join(taskArchiveRoot(dataDir), ".staging") + entries, err := os.ReadDir(parent) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + var errs []error + for _, entry := range entries { + if !entry.IsDir() { + continue + } + root := filepath.Join(parent, entry.Name()) + raw, err := os.ReadFile(filepath.Join(root, "journal.json")) + if err != nil { + errs = append(errs, err) + continue + } + var journal archiveStageJournal + if err := json.Unmarshal(raw, &journal); err != nil { + errs = append(errs, err) + continue + } + archive, getErr := pg.GetTaskArchive(journal.ArchiveID) + if getErr != nil { + errs = append(errs, getErr) + continue + } + if archive != nil && archive.State == pgdb.ArchiveReady { + if err := os.RemoveAll(root); err != nil { + errs = append(errs, err) + } + continue + } + stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload"), journal: journal} + if err := stage.rollback(); err != nil { + errs = append(errs, err) + } + if archive != nil && archive.ArchivePath != "" { + _ = os.Remove(archive.ArchivePath) + } + } + return errors.Join(errs...) +} + +func recoverTaskArchiveDeletePackages(dataDir string, pg *pgdb.DB) error { + root := taskArchiveRoot(dataDir) + entries, err := os.ReadDir(root) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + var errs []error + for _, entry := range entries { + if entry.IsDir() { + continue + } + name := entry.Name() + marker := strings.LastIndex(name, ".deleting-") + if marker < 0 { + continue + } + archiveID, err := strconv.ParseInt(name[marker+len(".deleting-"):], 10, 64) + if err != nil || archiveID <= 0 { + continue + } + staged := filepath.Join(root, name) + original := filepath.Join(root, name[:marker]) + archive, err := pg.GetTaskArchive(archiveID) + if err != nil { + errs = append(errs, err) + continue + } + if archive == nil { + if err := os.Remove(staged); err != nil && !os.IsNotExist(err) { + errs = append(errs, err) + } + continue + } + if archive.State == pgdb.DeleteQueued || archive.State == pgdb.Deleting || archive.State == pgdb.DeleteFailed { + // The idempotent delete worker consumes the staged path directly. + continue + } + if _, err := os.Lstat(original); err == nil { + errs = append(errs, fmt.Errorf("archive delete recovery destination exists: %s", original)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, err) + continue + } + if err := os.Rename(staged, original); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +func stageTaskArchivePackageDelete(archivePath string, archiveID int64) (string, bool, error) { + staged := archivePath + fmt.Sprintf(".deleting-%d", archiveID) + if _, err := os.Lstat(staged); err == nil { + if _, originalErr := os.Lstat(archivePath); originalErr == nil { + return staged, false, errors.New("归档包原文件和删除暂存文件同时存在") + } else if !os.IsNotExist(originalErr) { + return staged, false, originalErr + } + return staged, true, nil + } else if !os.IsNotExist(err) { + return staged, false, err + } + if err := os.Rename(archivePath, staged); err == nil { + return staged, true, nil + } else if !os.IsNotExist(err) { + return staged, false, err + } + return staged, false, nil +} + +func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchiveSnapshot) (originalSize, compressedSize int64, checksum string, err error) { + if snapshot == nil { + return 0, 0, "", errors.New("nil task archive snapshot") + } + if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil { + return 0, 0, "", err + } + manifest, err := json.Marshal(snapshot) + if err != nil { + return 0, 0, "", err + } + if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil { + return 0, 0, "", err + } + temporary := path + ".partial" + _ = os.Remove(temporary) + file, err := os.OpenFile(temporary, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode) + if err != nil { + return 0, 0, "", err + } + committed := false + defer func() { + _ = file.Close() + if !committed { + _ = os.Remove(temporary) + } + }() + hasher := sha256.New() + multi := io.MultiWriter(file, hasher) + encoder, err := zstd.NewWriter(multi, + zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(10)), + zstd.WithEncoderCRC(true), + ) + if err != nil { + return 0, 0, "", err + } + tarWriter := tar.NewWriter(encoder) + walkErr := filepath.WalkDir(payloadDir, func(current string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + relative, err := filepath.Rel(payloadDir, current) + if err != nil || relative == "." { + return err + } + info, err := entry.Info() + if err != nil { + return err + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("task archive refuses symlink: %s", current) + } + header, err := tar.FileInfoHeader(info, "") + if err != nil { + return err + } + header.Name = filepath.ToSlash(relative) + if err := tarWriter.WriteHeader(header); err != nil { + return err + } + if entry.IsDir() { + return nil + } + originalSize += info.Size() + input, err := os.Open(current) + if err != nil { + return err + } + _, copyErr := io.Copy(tarWriter, input) + closeErr := input.Close() + return errors.Join(copyErr, closeErr) + }) + if walkErr != nil { + _ = tarWriter.Close() + _ = encoder.Close() + return 0, 0, "", walkErr + } + if err := tarWriter.Close(); err != nil { + _ = encoder.Close() + return 0, 0, "", err + } + if err := encoder.Close(); err != nil { + return 0, 0, "", err + } + if err := file.Sync(); err != nil { + return 0, 0, "", err + } + if err := file.Close(); err != nil { + return 0, 0, "", err + } + stat, err := os.Stat(temporary) + if err != nil { + return 0, 0, "", err + } + compressedSize = stat.Size() + checksum = hex.EncodeToString(hasher.Sum(nil)) + if err := os.Rename(temporary, path); err != nil { + return 0, 0, "", err + } + if err := os.Chmod(path, archiveFileMode); err != nil { + return 0, 0, "", err + } + committed = true + return originalSize, compressedSize, checksum, nil +} + +func extractTaskArchivePackage(path, expectedSHA, destination string) error { + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + if expectedSHA != "" { + hasher := sha256.New() + if _, err := io.Copy(hasher, file); err != nil { + return err + } + if hex.EncodeToString(hasher.Sum(nil)) != strings.ToLower(expectedSHA) { + return errors.New("task archive checksum mismatch") + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + return err + } + } + decoder, err := zstd.NewReader(file) + if err != nil { + return err + } + defer decoder.Close() + tarReader := tar.NewReader(decoder) + if err := os.MkdirAll(destination, archiveDirMode); err != nil { + return err + } + var entries int + var total int64 + for { + header, err := tarReader.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return err + } + entries++ + if entries > maxArchiveFiles || header.Size < 0 || total+header.Size > maxArchiveBytes { + return errors.New("task archive exceeds extraction safety limits") + } + total += header.Size + clean := filepath.Clean(filepath.FromSlash(header.Name)) + if clean == "." || filepath.IsAbs(clean) || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) { + return fmt.Errorf("unsafe task archive path %q", header.Name) + } + target := filepath.Join(destination, clean) + if relative, err := filepath.Rel(destination, target); err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return fmt.Errorf("unsafe task archive target %q", header.Name) + } + switch header.Typeflag { + case tar.TypeDir: + if err := os.MkdirAll(target, archiveDirMode); err != nil { + return err + } + case tar.TypeReg, tar.TypeRegA: + if err := os.MkdirAll(filepath.Dir(target), archiveDirMode); err != nil { + return err + } + output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, fs.FileMode(header.Mode)&0o777) + if err != nil { + return err + } + _, copyErr := io.CopyN(output, tarReader, header.Size) + closeErr := output.Close() + if err := errors.Join(copyErr, closeErr); err != nil { + return err + } + default: + return fmt.Errorf("unsupported task archive entry type %d", header.Typeflag) + } + } + return nil +} diff --git a/server/task_archive_package_test.go b/server/task_archive_package_test.go new file mode 100644 index 0000000..74e86ad --- /dev/null +++ b/server/task_archive_package_test.go @@ -0,0 +1,209 @@ +package server + +import ( + "archive/tar" + "bytes" + "encoding/json" + "os" + "path/filepath" + "testing" + + pgdb "github.com/Autumn-27/artex/db" + "github.com/klauspost/compress/zstd" +) + +func TestTaskArchivePackageFilesRoundTrip(t *testing.T) { + dataDir := t.TempDir() + taskID := "42" + explorationID := int64(73) + taskFile := filepath.Join(dataDir, "tasks", taskID, "uploads", "evidence.txt") + transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker-1.jsonl") + unrelatedTranscript := filepath.Join(dataDir, "transcripts", "exp74-worker-1.jsonl") + for path, body := range map[string]string{ + taskFile: "task evidence", + transcriptFile: "transcript", + unrelatedTranscript: "leave me hot", + } { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + } + + stage, err := stageTaskArchiveFiles(dataDir, 1, taskID, explorationID) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(taskFile); !os.IsNotExist(err) { + t.Fatalf("task file still hot after staging: %v", err) + } + if _, err := os.Stat(transcriptFile); !os.IsNotExist(err) { + t.Fatalf("task transcript still hot after staging: %v", err) + } + if _, err := os.Stat(unrelatedTranscript); err != nil { + t.Fatalf("unrelated transcript was staged: %v", err) + } + + archivePath := taskArchivePath(dataDir, 1, taskID) + snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID} + original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot) + if err != nil { + t.Fatal(err) + } + if original == 0 || compressed == 0 || checksum == "" { + t.Fatalf("invalid package metrics original=%d compressed=%d checksum=%q", original, compressed, checksum) + } + if err := stage.commit(); err != nil { + t.Fatal(err) + } + extracted := filepath.Join(dataDir, "restore") + if err := extractTaskArchivePackage(archivePath, checksum, extracted); err != nil { + t.Fatal(err) + } + installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1) + if err != nil { + t.Fatal(err) + } + if err := installed.commit(); err != nil { + t.Fatal(err) + } + for path, want := range map[string]string{taskFile: "task evidence", transcriptFile: "transcript"} { + got, err := os.ReadFile(path) + if err != nil || string(got) != want { + t.Fatalf("restored %s = %q, %v; want %q", path, got, err, want) + } + } + if err := extractTaskArchivePackage(archivePath, "deadbeef", filepath.Join(dataDir, "bad-checksum")); err == nil { + t.Fatal("checksum mismatch was accepted") + } +} + +func TestTaskArchiveStageJournalRollsBackInterruptedMoves(t *testing.T) { + dataDir := t.TempDir() + taskFile := filepath.Join(dataDir, "tasks", "42", "evidence.txt") + transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker.jsonl") + for _, path := range []string{taskFile, transcriptFile} { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte(path), 0o600); err != nil { + t.Fatal(err) + } + } + stage, err := stageTaskArchiveFiles(dataDir, 9, "42", 73) + if err != nil { + t.Fatal(err) + } + raw, err := os.ReadFile(filepath.Join(stage.root, "journal.json")) + if err != nil { + t.Fatal(err) + } + var journal archiveStageJournal + if err := json.Unmarshal(raw, &journal); err != nil { + t.Fatal(err) + } + if len(journal.Moves) != 2 { + t.Fatalf("journal moves=%d, want 2", len(journal.Moves)) + } + recovered := &taskArchiveFileStage{root: stage.root, payload: stage.payload, journal: journal} + if err := recovered.rollback(); err != nil { + t.Fatal(err) + } + for _, path := range []string{taskFile, transcriptFile} { + if _, err := os.Stat(path); err != nil { + t.Fatalf("interrupted archive did not restore %s: %v", path, err) + } + } +} + +func TestTaskArchiveRestoreJournalRollsBackInterruptedInstall(t *testing.T) { + dataDir := t.TempDir() + extracted := filepath.Join(dataDir, "archives", "tasks", ".restore", "11-test") + source := filepath.Join(extracted, "files", "tasks", "42", "evidence.txt") + if err := os.MkdirAll(filepath.Dir(source), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(source, []byte("evidence"), 0o600); err != nil { + t.Fatal(err) + } + _, err := installTaskArchiveFiles(dataDir, extracted, "42", 11) + if err != nil { + t.Fatal(err) + } + destination := filepath.Join(dataDir, "tasks", "42", "evidence.txt") + if _, err := os.Stat(destination); err != nil { + t.Fatal(err) + } + raw, err := os.ReadFile(filepath.Join(extracted, "restore-journal.json")) + if err != nil { + t.Fatal(err) + } + var journal archiveRestoreJournal + if err := json.Unmarshal(raw, &journal); err != nil { + t.Fatal(err) + } + if journal.ArchiveID != 11 || len(journal.Moves) != 1 { + t.Fatalf("unexpected restore journal: %+v", journal) + } + recovered := &taskArchiveRestoreFiles{extracted: extracted, moves: journal.Moves} + if err := recovered.rollback(); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(destination); !os.IsNotExist(err) { + t.Fatalf("hot destination remains after restore rollback: %v", err) + } + if got, err := os.ReadFile(source); err != nil || string(got) != "evidence" { + t.Fatalf("restore source=%q err=%v", got, err) + } +} + +func TestTaskArchiveDeletePackageCanResumeFromStagedPath(t *testing.T) { + dataDir := t.TempDir() + archivePath := taskArchivePath(dataDir, 21, "42") + if err := os.MkdirAll(filepath.Dir(archivePath), archiveDirMode); err != nil { + t.Fatal(err) + } + staged := archivePath + ".deleting-21" + if err := os.WriteFile(staged, []byte("archive"), archiveFileMode); err != nil { + t.Fatal(err) + } + got, moved, err := stageTaskArchivePackageDelete(archivePath, 21) + if err != nil || !moved || got != staged { + t.Fatalf("resume staged package path=%q moved=%v err=%v", got, moved, err) + } +} + +func TestTaskArchivePackageRejectsTraversal(t *testing.T) { + var compressed bytes.Buffer + encoder, err := zstd.NewWriter(&compressed) + if err != nil { + t.Fatal(err) + } + tw := tar.NewWriter(encoder) + content := []byte("escape") + if err := tw.WriteHeader(&tar.Header{Name: "../escape.txt", Mode: 0o600, Size: int64(len(content))}); err != nil { + t.Fatal(err) + } + if _, err := tw.Write(content); err != nil { + t.Fatal(err) + } + if err := tw.Close(); err != nil { + t.Fatal(err) + } + if err := encoder.Close(); err != nil { + t.Fatal(err) + } + path := filepath.Join(t.TempDir(), "malicious.tar.zst") + if err := os.WriteFile(path, compressed.Bytes(), 0o600); err != nil { + t.Fatal(err) + } + root := t.TempDir() + if err := extractTaskArchivePackage(path, "", root); err == nil { + t.Fatal("path traversal archive was accepted") + } + if _, err := os.Stat(filepath.Join(filepath.Dir(root), "escape.txt")); !os.IsNotExist(err) { + t.Fatalf("traversal wrote outside destination: %v", err) + } +} diff --git a/server/task_archives.go b/server/task_archives.go new file mode 100644 index 0000000..c953a1a --- /dev/null +++ b/server/task_archives.go @@ -0,0 +1,656 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/Autumn-27/artex/agent" + pgdb "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +const taskArchivePollInterval = 2 * time.Second + +type archiveBatchRequest struct { + TaskIDs []string `json:"task_ids"` + ArchiveIDs []int64 `json:"archive_ids"` +} + +type archiveBatchItem struct { + ID string `json:"id"` + OK bool `json:"ok"` + Queued bool `json:"queued"` + Error string `json:"error,omitempty"` + Archive int64 `json:"archive_id,omitempty"` +} + +func (s *Server) startTaskArchiveWorker() { + if s.m == nil || s.m.pg == nil { + return + } + if err := s.m.pg.RecoverTaskArchiveJobs(); err != nil { + log.Printf("[task-archive] recover jobs: %v", err) + } + recoveryFailed := false + if err := recoverTaskArchiveStages(s.m.dir, s.m.pg); err != nil { + log.Printf("[task-archive] recover file staging: %v", err) + recoveryFailed = true + } + if err := recoverTaskArchiveRestoreStages(s.m.dir, s.m.pg); err != nil { + log.Printf("[task-archive] recover restore staging: %v", err) + recoveryFailed = true + } + if err := recoverTaskArchiveDeletePackages(s.m.dir, s.m.pg); err != nil { + log.Printf("[task-archive] recover delete packages: %v", err) + recoveryFailed = true + } + if recoveryFailed { + log.Printf("[task-archive] worker disabled because startup recovery is incomplete") + return + } + s.archiveWG.Add(1) + go func() { + defer s.archiveWG.Done() + ticker := time.NewTicker(taskArchivePollInterval) + defer ticker.Stop() + for { + if err := s.runOneTaskArchiveJob(); err != nil { + log.Printf("[task-archive] worker: %v", err) + } + select { + case <-s.ctx.Done(): + return + case <-s.archiveWake: + case <-ticker.C: + } + } + }() + s.notifyTaskArchiveWorker() +} + +func (s *Server) notifyTaskArchiveWorker() { + if s == nil || s.archiveWake == nil { + return + } + select { + case s.archiveWake <- struct{}{}: + default: + } +} + +func (s *Server) runOneTaskArchiveJob() error { + job, err := s.m.pg.ClaimTaskArchiveJob(s.ctx) + if err != nil || job == nil { + return err + } + var runErr error + switch job.State { + case pgdb.Archiving: + runErr = s.archiveTask(job) + case pgdb.Restoring: + runErr = s.restoreTaskArchive(job) + case pgdb.Deleting: + runErr = s.deleteTaskArchive(job) + default: + runErr = fmt.Errorf("unknown claimed archive state %q", job.State) + } + if runErr != nil { + if failErr := s.m.pg.FailTaskArchiveJob(job.ID, job.State, runErr); failErr != nil { + return errors.Join(runErr, failErr) + } + return fmt.Errorf("archive job %d (%s): %w", job.ID, job.State, runErr) + } + // Drain another queued item without waiting for the periodic poll. + s.notifyTaskArchiveWorker() + return nil +} + +func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) { + taskID := strconv.FormatInt(job.TaskID, 10) + if !s.beginTaskDelete(taskID) { + return errors.New("任务正在进行其他归档或删除操作") + } + committed := false + defer func() { + if !committed { + s.abortTaskDelete(taskID) + } + }() + s.cancelTaskChat(taskID, agent.AbortTaskDeleted) + drainCtx, cancel := context.WithTimeout(s.ctx, taskDeleteDrainTimeout) + defer cancel() + if err := s.waitTaskQuiescent(drainCtx, taskID); err != nil { + return errors.New("任务仍有运行中的 Agent，请先暂停后重试归档") + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 10) + snapshot, err := s.m.pg.SnapshotTaskArchive(job.TaskID) + if err != nil { + return err + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 25) + fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, snapshot.ExplorationID) + if err != nil { + return err + } + defer func() { + if !committed { + runErr = errors.Join(runErr, fileStage.rollback()) + } + }() + if s.m.traffic != nil && len(snapshot.Hosts) > 0 { + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_traffic", 38) + trafficCount, err := s.m.traffic.ExportHosts(snapshot.Hosts, filepath.Join(fileStage.payload, "traffic")) + if err != nil { + return err + } + snapshot.DataCounts["traffic"] = trafficCount + } + archivePath := taskArchivePath(s.m.dir, job.ID, taskID) + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50) + originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot) + if err != nil { + return err + } + removePackage := true + defer func() { + if removePackage { + _ = os.Remove(archivePath) + } + }() + var trafficStage *traffic.HostDeleteStage + if s.m.traffic != nil && len(snapshot.ExclusiveHosts) > 0 { + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_traffic", 68) + trafficStage, err = s.m.traffic.StageDeleteHostsExactForArchive(snapshot.ExclusiveHosts, job.ID, job.TaskID) + if err != nil { + return err + } + defer func() { + if !committed { + runErr = errors.Join(runErr, trafficStage.Rollback()) + } + }() + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_database", 78) + if err := s.m.pg.CompleteTaskArchive(job.ID, snapshot, archivePath, checksum, originalSize, compressedSize); err != nil { + return err + } + committed = true + removePackage = false + if trafficStage != nil { + if err := trafficStage.Commit(); err != nil { + warning := "归档已完成，但独占流量热存储清理失败（归档包仍可恢复）：" + err.Error() + log.Printf("[task-archive] task %s: %s", taskID, warning) + _ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning) + } + } + if err := fileStage.commit(); err != nil { + warning := "归档已完成，但暂存目录清理失败：" + err.Error() + log.Printf("[task-archive] task %s: %s", taskID, warning) + _ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning) + } + s.engine.StopTask(taskID) + s.taskAgentMu.Lock() + delete(s.taskAgents, taskID) + s.taskAgentMu.Unlock() + s.m.forgetTask(taskID, job.TaskID) + return nil +} + +func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) { + if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil { + return err + } + // A prior attempt may have committed PostgreSQL and files, then failed only + // while consuming the package. Finish that cleanup without replaying rows. + if restored, err := s.m.pg.IsTaskArchiveRestored(job.ID); err != nil { + return err + } else if restored { + if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) { + return err + } + if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil { + return err + } + s.loadRestoredTask(job.TaskID) + return nil + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "verify_package", 10) + restoreParent := filepath.Join(taskArchiveRoot(s.m.dir), ".restore") + if err := os.MkdirAll(restoreParent, archiveDirMode); err != nil { + return err + } + extracted, err := os.MkdirTemp(restoreParent, fmt.Sprintf("%d-", job.ID)) + if err != nil { + return err + } + defer func() { _ = os.RemoveAll(extracted) }() + if err := extractTaskArchivePackage(job.ArchivePath, job.SHA256, extracted); err != nil { + return err + } + manifestRaw, err := os.ReadFile(filepath.Join(extracted, "manifest.json")) + if err != nil { + return err + } + var snapshot pgdb.TaskArchiveSnapshot + if err := json.Unmarshal(manifestRaw, &snapshot); err != nil { + return err + } + if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion { + return pgdb.ErrTaskArchiveFormatMismatch + } + if s.m.traffic != nil { + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_traffic", 35) + if _, err := s.m.traffic.ImportArchive(filepath.Join(extracted, "traffic")); err != nil { + return err + } + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_files", 55) + files, err := installTaskArchiveFiles(s.m.dir, extracted, strconv.FormatInt(job.TaskID, 10), job.ID) + if err != nil { + return err + } + databaseRestored := false + defer func() { + if !databaseRestored { + runErr = errors.Join(runErr, files.rollback()) + } + }() + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70) + warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds) + if err != nil { + return err + } + for _, warning := range warnings { + log.Printf("[task-archive] restored task %d with warning: %s", job.TaskID, warning) + } + if len(warnings) > 0 { + if _, err := s.m.pg.Exploration(snapshot.ExplorationID).AppendActivity(pgdb.Activity{ + Worker: "system", Kind: "system", Summary: "任务还原完成，但部分关联已降级", Detail: strings.Join(warnings, "\n"), + }); err != nil { + log.Printf("[task-archive] persist restore warnings for task %d: %v", job.TaskID, err) + } + } + databaseRestored = true + if err := files.commit(); err != nil { + return err + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "consume_package", 95) + if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) { + return err + } + if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil { + return err + } + s.loadRestoredTask(job.TaskID) + return nil +} + +func (s *Server) loadRestoredTask(taskID int64) { + id := strconv.FormatInt(taskID, 10) + var restored *Task + for _, task := range s.m.LoadExisting() { + if task.ID == id { + restored = task + break + } + } + if restored == nil { + if task, ok := s.m.Task(id); ok { + restored = task + } + } + if restored == nil { + return + } + lifecycle := restored.lifecycleSnapshot() + if lifecycle.Paused { + s.engine.Pause(id, agent.AbortPausedOnReload) + } + if !isTerminalStatus(lifecycle.Status) { + s.engine.startDeadlineCoordinator(s.ctx, restored) + } +} + +func (s *Server) deleteTaskArchive(job *pgdb.TaskArchive) error { + if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil { + return err + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_package_delete", 35) + staged, moved, err := stageTaskArchivePackageDelete(job.ArchivePath, job.ID) + if err != nil { + return err + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "delete_metadata", 70) + if err := s.m.pg.DeleteTaskArchiveStub(job.ID); err != nil { + if moved { + _ = os.Rename(staged, job.ArchivePath) + } + return err + } + if moved { + if err := os.Remove(staged); err != nil && !os.IsNotExist(err) { + log.Printf("[task-archive] archived task %d metadata deleted; stale package cleanup failed: %v", job.TaskID, err) + } + } + return nil +} + +func validateArchivePath(dataDir, candidate string) error { + if strings.TrimSpace(candidate) == "" { + return errors.New("归档包路径为空") + } + root, err := filepath.Abs(taskArchiveRoot(dataDir)) + if err != nil { + return err + } + path, err := filepath.Abs(candidate) + if err != nil { + return err + } + relative, err := filepath.Rel(root, path) + if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return errors.New("归档包路径不在受管目录内") + } + return nil +} + +func (s *Server) listTaskArchives(w http.ResponseWriter, r *http.Request) { + page, _ := strconv.Atoi(r.URL.Query().Get("page")) + size, _ := strconv.Atoi(r.URL.Query().Get("size")) + result, err := s.m.pg.ListTaskArchives(r.URL.Query().Get("q"), r.URL.Query().Get("state"), page, size) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, result) +} + +func (s *Server) getTaskArchive(w http.ResponseWriter, r *http.Request) { + id, ok := pathInt(r, "id") + if !ok { + writeErr(w, 400, "归档 id 无效") + return + } + item, err := s.m.pg.GetTaskArchive(id) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if item == nil { + writeErr(w, 404, "归档不存在") + return + } + writeJSON(w, 200, item) +} + +func (s *Server) queueTaskArchive(w http.ResponseWriter, r *http.Request) { + id, ok := canonicalTaskID(r.PathValue("id")) + if !ok { + writeErr(w, 400, "任务 id 无效") + return + } + numeric, _ := strconv.ParseInt(id, 10, 64) + item, err := s.m.pg.QueueTaskArchive(numeric) + if err != nil { + writeArchiveError(w, err) + return + } + s.notifyTaskArchiveWorker() + writeJSON(w, http.StatusAccepted, item) +} + +func (s *Server) queueTaskArchiveRestore(w http.ResponseWriter, r *http.Request) { + id, ok := pathInt(r, "id") + if !ok { + writeErr(w, 400, "归档 id 无效") + return + } + item, err := s.m.pg.QueueTaskArchiveRestore(id) + if err != nil { + writeArchiveError(w, err) + return + } + s.notifyTaskArchiveWorker() + writeJSON(w, http.StatusAccepted, item) +} + +func (s *Server) queueTaskArchiveDelete(w http.ResponseWriter, r *http.Request) { + id, ok := pathInt(r, "id") + if !ok { + writeErr(w, 400, "归档 id 无效") + return + } + item, err := s.m.pg.QueueTaskArchiveDelete(id) + if err != nil { + writeArchiveError(w, err) + return + } + s.notifyTaskArchiveWorker() + writeJSON(w, http.StatusAccepted, item) +} + +func (s *Server) queueTaskArchivesBatch(w http.ResponseWriter, r *http.Request) { + var request archiveBatchRequest + if err := decode(r, &request); err != nil { + writeErr(w, 400, err.Error()) + return + } + parsed := normalizeBatchTaskIDs(request.TaskIDs) + if len(parsed) == 0 { + writeErr(w, 400, "task_ids 不能为空") + return + } + if len(parsed) > 100 { + writeErr(w, 400, "一次最多处理 100 个任务") + return + } + ids := make([]string, 0, len(parsed)) + items := make([]archiveBatchItem, 0, len(parsed)) + for _, item := range parsed { + if !item.valid { + items = append(items, archiveBatchItem{ID: item.id, Error: "任务 id 无效"}) + continue + } + ids = append(ids, item.id) + } + ids = s.orderLiveTaskArchiveBatch(ids) + for _, id := range ids { + numeric, _ := strconv.ParseInt(id, 10, 64) + archive, err := s.m.pg.QueueTaskArchive(numeric) + item := archiveBatchItem{ID: id, OK: err == nil, Queued: err == nil} + if err != nil { + item.Error = err.Error() + } else { + item.Archive = archive.ID + } + items = append(items, item) + } + s.notifyTaskArchiveWorker() + writeJSON(w, http.StatusAccepted, map[string]any{"items": items}) +} + +func (s *Server) queueTaskArchiveActionBatch(w http.ResponseWriter, r *http.Request, action string) { + var request archiveBatchRequest + if err := decode(r, &request); err != nil { + writeErr(w, 400, err.Error()) + return + } + ids, err := normalizeArchiveIDs(request.ArchiveIDs) + if err != nil { + writeErr(w, 400, err.Error()) + return + } + ids = s.orderColdTaskArchiveBatch(ids, action == "restore") + items := make([]archiveBatchItem, 0, len(ids)) + for _, id := range ids { + var archive *pgdb.TaskArchive + if action == "restore" { + archive, err = s.m.pg.QueueTaskArchiveRestore(id) + } else { + archive, err = s.m.pg.QueueTaskArchiveDelete(id) + } + item := archiveBatchItem{ID: strconv.FormatInt(id, 10), Archive: id, OK: err == nil, Queued: err == nil} + if err != nil { + item.Error = err.Error() + } else if archive != nil { + item.Archive = archive.ID + } + items = append(items, item) + } + s.notifyTaskArchiveWorker() + writeJSON(w, http.StatusAccepted, map[string]any{"items": items}) +} + +func (s *Server) restoreTaskArchivesBatch(w http.ResponseWriter, r *http.Request) { + s.queueTaskArchiveActionBatch(w, r, "restore") +} + +func (s *Server) deleteTaskArchivesBatch(w http.ResponseWriter, r *http.Request) { + s.queueTaskArchiveActionBatch(w, r, "delete") +} + +func normalizeArchiveIDs(ids []int64) ([]int64, error) { + if len(ids) == 0 { + return nil, errors.New("archive_ids 不能为空") + } + if len(ids) > 100 { + return nil, errors.New("一次最多处理 100 个归档") + } + seen := map[int64]bool{} + out := make([]int64, 0, len(ids)) + for _, id := range ids { + if id <= 0 { + return nil, fmt.Errorf("归档 id %d 无效", id) + } + if !seen[id] { + seen[id] = true + out = append(out, id) + } + } + return out, nil +} + +func (s *Server) orderLiveTaskArchiveBatch(ids []string) []string { + sources := map[string][]string{} + for _, id := range ids { + if task, ok := s.m.Task(id); ok { + for _, sourceID := range task.lifecycleSnapshot().SourceTaskIDs { + sources[id] = append(sources[id], strconv.FormatInt(sourceID, 10)) + } + } + } + return topoTaskIDs(ids, sources, false) +} + +func (s *Server) orderColdTaskArchiveBatch(ids []int64, restore bool) []int64 { + taskToArchive := map[string]int64{} + sources := map[string][]string{} + var taskIDs []string + for _, id := range ids { + archive, err := s.m.pg.GetTaskArchive(id) + if err != nil || archive == nil { + continue + } + taskID := strconv.FormatInt(archive.TaskID, 10) + taskToArchive[taskID] = id + taskIDs = append(taskIDs, taskID) + for _, sourceID := range archive.SourceTaskIDs { + sources[taskID] = append(sources[taskID], strconv.FormatInt(sourceID, 10)) + } + } + orderedTasks := topoTaskIDs(taskIDs, sources, restore) + ordered := make([]int64, 0, len(ids)) + seen := map[int64]bool{} + for _, taskID := range orderedTasks { + id := taskToArchive[taskID] + ordered = append(ordered, id) + seen[id] = true + } + for _, id := range ids { + if !seen[id] { + ordered = append(ordered, id) + } + } + return ordered +} + +// topoTaskIDs orders source->dependent for restore and dependent->source for +// archive/delete. Cycles are impossible at creation, but stable fallback keeps a +// damaged database operable. +func topoTaskIDs(ids []string, sources map[string][]string, restore bool) []string { + set := map[string]bool{} + for _, id := range ids { + set[id] = true + } + adjacency := map[string][]string{} + indegree := map[string]int{} + for id := range set { + indegree[id] = 0 + } + for dependent, sourceIDs := range sources { + for _, source := range sourceIDs { + if !set[dependent] || !set[source] { + continue + } + from, to := dependent, source + if restore { + from, to = source, dependent + } + adjacency[from] = append(adjacency[from], to) + indegree[to]++ + } + } + queue := []string{} + for _, id := range ids { + if indegree[id] == 0 { + queue = append(queue, id) + } + } + ordered := make([]string, 0, len(ids)) + for len(queue) > 0 { + id := queue[0] + queue = queue[1:] + ordered = append(ordered, id) + for _, next := range adjacency[id] { + indegree[next]-- + if indegree[next] == 0 { + queue = append(queue, next) + } + } + } + if len(ordered) != len(ids) { + seen := map[string]bool{} + for _, id := range ordered { + seen[id] = true + } + for _, id := range ids { + if !seen[id] { + ordered = append(ordered, id) + } + } + } + return ordered +} + +func writeArchiveError(w http.ResponseWriter, err error) { + status := http.StatusConflict + if errors.Is(err, pgdb.ErrTaskArchiveNotFound) { + status = http.StatusNotFound + } else if !errors.Is(err, pgdb.ErrTaskArchiveIneligible) && + !errors.Is(err, pgdb.ErrTaskArchiveQueued) && + !errors.Is(err, pgdb.ErrTaskArchiveDependent) && + !errors.Is(err, pgdb.ErrTaskArchiveState) && + !errors.Is(err, pgdb.ErrTaskArchiveDeleteBlocked) { + status = http.StatusInternalServerError + } + writeErr(w, status, err.Error()) +} diff --git a/server/task_archives_test.go b/server/task_archives_test.go new file mode 100644 index 0000000..3d02999 --- /dev/null +++ b/server/task_archives_test.go @@ -0,0 +1,88 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" +) + +func TestTaskArchiveAPIQueueListAndLimits(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("archive api", "verify 202", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + taskID, err := strconv.ParseInt(task.ID, 10, 64) + if err != nil { + t.Fatal(err) + } + defer func() { + _, _ = m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, taskID) + _ = m.pg.DeleteTask(taskID) + }() + + // Start with an already-cancelled context so the background archive worker + // exits before these route-contract assertions enqueue work. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + dataDir := t.TempDir() + s := New(ctx, m, dataDir, dataDir, dataDir) + s.archiveWG.Wait() + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatal(err) + } + request := func(method, path, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest(method, path, strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + s.Handler().ServeHTTP(recorder, req) + return recorder + } + + recorder := request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}") + if recorder.Code != http.StatusConflict { + t.Fatalf("running archive status=%d body=%s", recorder.Code, recorder.Body.String()) + } + if err := m.pg.SetPaused(taskID, true); err != nil { + t.Fatal(err) + } + recorder = request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}") + if recorder.Code != http.StatusAccepted { + t.Fatalf("archive queue status=%d body=%s", recorder.Code, recorder.Body.String()) + } + var queued struct { + ID int64 `json:"id"` + State string `json:"state"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &queued); err != nil { + t.Fatal(err) + } + if queued.ID <= 0 || queued.State != "archive_queued" { + t.Fatalf("unexpected queue response: %+v", queued) + } + + recorder = request(http.MethodGet, "/api/task-archives?page=1&size=10&q="+task.ID, "") + if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), fmt.Sprintf(`"task_id":%d`, taskID)) { + t.Fatalf("archive list status=%d body=%s", recorder.Code, recorder.Body.String()) + } + ids := make([]string, 101) + for index := range ids { + ids[index] = strconv.Itoa(index + 1) + } + body, _ := json.Marshal(map[string]any{"task_ids": ids}) + recorder = request(http.MethodPost, "/api/tasks/archive/batch", string(body)) + if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), "100") { + t.Fatalf("oversized batch status=%d body=%s", recorder.Code, recorder.Body.String()) + } +} diff --git a/server/task_llm_test.go b/server/task_llm_test.go index a72ae34..48da7d3 100644 --- a/server/task_llm_test.go +++ b/server/task_llm_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "iter" + "strings" "sync" "testing" "time" @@ -148,6 +149,22 @@ func TestTaskDTOUsesEmptyArrays(t *testing.T) { } } +func TestApplyTaskArchiveBlocker(t *testing.T) { + t.Parallel() + dto := taskDTO(&Task{ID: "7"}, "paused") + applyTaskArchiveBlocker(&dto, map[int64]int64{7: 11}) + if dto.ArchiveBlockedBy != "11" { + t.Fatalf("archive blocker=%q, want 11", dto.ArchiveBlockedBy) + } + raw, err := json.Marshal(dto) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(raw), `"archive_blocked_by_task_id":"11"`) { + t.Fatalf("archive blocker missing from payload: %s", raw) + } +} + func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) { t.Parallel() task := &Task{ID: "7"} diff --git a/traffic/archive.go b/traffic/archive.go new file mode 100644 index 0000000..8f57dd1 --- /dev/null +++ b/traffic/archive.go @@ -0,0 +1,237 @@ +package traffic + +import ( + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "strings" +) + +// ArchiveSnapshot is the portable traffic subset embedded in one task archive. +// Large content-addressed bodies are copied alongside this manifest rather than +// base64-encoded, preserving deduplication and allowing the tar writer to stream. +type ArchiveSnapshot struct { + Version int `json:"version"` + Exchanges []ArchiveExchange `json:"exchanges"` + Blobs []string `json:"blobs"` +} + +type ArchiveExchange struct { + ID string `json:"id"` + TS int64 `json:"ts"` + Host string `json:"host"` + Method string `json:"method"` + URLTemplate string `json:"url_template"` + URL string `json:"url"` + Status int `json:"status"` + ContentType string `json:"content_type"` + ReqLen int `json:"req_len"` + RespLen int `json:"resp_len"` + ReqHead string `json:"req_head"` + ReqBody []byte `json:"req_body,omitempty"` + ReqBlob string `json:"req_blob,omitempty"` + RespHead string `json:"resp_head"` + RespBody []byte `json:"resp_body,omitempty"` + RespBlob string `json:"resp_blob,omitempty"` +} + +// ExportHosts writes an exact-host snapshot to dir. It holds the traffic writer +// lock while reading SQLite and blobs, so each body and its index row come from +// one consistent point in time. +func (t *Traffic) ExportHosts(hosts []string, dir string) (int64, error) { + if t == nil || len(hosts) == 0 { + return 0, nil + } + if err := os.MkdirAll(filepath.Join(dir, "blobs"), 0o700); err != nil { + return 0, err + } + t.wmu.Lock() + defer t.wmu.Unlock() + unique := uniqueArchiveHosts(hosts) + placeholders := make([]string, len(unique)) + args := make([]any, len(unique)) + for i, host := range unique { + placeholders[i] = "?" + args[i] = host + } + rows, err := t.db.Query(`SELECT id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path +FROM exchanges WHERE host IN (`+strings.Join(placeholders, ",")+`) ORDER BY ts,id`, args...) + if err != nil { + return 0, err + } + defer rows.Close() + snapshot := ArchiveSnapshot{Version: 1} + blobs := map[string]struct{}{} + for rows.Next() { + var item ArchiveExchange + var legacyPath string + if err := rows.Scan(&item.ID, &item.TS, &item.Host, &item.Method, &item.URLTemplate, + &item.URL, &item.Status, &item.ContentType, &item.ReqLen, &item.RespLen, &legacyPath); err != nil { + return 0, err + } + var reqBlob, respBlob sql.NullString + err := t.db.QueryRow(`SELECT req_head,req_body,req_blob,resp_head,resp_body,resp_blob +FROM exchange_bodies WHERE id=?`, item.ID).Scan(&item.ReqHead, &item.ReqBody, &reqBlob, &item.RespHead, &item.RespBody, &respBlob) + if errors.Is(err, sql.ErrNoRows) && strings.TrimSpace(legacyPath) != "" { + req, readErr := os.ReadFile(filepath.Join(t.dir, legacyPath, "request.http")) + if readErr != nil && !os.IsNotExist(readErr) { + return 0, readErr + } + resp, readErr := os.ReadFile(filepath.Join(t.dir, legacyPath, "response.http")) + if readErr != nil && !os.IsNotExist(readErr) { + return 0, readErr + } + item.ReqHead, item.RespHead = string(req), string(resp) + } else if err != nil { + return 0, err + } + item.ReqBlob, item.RespBlob = reqBlob.String, respBlob.String + for _, hash := range []string{item.ReqBlob, item.RespBlob} { + if hash != "" { + blobs[hash] = struct{}{} + } + } + snapshot.Exchanges = append(snapshot.Exchanges, item) + } + if err := rows.Err(); err != nil { + return 0, err + } + for hash := range blobs { + source, err := t.blobPath(hash) + if err != nil { + return 0, err + } + data, err := os.ReadFile(source) + if err != nil { + return 0, err + } + if err := os.WriteFile(filepath.Join(dir, "blobs", hash+".bin"), data, 0o600); err != nil { + return 0, err + } + snapshot.Blobs = append(snapshot.Blobs, hash) + } + sort.Strings(snapshot.Blobs) + raw, err := json.Marshal(snapshot) + if err != nil { + return 0, err + } + if err := os.WriteFile(filepath.Join(dir, "traffic.json"), raw, 0o600); err != nil { + return 0, err + } + return int64(len(snapshot.Exchanges)), nil +} + +// ImportArchive imports only missing exchange IDs. Current hot rows always win, +// and repeated restore attempts are safe after a partial external failure. +func (t *Traffic) ImportArchive(dir string) (int64, error) { + if t == nil { + return 0, nil + } + raw, err := os.ReadFile(filepath.Join(dir, "traffic.json")) + if os.IsNotExist(err) { + return 0, nil + } + if err != nil { + return 0, err + } + var snapshot ArchiveSnapshot + if err := json.Unmarshal(raw, &snapshot); err != nil { + return 0, err + } + if snapshot.Version != 1 { + return 0, fmt.Errorf("unsupported traffic archive version %d", snapshot.Version) + } + t.wmu.Lock() + defer t.wmu.Unlock() + for _, hash := range snapshot.Blobs { + if !blobHashRe.MatchString(hash) { + return 0, fmt.Errorf("invalid archived traffic blob %q", hash) + } + data, err := os.ReadFile(filepath.Join(dir, "blobs", hash+".bin")) + if err != nil { + return 0, err + } + sum := sha256.Sum256(data) + if hex.EncodeToString(sum[:]) != hash { + return 0, fmt.Errorf("traffic blob checksum mismatch: %s", hash) + } + destination := filepath.Join(t.dir, "_blobs", "sha256", hash[:2], hash+".bin") + if _, err := os.Stat(destination); os.IsNotExist(err) { + if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil { + return 0, err + } + if err := os.WriteFile(destination, data, 0o644); err != nil { + return 0, err + } + } + } + tx, err := t.db.Begin() + if err != nil { + return 0, err + } + defer tx.Rollback() //nolint:errcheck + var imported int64 + for _, item := range snapshot.Exchanges { + res, err := tx.Exec(`INSERT OR IGNORE INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,'')`, item.ID, item.TS, item.Host, item.Method, item.URLTemplate, + item.URL, item.Status, item.ContentType, item.ReqLen, item.RespLen) + if err != nil { + return imported, err + } + n, _ := res.RowsAffected() + if n == 0 { + continue + } + if _, err := tx.Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,req_blob,resp_head,resp_body,resp_blob) +VALUES(?,?,?,?,?,?,?)`, item.ID, item.ReqHead, item.ReqBody, nullIfEmpty(item.ReqBlob), + item.RespHead, item.RespBody, nullIfEmpty(item.RespBlob)); err != nil { + return imported, err + } + for _, hash := range []string{item.ReqBlob, item.RespBlob} { + if hash != "" { + if _, err := tx.Exec(`INSERT OR IGNORE INTO blob_refs(hash,exchange_id) VALUES(?,?)`, hash, item.ID); err != nil { + return imported, err + } + } + } + if t.fts { + var rowID int64 + if err := tx.QueryRow(`SELECT rowid FROM exchanges WHERE id=?`, item.ID).Scan(&rowID); err != nil { + return imported, err + } + content := strings.Join([]string{item.URL, item.ReqHead, string(item.ReqBody), item.RespHead, string(item.RespBody)}, "\n") + if _, err := tx.Exec(`INSERT INTO ex_fts(rowid,content) VALUES(?,?)`, rowID, content); err != nil { + return imported, err + } + } + imported++ + } + if err := tx.Commit(); err != nil { + return imported, err + } + return imported, nil +} + +func uniqueArchiveHosts(hosts []string) []string { + seen := make(map[string]struct{}, len(hosts)) + out := make([]string, 0, len(hosts)) + for _, host := range hosts { + host = strings.ToLower(strings.TrimSpace(host)) + if host == "" { + continue + } + if _, exists := seen[host]; exists { + continue + } + seen[host] = struct{}{} + out = append(out, host) + } + sort.Strings(out) + return out +} diff --git a/traffic/archive_test.go b/traffic/archive_test.go new file mode 100644 index 0000000..222864f --- /dev/null +++ b/traffic/archive_test.go @@ -0,0 +1,158 @@ +package traffic + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestTrafficArchiveRoundTripAndCurrentRowWins(t *testing.T) { + trafficDir := t.TempDir() + tr, err := Open(trafficDir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + const ( + id = "archive-exchange-1" + host = "archive.example" + ) + requestHead := "POST /v1/test HTTP/1.1\nHost: archive.example" + responseHead := "HTTP/1.1 200 OK\nContent-Type: application/json" + 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, "POST", "/v1/test", "https://archive.example/v1/test", 200, + "application/json", 7, 11, ""); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body) +VALUES(?,?,?,?,?)`, id, requestHead, []byte("payload"), responseHead, []byte(`{"ok":true}`)); err != nil { + t.Fatal(err) + } + + archiveDir := filepath.Join(t.TempDir(), "traffic") + count, err := tr.ExportHosts([]string{host, host}, archiveDir) + if err != nil || count != 1 { + t.Fatalf("ExportHosts count=%d err=%v", count, err) + } + if _, err := os.Stat(filepath.Join(archiveDir, "traffic.json")); 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 := stage.Commit(); err != nil { + t.Fatal(err) + } + if _, _, err := tr.Get(id); err == nil { + t.Fatal("deleted exchange remained readable") + } + imported, err := tr.ImportArchive(archiveDir) + if err != nil || imported != 1 { + t.Fatalf("ImportArchive imported=%d err=%v", imported, err) + } + req, resp, err := tr.Get(id) + if err != nil || !strings.Contains(req, "payload") || !strings.Contains(resp, `{"ok":true}`) { + t.Fatalf("restored exchange req=%q resp=%q err=%v", req, resp, err) + } + if _, err := tr.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id=?`, []byte(`{"current":true}`), id); err != nil { + t.Fatal(err) + } + imported, err = tr.ImportArchive(archiveDir) + if err != nil || imported != 0 { + t.Fatalf("idempotent import imported=%d err=%v", imported, err) + } + _, resp, err = tr.Get(id) + if err != nil || !strings.Contains(resp, `{"current":true}`) || strings.Contains(resp, `{"ok":true}`) { + t.Fatalf("archive overwrote current exchange resp=%q err=%v", resp, err) + } +} + +func TestRecoverArchiveHostDeleteStageRollsBackBeforePostgresCommit(t *testing.T) { + trafficDir := t.TempDir() + tr, err := Open(trafficDir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + const host = "rollback-archive.example" + insertArchiveRecoveryExchange(t, tr, "rollback-exchange", host) + legacyDir := filepath.Join(trafficDir, sanitize(host)) + if err := os.MkdirAll(legacyDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(legacyDir, "request.http"), []byte("request"), 0o600); err != nil { + t.Fatal(err) + } + stage, err := tr.StageDeleteHostsExactForArchive([]string{host}, 17, 42) + if err != nil { + t.Fatal(err) + } + simulateTrafficStageProcessExit(t, stage) + if err := tr.RecoverHostDeleteStages(func(id, taskID int64) (bool, error) { + if id != 17 || taskID != 42 { + t.Fatalf("archive id=%d task=%d, want 17/42", id, taskID) + } + return false, nil + }); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(legacyDir, "request.http")); err != nil { + t.Fatalf("legacy traffic tree not restored: %v", err) + } + var count int + if err := tr.DB().QueryRow(`SELECT count(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("exchange count=%d err=%v, want 1", count, err) + } +} + +func TestRecoverArchiveHostDeleteStageCompletesAfterPostgresCommit(t *testing.T) { + trafficDir := t.TempDir() + tr, err := Open(trafficDir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + const host = "committed-archive.example" + insertArchiveRecoveryExchange(t, tr, "committed-exchange", host) + legacyDir := filepath.Join(trafficDir, sanitize(host)) + if err := os.MkdirAll(legacyDir, 0o755); err != nil { + t.Fatal(err) + } + stage, err := tr.StageDeleteHostsExactForArchive([]string{host}, 18, 43) + if err != nil { + t.Fatal(err) + } + simulateTrafficStageProcessExit(t, stage) + if err := tr.RecoverHostDeleteStages(func(id, taskID int64) (bool, error) { return id == 18 && taskID == 43, nil }); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(legacyDir); !os.IsNotExist(err) { + t.Fatalf("committed legacy traffic tree remains: %v", err) + } + var count int + if err := tr.DB().QueryRow(`SELECT count(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 0 { + t.Fatalf("exchange count=%d err=%v, want 0", count, err) + } +} + +func insertArchiveRecoveryExchange(t *testing.T, tr *Traffic, id, host string) { + t.Helper() + 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", "/", "https://"+host+"/", 200, "text/plain", 0, 0, ""); err != nil { + t.Fatal(err) + } +} + +func simulateTrafficStageProcessExit(t *testing.T, stage *HostDeleteStage) { + t.Helper() + if err := stage.tx.Rollback(); err != nil { + t.Fatal(err) + } + stage.done = true + stage.traffic.wmu.Unlock() +} diff --git a/traffic/traffic.go b/traffic/traffic.go index e827443..6dfbce3 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -205,7 +205,7 @@ func (t *Traffic) Close() error { t.reaping.Wait() return t.db.Close() } -func (t *Traffic) DB() *sql.DB { return t.db } +func (t *Traffic) DB() *sql.DB { return t.db } // sink is the go-mitmproxy addon that records completed exchanges. type sink struct { @@ -870,6 +870,21 @@ type stagedTrafficPath struct { staged string } +type hostDeleteStageJournal struct { + Version int `json:"version"` + ArchiveID int64 `json:"archive_id,omitempty"` + TaskID int64 `json:"task_id,omitempty"` + Hosts []string `json:"hosts,omitempty"` + Moves []hostDeleteStageMove `json:"moves"` +} + +type hostDeleteStageMove struct { + Source string `json:"source"` + Staged string `json:"staged"` +} + +const hostDeleteStageJournalName = "journal.json" + // stageTrees moves host directories aside onto the same filesystem. The rename is // atomic and instant, which buys two things the unlink cannot: the deletion stays // reversible until the transaction commits, and the tree stops being visible to @@ -877,7 +892,12 @@ type stagedTrafficPath struct { // directories, so anything under _delete_staging is already out of the live set). // An empty stageDir is returned when there was nothing to stage. func (t *Traffic) stageTrees(dirs []string) (stageDir string, moves []stagedTrafficPath, err error) { + return t.stageTreesForArchive(dirs, nil, 0, 0) +} + +func (t *Traffic) stageTreesForArchive(dirs, hosts []string, archiveID, taskID int64) (stageDir string, moves []stagedTrafficPath, err error) { seen := make(map[string]struct{}, len(dirs)) + planned := make([]stagedTrafficPath, 0, len(dirs)) for _, source := range dirs { if _, dup := seen[source]; dup { continue @@ -889,16 +909,29 @@ func (t *Traffic) stageTrees(dirs []string) (stageDir string, moves []stagedTraf } return stageDir, moves, fmt.Errorf("检查历史流量目录 %s: %w", source, err) } - if stageDir == "" { - parent := filepath.Join(t.dir, "_delete_staging") - if err := os.MkdirAll(parent, 0o700); err != nil { - return stageDir, moves, fmt.Errorf("创建流量暂存目录: %w", err) - } - if stageDir, err = os.MkdirTemp(parent, "hosts-"); err != nil { - return stageDir, moves, fmt.Errorf("创建流量暂存目录: %w", err) - } - } - staged := filepath.Join(stageDir, fmt.Sprintf("%d-%s", len(moves), filepath.Base(source))) + planned = append(planned, stagedTrafficPath{source: source}) + } + if len(planned) == 0 { + return "", nil, nil + } + parent := filepath.Join(t.dir, "_delete_staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return "", nil, fmt.Errorf("创建流量暂存目录: %w", err) + } + if stageDir, err = os.MkdirTemp(parent, "hosts-"); err != nil { + return "", nil, fmt.Errorf("创建流量暂存目录: %w", err) + } + journal := hostDeleteStageJournal{Version: 1, ArchiveID: archiveID, TaskID: taskID, Hosts: uniqueArchiveHosts(hosts)} + for i := range planned { + planned[i].staged = filepath.Join(stageDir, fmt.Sprintf("%d-%s", i, filepath.Base(planned[i].source))) + journal.Moves = append(journal.Moves, hostDeleteStageMove{Source: planned[i].source, Staged: planned[i].staged}) + } + if err := writeHostDeleteStageJournal(filepath.Join(stageDir, hostDeleteStageJournalName), journal); err != nil { + _ = os.RemoveAll(stageDir) + return "", nil, err + } + for _, move := range planned { + source, staged := move.source, move.staged if err := os.Rename(source, staged); err != nil { return stageDir, moves, fmt.Errorf("移出历史流量目录 %s: %w", source, err) } @@ -907,6 +940,26 @@ func (t *Traffic) stageTrees(dirs []string) (stageDir string, moves []stagedTraf return stageDir, moves, nil } +func writeHostDeleteStageJournal(path string, journal hostDeleteStageJournal) error { + raw, err := json.Marshal(journal) + if err != nil { + return err + } + file, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) + if err != nil { + return err + } + if _, err := file.Write(raw); err != nil { + _ = file.Close() + return err + } + if err := file.Sync(); err != nil { + _ = file.Close() + return err + } + return file.Close() +} + // restoreTrees puts staged directories back where they came from, newest move // first, and drops the staging directory once everything is home. func restoreTrees(stageDir string, moves []stagedTrafficPath) error { @@ -992,6 +1045,20 @@ func (s *HostDeleteStage) Deleted() int64 { // 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) { + return t.stageDeleteHostsExact(hosts, 0, 0) +} + +// StageDeleteHostsExactForArchive ties the reversible traffic deletion to a +// persistent task archive. Startup recovery uses archiveCommitted to decide +// whether an interrupted stage must be completed or rolled back. +func (t *Traffic) StageDeleteHostsExactForArchive(hosts []string, archiveID, taskID int64) (*HostDeleteStage, error) { + if archiveID <= 0 || taskID <= 0 { + return nil, errors.New("archive and task ids must be positive") + } + return t.stageDeleteHostsExact(hosts, archiveID, taskID) +} + +func (t *Traffic) stageDeleteHostsExact(hosts []string, archiveID, taskID int64) (*HostDeleteStage, error) { t.wmu.Lock() stage := &HostDeleteStage{traffic: t} fail := func(cause error) (*HostDeleteStage, error) { @@ -1037,13 +1104,156 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error stage.deleted += n } - stage.stageDir, stage.moves, err = t.stageTrees(legacy) + stage.stageDir, stage.moves, err = t.stageTreesForArchive(legacy, unique, archiveID, taskID) if err != nil { return fail(err) } return stage, nil } +// RecoverHostDeleteStages resolves filesystem moves left by a process exit. +// SQLite rolls open transactions back on restart, while the supplied callback +// identifies task archives whose PostgreSQL compaction already committed. +func (t *Traffic) RecoverHostDeleteStages(archiveCommitted func(int64, int64) (bool, error)) error { + if t == nil { + return nil + } + t.wmu.Lock() + defer t.wmu.Unlock() + parent := filepath.Join(t.dir, "_delete_staging") + entries, err := os.ReadDir(parent) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + var errs []error + needsGC := false + for _, entry := range entries { + if !entry.IsDir() { + continue + } + stageDir := filepath.Join(parent, entry.Name()) + raw, err := os.ReadFile(filepath.Join(stageDir, hostDeleteStageJournalName)) + if os.IsNotExist(err) { + // Older versions did not persist enough information for lossless + // recovery. Keep the directory for manual inspection. + continue + } + if err != nil { + errs = append(errs, err) + continue + } + var journal hostDeleteStageJournal + if err := json.Unmarshal(raw, &journal); err != nil { + errs = append(errs, fmt.Errorf("读取流量暂存日志 %s: %w", stageDir, err)) + continue + } + if journal.Version != 1 { + errs = append(errs, fmt.Errorf("流量暂存日志 %s 的版本 %d 不受支持", stageDir, journal.Version)) + continue + } + moves := make([]stagedTrafficPath, 0, len(journal.Moves)) + for _, move := range journal.Moves { + if !pathWithin(t.dir, move.Source) || !pathWithin(stageDir, move.Staged) { + errs = append(errs, fmt.Errorf("流量暂存日志包含越界路径: %s", stageDir)) + moves = nil + break + } + if _, err := os.Lstat(move.Staged); err == nil { + moves = append(moves, stagedTrafficPath{source: move.Source, staged: move.Staged}) + } else if !os.IsNotExist(err) { + errs = append(errs, err) + moves = nil + break + } + } + if moves == nil { + continue + } + committed := false + if journal.ArchiveID > 0 { + if archiveCommitted == nil { + errs = append(errs, fmt.Errorf("流量归档 %d 无状态解析器", journal.ArchiveID)) + continue + } + committed, err = archiveCommitted(journal.ArchiveID, journal.TaskID) + if err != nil { + errs = append(errs, err) + continue + } + } else if len(journal.Hosts) > 0 { + committed, err = t.hostsHaveNoExchanges(journal.Hosts) + if err != nil { + errs = append(errs, err) + continue + } + } + if !committed { + if err := restoreTrees(stageDir, moves); err != nil { + errs = append(errs, err) + } + continue + } + if err := t.deleteArchivedHosts(journal.Hosts); err != nil { + errs = append(errs, err) + continue + } + if err := os.RemoveAll(stageDir); err != nil { + errs = append(errs, err) + continue + } + needsGC = true + } + if needsGC { + if err := t.gcBlobs(); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +func (t *Traffic) hostsHaveNoExchanges(hosts []string) (bool, error) { + for _, host := range hosts { + var count int + if err := t.db.QueryRow(`SELECT count(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil { + return false, err + } + if count > 0 { + return false, nil + } + } + return true, nil +} + +func (t *Traffic) deleteArchivedHosts(hosts []string) error { + tx, err := t.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck + for _, host := range hosts { + if _, err := t.deleteWhere(tx, `host=?`, host); err != nil { + return err + } + } + return tx.Commit() +} + +func pathWithin(root, candidate string) bool { + rootAbs, err := filepath.Abs(root) + if err != nil { + return false + } + candidateAbs, err := filepath.Abs(candidate) + if err != nil { + return false + } + relative, err := filepath.Rel(rootAbs, candidateAbs) + return err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator)) +} + // Rollback discards the staged deletion, leaving the traffic store untouched. func (s *HostDeleteStage) Rollback() error { if s == nil || s.done { diff --git a/web/src/app/(main)/function/findings/page.tsx b/web/src/app/(main)/function/findings/page.tsx index 0ba7988..5ad51ca 100644 --- a/web/src/app/(main)/function/findings/page.tsx +++ b/web/src/app/(main)/function/findings/page.tsx @@ -13,6 +13,7 @@ import { FileTextIcon, FlaskConicalIcon, InfoIcon, + SearchIcon, ShieldAlertIcon, Trash2Icon, TriangleAlertIcon, @@ -47,6 +48,7 @@ import { } from "@/components/ui/dialog"; import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; +import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { Label } from "@/components/ui/label"; import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; @@ -55,6 +57,7 @@ import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@ import { Textarea } from "@/components/ui/textarea"; import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { api } from "@/lib/api"; +import { getLocalStorageValue, setLocalStorageValue } from "@/lib/local-storage.client"; import { statusMeta } from "@/lib/status"; import type { Finding, FindingGroup, FindingStats, FindingStatus, Severity } from "@/lib/types"; import { cn } from "@/lib/utils"; @@ -73,6 +76,7 @@ const FINDING_STATUSES: FindingStatus[] = [ ]; const UNASSIGNED_TASK = "__unassigned__"; +const FINDING_LIST_PREFERENCE_KEY = "artex_finding_list_preferences"; interface GroupFindingsState { items: Finding[]; @@ -125,6 +129,8 @@ export default function FindingsPage() { const [vulnclass, setVulnclass] = React.useState("all"); const [task, setTask] = React.useState("all"); const [sort, setSort] = React.useState<"severity" | "time">("severity"); + const [search, setSearch] = React.useState(""); + const [query, setQuery] = React.useState(""); const [expanded, setExpanded] = React.useState(null); const [groups, setGroups] = React.useState([]); const [groupTotal, setGroupTotal] = React.useState(0); @@ -132,15 +138,54 @@ export default function FindingsPage() { const [groupFindings, setGroupFindings] = React.useState>({}); const [total, setTotal] = React.useState(0); const [stats, setStats] = React.useState(EMPTY_STATS); + const [statsLoaded, setStatsLoaded] = React.useState(false); + const [preferencesHydrated, setPreferencesHydrated] = React.useState(false); const [page, setPage] = React.useState(1); const [pageSize, setPageSize] = React.useState(10); const [deepenFinding, setDeepenFinding] = React.useState(null); const [deepenDescription, setDeepenDescription] = React.useState(""); const [deepening, setDeepening] = React.useState(false); - const filterFingerprint = `${severity}:${status}:${vulnclass}:${task}:${sort}`; + const filterFingerprint = JSON.stringify([severity, status, vulnclass, task, sort, query]); const activeFilterFingerprint = React.useRef(filterFingerprint); activeFilterFingerprint.current = filterFingerprint; + React.useEffect(() => { + const raw = getLocalStorageValue(FINDING_LIST_PREFERENCE_KEY); + if (raw) { + try { + const parsed = JSON.parse(raw) as { + severity?: unknown; + status?: unknown; + vulnclass?: unknown; + task?: unknown; + sort?: unknown; + }; + if (parsed.severity === "all" || SEVERITIES.includes(parsed.severity as Severity)) { + setSeverity(parsed.severity as "all" | Severity); + } + if (parsed.status === "all" || FINDING_STATUSES.includes(parsed.status as FindingStatus)) { + setStatus(parsed.status as "all" | FindingStatus); + } + if (typeof parsed.vulnclass === "string" && parsed.vulnclass) setVulnclass(parsed.vulnclass); + if (typeof parsed.task === "string" && parsed.task) setTask(parsed.task); + if (parsed.sort === "severity" || parsed.sort === "time") setSort(parsed.sort); + } catch { + // Ignore malformed or legacy preferences and retain the defaults. + } + } + setPreferencesHydrated(true); + }, []); + + React.useEffect(() => { + if (!preferencesHydrated) return; + setLocalStorageValue(FINDING_LIST_PREFERENCE_KEY, JSON.stringify({ severity, status, vulnclass, task, sort })); + }, [preferencesHydrated, severity, sort, status, task, vulnclass]); + + React.useEffect(() => { + const timer = window.setTimeout(() => setQuery(search.trim()), 300); + return () => window.clearTimeout(timer); + }, [search]); + const setFindings = React.useCallback((update: (current: Finding[]) => Finding[]) => { setGroupFindings((current) => { const next: Record = {}; @@ -191,7 +236,7 @@ export default function FindingsPage() { await api.exportFindings({ format: exportFormat, scope: exportScope, - filters: { severity, status, vulnclass, task, sort }, + filters: { severity, status, vulnclass, task, query, sort }, ids: [...selectedIds], }); setExportOpen(false); @@ -224,6 +269,7 @@ export default function FindingsPage() { status, vulnclass, task, + query, sort, }); if (request !== groupsRequest.current || activeFilterFingerprint.current !== requestFilter) return; @@ -233,7 +279,7 @@ export default function FindingsPage() { } catch { // Polling keeps the last successful snapshot visible. } - }, [filterFingerprint, page, pageSize, severity, status, vulnclass, task, sort]); + }, [filterFingerprint, page, pageSize, severity, status, vulnclass, task, query, sort]); const loadGroup = React.useCallback( async (key: string, groupPage: number, groupPageSize: number) => { @@ -260,6 +306,7 @@ export default function FindingsPage() { status, vulnclass, task: key, + query, sort, }); if (groupRequests.current[key] !== request || activeFilterFingerprint.current !== requestFilter) return; @@ -291,7 +338,7 @@ export default function FindingsPage() { })); } }, - [filterFingerprint, severity, status, vulnclass, sort], + [filterFingerprint, severity, status, vulnclass, query, sort], ); React.useEffect(() => { @@ -367,7 +414,10 @@ export default function FindingsPage() { api .findingStats() .then((s) => { - if (alive) setStats(s); + if (alive) { + setStats(s); + setStatsLoaded(true); + } }) .catch(() => { // Keep the previous aggregate snapshot until the next poll. @@ -381,6 +431,18 @@ export default function FindingsPage() { }; }, []); + React.useEffect(() => { + if (!statsLoaded) return; + if (vulnclass !== "all" && !stats.vulnclasses.includes(vulnclass)) setVulnclass("all"); + if ( + task !== "all" && + task !== UNASSIGNED_TASK && + !(stats.tasks ?? []).some((option) => String(option.id) === task) + ) { + setTask("all"); + } + }, [stats, statsLoaded, task, vulnclass]); + // updateStatus optimistically flips one finding's triage state, reverting on error. const updateStatus = React.useCallback( async (f: Finding, next: FindingStatus) => { @@ -566,6 +628,19 @@ export default function FindingsPage() {
+ + setSearch(event.target.value)} + placeholder="检索漏洞内容" + aria-label="检索漏洞内容" + /> + + + + (null); @@ -198,6 +211,7 @@ function TaskDetailInner() { const [tab, setTab] = React.useState("sessions"); const [interceptPendingCount, setInterceptPendingCount] = React.useState(0); const [profiles, setProfiles] = React.useState([]); + const [archiving, setArchiving] = React.useState(false); React.useEffect(() => { api @@ -272,6 +286,19 @@ function TaskDetailInner() { } } + async function archiveTask() { + if (!task || archiving) return; + setArchiving(true); + try { + await api.archiveTask(task.id); + toast.success("任务已加入归档队列"); + router.push("/function/tasks"); + } catch (error) { + toast.error(`归档失败：${(error as Error).message}`); + setArchiving(false); + } + } + if (!task) { return (
@@ -289,6 +316,12 @@ function TaskDetailInner() { const completed = task.status === "done"; const terminal = ["done", "failed", "timeout"].includes(task.status); + const archiveLifecycleEligible = terminal || paused || task.status === "paused"; + const canArchive = archiveLifecycleEligible && !task.archive_blocked_by_task_id; + let archiveDisabledReason = task.queued ? "排队中的任务必须先暂停" : "运行中的任务必须先暂停"; + if (archiveLifecycleEligible && task.archive_blocked_by_task_id) { + archiveDisabledReason = `任务被未归档任务 #${task.archive_blocked_by_task_id} 直接继承，请先归档依赖任务`; + } const engineMode = paused ? "paused" : (task.engine_mode ?? "idle"); let controlVariant: "default" | "secondary" | "outline" = "outline"; let controlIcon = ; @@ -302,6 +335,16 @@ function TaskDetailInner() { controlIcon = ; controlLabel = "恢复"; } + const archiveTrigger = ( + + ); return ( @@ -323,6 +366,31 @@ function TaskDetailInner() { + {canArchive ? ( + + {archiveTrigger} + + + 归档任务 #{task.id}？ + + 任务图谱、关联记录、独占资产与流量、工作文件和 LLM + 历史将压缩到冷存储。归档完成后可在任务列表的“已归档”页还原。 + + + + 取消 + void archiveTask()}>确认归档 + + + + ) : ( + + + {archiveTrigger} + + {archiveDisabledReason} + + )} - )} -
- - - - {filtered.length}/{tasks.length} 条 - - {selectedIds.size > 0 && ( - <> - 已选 {selectedIds.size} 个 - - {pausableTaskIDs.length > 0 && ( - + )} +
+ + + + {filtered.length}/{tasks.length} 条 + + {selectedIds.size > 0 && ( + <> + 已选 {selectedIds.size} 个 + + {pausableTaskIDs.length > 0 && ( + )} - 暂停 {pausableTaskIDs.length} - - )} - {resumableTaskIDs.length > 0 && ( - )} - 继续 {resumableTaskIDs.length} - - )} - - - - )} - - - - - - {tasks.length === 0 ? ( -
- 暂无任务，点击右上角「新建任务」开始。 -
- ) : filtered.length === 0 ? ( -
- 没有匹配的任务。 -
- ) : ( - - - - - toggleSelectedPage(checked === true)} - aria-label="选择本页全部任务" + {archivableTaskIDs.length > 0 && ( + + + 归档 {archivableTaskIDs.length} + + } + /> + )} + - - - 名称 - 描述 - 目标 - - 目标进度 - - - Token - - 操作 - - - - - {paginated.map((task) => ( - - ))} - -
- )} - { - setPageSize(nextPageSize); - setPage(1); - }} - /> - - + + + )} + + + + + + {tasks.length === 0 ? ( +
+ 暂无任务，点击右上角「新建任务」开始。 +
+ ) : filtered.length === 0 ? ( +
+ 没有匹配的任务。 +
+ ) : ( + + + + + toggleSelectedPage(checked === true)} + aria-label="选择本页全部任务" + /> + + + 名称 + 描述 + 目标 + + 目标进度 + + + Token + + 操作 + + + + + {paginated.map((task) => ( + + ))} + +
+ )} + { + setPageSize(nextPageSize); + setPage(1); + }} + /> + + + + + + + ); } @@ -1046,6 +1157,7 @@ const TaskRow = React.memo(function TaskRow({ onControl, onRename, onTogglePinned, + onArchive, selected, onSelectedChange, }: { @@ -1055,6 +1167,7 @@ const TaskRow = React.memo(function TaskRow({ onControl: (id: string, action: "pause" | "resume") => Promise; onRename: (task: Task, name: string) => Promise; onTogglePinned: (task: Task) => Promise; + onArchive: (task: Task) => Promise; selected: boolean; onSelectedChange: (id: string, checked: boolean) => void; }) { @@ -1140,6 +1253,7 @@ const TaskRow = React.memo(function TaskRow({ + @@ -1147,13 +1261,7 @@ const TaskRow = React.memo(function TaskRow({ ); }); -function TaskNameEditor({ - task, - onRename, -}: { - task: Task; - onRename: (task: Task, name: string) => Promise; -}) { +function TaskNameEditor({ task, onRename }: { task: Task; onRename: (task: Task, name: string) => Promise }) { const [editing, setEditing] = React.useState(false); const [name, setName] = React.useState(task.name ?? ""); const [saving, setSaving] = React.useState(false); @@ -1253,13 +1361,7 @@ function taskPinIcon(pinning: boolean, pinned: boolean) { return ; } -function TaskPinAction({ - task, - onTogglePinned, -}: { - task: Task; - onTogglePinned: (task: Task) => Promise; -}) { +function TaskPinAction({ task, onTogglePinned }: { task: Task; onTogglePinned: (task: Task) => Promise }) { const [pinning, setPinning] = React.useState(false); const pinned = taskIsPinned(task); @@ -1326,6 +1428,553 @@ function TaskControlButton({ ); } +function archiveBlockReason(task: Task): string { + if (task.queued) return "排队中的任务必须先暂停"; + if (!ARCHIVABLE_STATUSES.has(task.status)) return "运行中或尚未结束的任务必须先暂停"; + if (task.archive_blocked_by_task_id) { + return `任务被未归档任务 #${task.archive_blocked_by_task_id} 直接继承，请先归档依赖任务`; + } + return ""; +} + +function ArchiveConfirmDialog({ + count, + trigger, + onConfirm, +}: { + count: number; + trigger: React.ReactNode; + onConfirm: () => Promise; +}) { + const [open, setOpen] = React.useState(false); + const [pending, setPending] = React.useState(false); + return ( + !pending && setOpen(next)}> + {trigger} + + + {count === 1 ? "归档任务" : `归档 ${count} 个任务`} + + 归档会停止任务调度，将图谱、LLM 历史、文件以及独占资产和流量压缩到本地冷存储。归档完成后可从“已归档”中还原。 + + + + 取消 + { + event.preventDefault(); + setPending(true); + try { + await onConfirm(); + setOpen(false); + } finally { + setPending(false); + } + }} + > + {pending && } + 确认归档 + + + + + ); +} + +function TaskArchiveAction({ task, onArchive }: { task: Task; onArchive: (task: Task) => Promise }) { + const reason = archiveBlockReason(task); + const trigger = ( + + ); + if (reason) { + return ( + + + {trigger} + + {reason} + + ); + } + return onArchive(task)} />; +} + +const ARCHIVE_PROCESSING_STATES = new Set([ + "archive_queued", + "archiving", + "restore_queued", + "restoring", + "delete_queued", + "deleting", +]); + +const ARCHIVE_FAILED_STATES = new Set(["archive_failed", "restore_failed", "delete_failed"]); + +function archiveStateLabel(state: TaskArchiveState): string { + const labels: Record = { + archive_queued: "等待归档", + archiving: "归档中", + archive_failed: "归档失败", + ready: "可还原", + restore_queued: "等待还原", + restoring: "还原中", + restore_failed: "还原失败", + delete_queued: "等待删除", + deleting: "删除中", + delete_failed: "删除失败", + }; + return labels[state]; +} + +function ArchiveStateBadge({ state }: { state: TaskArchiveState }) { + let variant: "default" | "secondary" | "destructive" | "outline" = "outline"; + if (state === "ready") variant = "secondary"; + if (ARCHIVE_PROCESSING_STATES.has(state)) variant = "default"; + if (ARCHIVE_FAILED_STATES.has(state)) variant = "destructive"; + return {archiveStateLabel(state)}; +} + +function formatArchiveBytes(bytes: number): string { + if (bytes <= 0) return "—"; + if (bytes >= 1024 ** 3) return `${(bytes / 1024 ** 3).toFixed(1)} GB`; + if (bytes >= 1024 ** 2) return `${(bytes / 1024 ** 2).toFixed(1)} MB`; + if (bytes >= 1024) return `${(bytes / 1024).toFixed(1)} KB`; + return `${bytes} B`; +} + +function archiveCompressionLabel(archive: TaskArchive): string { + if (archive.original_size <= 0 || archive.compressed_size <= 0) return "—"; + const saved = Math.max(0, 100 - (archive.compressed_size / archive.original_size) * 100); + return `${formatArchiveBytes(archive.compressed_size)} · 节省 ${saved.toFixed(0)}%`; +} + +function archiveDataTotal(archive: TaskArchive): number { + return Object.values(archive.data_counts).reduce((sum, value) => sum + (Number(value) || 0), 0); +} + +function archiveDate(value?: string): string { + if (!value) return "—"; + const date = new Date(value); + if (Number.isNaN(date.getTime())) return "—"; + return date.toLocaleString(); +} + +function ArchiveDeleteDialog({ + archives, + onConfirm, + trigger, +}: { + archives: TaskArchive[]; + onConfirm: () => Promise; + trigger: React.ReactNode; +}) { + const [open, setOpen] = React.useState(false); + const [pending, setPending] = React.useState(false); + const rows = archives.reduce((sum, archive) => sum + archiveDataTotal(archive), 0); + const bytes = archives.reduce((sum, archive) => sum + archive.compressed_size, 0); + return ( + !pending && setOpen(next)}> + {trigger} + + + 永久删除 {archives.length} 个任务归档？ + + 将永久删除约 {formatArchiveBytes(bytes)} 的归档包及 {rows.toLocaleString()} 条关联数据快照。此操作不可恢复。 + + + + 取消 + { + event.preventDefault(); + setPending(true); + try { + await onConfirm(); + setOpen(false); + } finally { + setPending(false); + } + }} + > + {pending && } + 永久删除 + + + + + ); +} + +function TaskArchivesPanel({ onChanged }: { onChanged: () => void }) { + const [archives, setArchives] = React.useState([]); + const [query, setQuery] = React.useState(""); + const [stateFilter, setStateFilter] = React.useState("all"); + const [page, setPage] = React.useState(1); + const [pageSize, setPageSize] = React.useState(20); + const [total, setTotal] = React.useState(0); + const [selected, setSelected] = React.useState>(() => new Set()); + const [loading, setLoading] = React.useState(true); + const pendingRestoreIDs = React.useRef(new Set()); + + const load = React.useCallback(async () => { + try { + const result = await api.taskArchives({ + page, + size: pageSize, + q: query.trim(), + state: stateFilter === "all" ? undefined : stateFilter, + }); + setArchives(result.items); + setTotal(result.total); + setSelected((current) => { + const visible = new Set(result.items.map((archive) => archive.id)); + return new Set([...current].filter((id) => visible.has(id))); + }); + const pending = [...pendingRestoreIDs.current]; + if (pending.length > 0) { + const states = await Promise.allSettled(pending.map((id) => api.taskArchive(id))); + let restored = false; + states.forEach((state, index) => { + if (state.status !== "rejected" || !(state.reason instanceof Error)) return; + if (!state.reason.message.includes("归档不存在")) return; + pendingRestoreIDs.current.delete(pending[index]); + restored = true; + }); + if (restored) onChanged(); + } + } catch (error) { + toast.error(`读取归档列表失败：${(error as Error).message}`); + } finally { + setLoading(false); + } + }, [onChanged, page, pageSize, query, stateFilter]); + + React.useEffect(() => { + void load(); + const timer = setInterval(() => void load(), 2_000); + return () => clearInterval(timer); + }, [load]); + + const selectable = React.useMemo( + () => archives.filter((archive) => !ARCHIVE_PROCESSING_STATES.has(archive.state)), + [archives], + ); + const selectedArchives = React.useMemo( + () => archives.filter((archive) => selected.has(archive.id)), + [archives, selected], + ); + const restorable = selectedArchives.filter( + (archive) => archive.state === "ready" || archive.state === "restore_failed", + ); + const deletable = selectedArchives.filter( + (archive) => archive.state === "ready" || archive.state === "delete_failed", + ); + const selectedAll = selectable.length > 0 && selectable.every((archive) => selected.has(archive.id)); + const selectedSome = selectable.some((archive) => selected.has(archive.id)); + + const afterAction = React.useCallback(() => { + setSelected(new Set()); + void load(); + onChanged(); + }, [load, onChanged]); + + async function restoreMany(items: TaskArchive[]) { + try { + const result = + items.length === 1 + ? { + items: [ + { + id: String(items[0].id), + ok: true, + queued: true, + archive_id: (await api.restoreTaskArchive(items[0].id)).id, + }, + ], + } + : await api.restoreTaskArchives(items.map((archive) => archive.id)); + const succeeded = result.items.filter((item) => item.ok).length; + const failed = result.items.length - succeeded; + for (const item of result.items) { + const archiveID = Number(item.archive_id ?? item.id); + if (item.ok && Number.isSafeInteger(archiveID) && archiveID > 0) pendingRestoreIDs.current.add(archiveID); + } + if (succeeded > 0) toast.success(`已将 ${succeeded} 个任务加入还原队列`); + if (failed > 0) toast.error(`${failed} 个任务无法还原`); + afterAction(); + } catch (error) { + toast.error(`还原失败：${(error as Error).message}`); + } + } + + async function deleteMany(items: TaskArchive[]) { + try { + const result = + items.length === 1 + ? { + items: [ + { + id: String(items[0].id), + ok: true, + queued: true, + archive_id: (await api.deleteTaskArchive(items[0].id)).id, + }, + ], + } + : await api.deleteTaskArchives(items.map((archive) => archive.id)); + const succeeded = result.items.filter((item) => item.ok).length; + const failed = result.items.length - succeeded; + if (succeeded > 0) toast.success(`已将 ${succeeded} 个归档加入永久删除队列`); + if (failed > 0) toast.error(`${failed} 个归档无法删除`); + afterAction(); + } catch (error) { + toast.error(`永久删除失败：${(error as Error).message}`); + throw error; + } + } + + async function retry(archive: TaskArchive) { + try { + if (archive.state === "archive_failed") await api.archiveTask(String(archive.task_id)); + if (archive.state === "restore_failed") await api.restoreTaskArchive(archive.id); + if (archive.state === "delete_failed") await api.deleteTaskArchive(archive.id); + toast.success("已重新加入处理队列"); + afterAction(); + } catch (error) { + toast.error(`重试失败：${(error as Error).message}`); + } + } + + return ( + + +
+
+ + { + setQuery(event.target.value); + setPage(1); + }} + placeholder="搜索任务名称 / 描述 / ID" + className="pl-8" + /> +
+ + {total} 个归档 + {selectedArchives.length > 0 && ( + <> + 已选 {selectedArchives.length} 个 + {restorable.length > 0 && ( + + )} + {deletable.length > 0 && ( + deleteMany(deletable)} + trigger={ + + } + /> + )} + + )} +
+ {loading ? ( +
+ +
+ ) : archives.length === 0 ? ( + + + 暂无任务归档 + 暂停或终态任务可从当前任务列表归档。 + + + ) : ( + + + + + + setSelected(checked === true ? new Set(selectable.map((archive) => archive.id)) : new Set()) + } + aria-label="选择当前页归档" + /> + + 任务 + 原状态 + 分类 + 归档时间 + 压缩大小 + 数据量 + 处理状态 + 操作 + + + + {archives.map((archive) => { + const processing = ARCHIVE_PROCESSING_STATES.has(archive.state); + const canRestore = archive.state === "ready" || archive.state === "restore_failed"; + const canDelete = archive.state === "ready" || archive.state === "delete_failed"; + return ( + + + + setSelected((current) => { + const next = new Set(current); + if (checked === true) next.add(archive.id); + else next.delete(archive.id); + return next; + }) + } + aria-label={`选择任务归档 #${archive.task_id}`} + /> + + +
+ + {archive.task_name || archive.task_description || `任务 #${archive.task_id}`} + + + #{archive.task_id} · {archive.task_description} + +
+ + + + + {archive.category_name || "未分类"} + + {archiveDate(archive.archived_at || archive.requested_at)} + + + {archiveCompressionLabel(archive)} + + {archiveDataTotal(archive).toLocaleString()} + +
+ + {processing && } + {archive.error && ( +

{archive.error}

+ )} + {[...new Set(archive.warnings ?? [])].map((warning) => ( +

+ {warning} +

+ ))} + + {archive.phase} · {archive.progress}% + +
+ + +
+ {ARCHIVE_FAILED_STATES.has(archive.state) && ( + + )} + {canRestore && ( + + )} + {canDelete && ( + deleteMany([archive])} + trigger={ + + } + /> + )} +
+ + + ); + })} + +
+ )} + { + setPageSize(size); + setPage(1); + }} + /> + + + ); +} + const emptyDeleteOptions = (): DeleteTaskOptions => ({ delete_assets: false, delete_traffic: false, diff --git a/web/src/app/(main)/function/traffic/page.tsx b/web/src/app/(main)/function/traffic/page.tsx index b565110..df0c5a3 100644 --- a/web/src/app/(main)/function/traffic/page.tsx +++ b/web/src/app/(main)/function/traffic/page.tsx @@ -3,6 +3,8 @@ import * as React from "react"; import { + ArrowDownWideNarrowIcon, + ArrowUpNarrowWideIcon, CheckIcon, ChevronLeftIcon, ChevronRightIcon, @@ -291,6 +293,7 @@ function HttpCodeBlock({ raw }: { raw: string }) { // single page, which would only ever list the methods on that page. const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]; const PAGE_SIZES = [25, 50, 100, 200]; +type HostCountSortDirection = "asc" | "desc"; export default function TrafficPage() { const [page, setPage] = React.useState(0); @@ -309,6 +312,7 @@ export default function TrafficPage() { const [hosts, setHosts] = React.useState([]); // target picker const [selectedHosts, setSelectedHosts] = React.useState([]); // checked in picker const [pickerOpen, setPickerOpen] = React.useState(false); + const [hostCountSortDirection, setHostCountSortDirection] = React.useState("desc"); const [deleteMode, setDeleteMode] = React.useState<"filter" | "selected" | null>(null); // null = dialog closed const [deleting, setDeleting] = React.useState(false); @@ -367,6 +371,14 @@ export default function TrafficPage() { // Delete traffic for the current host filter (substring) or the checked // hosts (exact batch), then refetch. const allSelected = hosts.length > 0 && hosts.every((h) => selectedHosts.includes(h.host)); + const sortedHosts = React.useMemo( + () => + [...hosts].sort((a, b) => { + const countOrder = hostCountSortDirection === "asc" ? a.count - b.count : b.count - a.count; + return countOrder || a.host.localeCompare(b.host); + }), + [hosts, hostCountSortDirection], + ); const confirmDelete = () => { setDeleting(true); @@ -455,28 +467,49 @@ export default function TrafficPage() {
按目标批量删除 - {hosts.length > 0 && ( - - )} +
+ {hosts.length > 0 && ( + + + + + 按数据包数量{hostCountSortDirection === "desc" ? "倒序" : "正序"} + + )} + {hosts.length > 0 && ( + + )} +
{hosts.length === 0 ? (
暂无流量记录
) : ( - hosts.map((h, index) => ( + sortedHosts.map((h, index) => (
)} +
@@ -549,7 +571,7 @@ function AnswerBlock({ step, getDetail }: { step: Activity; getDetail: (seq: num }; }, [inView, step.seq, getDetail, step.summary]); return ( -
+
)}
+
); } diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index c28c328..9a9f3b7 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -10,6 +10,7 @@ import type { Agent, AgentDetail, AgentTrigger, + ArchiveBatchItem, Asset, Audit, BatchCategoryItem, @@ -59,6 +60,8 @@ import type { SSTask, Stats, Task, + TaskArchive, + TaskArchivePage, TaskAssetMutation, TaskAssetScopeMutation, TaskCategory, @@ -253,6 +256,27 @@ export const api = { post<{ id: string; paused: boolean; queued: boolean; status: string }>(`/tasks/${id}/control`, { action }), controlTasksBatch: (taskIds: string[], action: "pause" | "resume") => post<{ items: BatchControlItem[] }>("/tasks/control/batch", { task_ids: taskIds, action }), + taskArchives: (input?: { page?: number; size?: number; q?: string; state?: string }) => { + const query = new URLSearchParams(); + query.set("page", String(input?.page ?? 1)); + query.set("size", String(input?.size ?? 20)); + if (input?.q) query.set("q", input.q); + if (input?.state) query.set("state", input.state); + return get(`/task-archives?${query.toString()}`).then((response) => ({ + ...response, + items: arr(response.items), + })); + }, + taskArchive: (id: number) => get(`/task-archives/${id}`), + archiveTask: (id: string) => post(`/tasks/${id}/archive`), + archiveTasks: (taskIds: string[]) => + post<{ items: ArchiveBatchItem[] }>("/tasks/archive/batch", { task_ids: taskIds }), + restoreTaskArchive: (id: number) => post(`/task-archives/${id}/restore`), + restoreTaskArchives: (archiveIds: number[]) => + post<{ items: ArchiveBatchItem[] }>("/task-archives/restore/batch", { archive_ids: archiveIds }), + deleteTaskArchive: (id: number) => del(`/task-archives/${id}`), + deleteTaskArchives: (archiveIds: number[]) => + post<{ items: ArchiveBatchItem[] }>("/task-archives/delete/batch", { archive_ids: archiveIds }), controlIntent: (taskId: string, intentId: string, action: "pause" | "resume" | "cancel", reason?: string) => post<{ id: number; @@ -429,6 +453,7 @@ export const api = { if (q.status && q.status !== "all") p.set("status", q.status); if (q.vulnclass && q.vulnclass !== "all") p.set("vulnclass", q.vulnclass); if (q.task && q.task !== "all") p.set("task_id", q.task); + if (q.query?.trim()) p.set("q", q.query.trim()); if (q.sort) p.set("sort", q.sort); return get(`/exploration/findings?${p.toString()}`); }, @@ -438,6 +463,7 @@ export const api = { if (q.status && q.status !== "all") p.set("status", q.status); if (q.vulnclass && q.vulnclass !== "all") p.set("vulnclass", q.vulnclass); if (q.task && q.task !== "all") p.set("task_id", q.task); + if (q.query?.trim()) p.set("q", q.query.trim()); if (q.sort) p.set("sort", q.sort); return get(`/exploration/findings/groups?${p.toString()}`); }, @@ -459,6 +485,7 @@ export const api = { if (f.status && f.status !== "all") p.set("status", f.status); if (f.vulnclass && f.vulnclass !== "all") p.set("vulnclass", f.vulnclass); if (f.task && f.task !== "all") p.set("task_id", f.task); + if (f.query?.trim()) p.set("q", f.query.trim()); if (f.sort) p.set("sort", f.sort); } const token = getToken(); diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 2b6f536..6002dae 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -10,6 +10,7 @@ import { normalizeCompanyScopeValue, } from "../company-scope"; import type { + ArchiveBatchItem, Asset, BatchControlItem, Company, @@ -18,6 +19,7 @@ import type { IntentAsset, ScopeRow, Task, + TaskArchive, TaskAssetMutation, TaskAssetScopeMutation, TaskCategory, @@ -40,6 +42,15 @@ const mockConversations = structuredClone(D.conversations); const mockIntents = structuredClone(D.intents); const mockCompanies = structuredClone(D.companies); const mockAssets = structuredClone(D.assets); +type MockTaskArchiveSnapshot = { + task: Task; + numericTaskID: number; + assetIDs: number[]; + assetSources: Array<[number, MockTaskAssetSource]>; +}; +type MockTaskArchive = TaskArchive & { snapshot: MockTaskArchiveSnapshot }; +const mockTaskArchives: MockTaskArchive[] = []; +let nextMockTaskArchiveID = 1; const mockTaskAssetIDs = new Map(D.tasks.map((task, index) => [task.id, index + 1])); const mockTaskScopes = new Map(); type MockTaskAssetSource = Pick; @@ -51,6 +62,195 @@ function mockTaskAssetSourceKey(taskID: string, assetID: number): string { return JSON.stringify([taskID, assetID]); } +function publicMockTaskArchive(archive: MockTaskArchive): TaskArchive { + const { snapshot: _snapshot, ...item } = archive; + return structuredClone(item); +} + +function mockArchiveTaskID(taskID: string): number { + const numeric = Number(taskID); + if (Number.isSafeInteger(numeric) && numeric > 0) return numeric; + return mockTaskAssetID(taskID) ?? 0; +} + +function mockTaskArchiveBlocker(taskID: string): string | undefined { + return mockTasks.find( + (candidate) => + candidate.id !== taskID && + (candidate.source_task_ids ?? []).map(String).includes(taskID) && + !mockTaskArchives.some( + (archive) => + archive.snapshot.task.id === candidate.id && + (archive.state === "archive_queued" || archive.state === "archiving"), + ), + )?.id; +} + +function publicMockTask(task: Task): Task { + const item = structuredClone(task); + const blocker = mockTaskArchiveBlocker(task.id); + if (blocker) item.archive_blocked_by_task_id = blocker; + else delete item.archive_blocked_by_task_id; + return item; +} + +function mockArchiveTask(taskID: string): MockTaskArchive { + const task = mockTasks.find((item) => item.id === taskID); + if (!task) throw new Error("任务不存在"); + if (!["paused", "done", "failed", "timeout"].includes(task.status) && !task.paused) { + throw new Error(task.queued ? "排队中的任务必须先暂停" : "运行中的任务必须先暂停"); + } + const existing = mockTaskArchives.find((item) => item.task_id === mockArchiveTaskID(taskID)); + if (existing) throw new Error("任务已经在归档队列中"); + const dependent = mockTasks.find( + (candidate) => + candidate.id !== taskID && + (candidate.source_task_ids ?? []).map(String).includes(taskID) && + !mockTaskArchives.some( + (archive) => + archive.snapshot.task.id === candidate.id && + (archive.state === "archive_queued" || archive.state === "archiving"), + ), + ); + if (dependent) throw new Error(`任务被未归档任务 #${dependent.id} 直接继承，暂不能归档`); + + const numericTaskID = mockArchiveTaskID(taskID); + const assetIDs = mockAssets.filter((asset) => asset.task_ids.includes(numericTaskID)).map((asset) => asset.id); + const assetSources: Array<[number, MockTaskAssetSource]> = []; + for (const assetID of assetIDs) { + const source = mockTaskAssetSources.get(mockTaskAssetSourceKey(taskID, assetID)); + if (source) assetSources.push([assetID, structuredClone(source)]); + } + const findings = mockFindings.filter((finding) => finding.task_id === taskID); + const llmRecords = mockLLMRecords.filter((record) => record.task_id === taskID); + const tokens = task.tokens ?? { input_tokens: 0, output_tokens: 0, cache_read_tokens: 0, cache_write_tokens: 0 }; + const now = new Date().toISOString(); + const archive: MockTaskArchive = { + id: nextMockTaskArchiveID++, + task_id: numericTaskID, + state: "archive_queued", + phase: "等待归档", + progress: 0, + format_version: 1, + original_size: 0, + compressed_size: 0, + task_name: task.name ?? "", + task_description: task.description, + task_goal: task.goal, + original_status: task.status, + category_id: task.category_id, + category_name: task.category_name, + source_task_ids: (task.source_task_ids ?? []).map(mockArchiveTaskID).filter((id) => id > 0), + remaining_timeout_seconds: 0, + data_counts: { + assets: assetIDs.length, + findings: findings.length, + llm_records: llmRecords.length, + }, + aggregate_stats: { + tokens: { + calls: llmRecords.length, + input_tokens: tokens.input_tokens, + output_tokens: tokens.output_tokens, + cache_read_tokens: tokens.cache_read_tokens, + cache_write_tokens: tokens.cache_write_tokens, + }, + findings: findings.reduce>((counts, finding) => { + counts[finding.severity] = (counts[finding.severity] ?? 0) + 1; + return counts; + }, {}), + }, + requested_at: now, + created_at: now, + updated_at: now, + snapshot: { task: structuredClone(task), numericTaskID, assetIDs, assetSources }, + }; + mockTaskArchives.unshift(archive); + setTimeout(() => { + if (archive.state !== "archive_queued") return; + archive.state = "archiving"; + archive.phase = "压缩任务数据"; + archive.progress = 55; + archive.updated_at = new Date().toISOString(); + }, 100); + setTimeout(() => { + if (archive.state !== "archiving" && archive.state !== "archive_queued") return; + archive.state = "ready"; + archive.phase = "归档完成"; + archive.progress = 100; + archive.archived_at = new Date().toISOString(); + archive.updated_at = archive.archived_at; + archive.original_size = Math.max(4096, JSON.stringify(archive.snapshot).length * 4); + archive.compressed_size = Math.max(1024, Math.round(archive.original_size * 0.32)); + const index = mockTasks.findIndex((item) => item.id === taskID); + if (index >= 0) mockTasks.splice(index, 1); + for (const asset of mockAssets) asset.task_ids = asset.task_ids.filter((id) => id !== numericTaskID); + deleteMockTaskAssetSources(taskID); + if (mockActiveTask === taskID) { + mockActiveTask = mockTasks[0]?.id ?? ""; + for (const item of mockTasks) item.active = item.id === mockActiveTask; + } + }, 500); + return archive; +} + +function mockRestoreArchive(archive: MockTaskArchive): void { + if (archive.state !== "ready" && archive.state !== "restore_failed") throw new Error("当前归档状态不可还原"); + archive.state = "restore_queued"; + archive.phase = "等待还原"; + archive.progress = 0; + archive.error = undefined; + archive.updated_at = new Date().toISOString(); + setTimeout(() => { + if (archive.state !== "restore_queued") return; + archive.state = "restoring"; + archive.phase = "恢复任务数据"; + archive.progress = 60; + archive.updated_at = new Date().toISOString(); + }, 100); + setTimeout(() => { + if (archive.state !== "restoring" && archive.state !== "restore_queued") return; + const restored = structuredClone(archive.snapshot.task); + restored.active = false; + restored.queued = false; + if (!mockTasks.some((task) => task.id === restored.id)) mockTasks.push(restored); + mockTaskAssetIDs.set(restored.id, archive.snapshot.numericTaskID); + for (const assetID of archive.snapshot.assetIDs) { + const asset = mockAssets.find((candidate) => candidate.id === assetID); + if (asset && !asset.task_ids.includes(archive.snapshot.numericTaskID)) { + asset.task_ids.push(archive.snapshot.numericTaskID); + } + } + for (const [assetID, source] of archive.snapshot.assetSources) setMockTaskAssetSource(restored.id, assetID, source); + const index = mockTaskArchives.indexOf(archive); + if (index >= 0) mockTaskArchives.splice(index, 1); + sortMockTasks(); + }, 500); +} + +function mockDeleteArchive(archive: MockTaskArchive): void { + if (archive.state !== "ready" && archive.state !== "delete_failed") throw new Error("当前归档状态不可永久删除"); + const dependent = mockTaskArchives.find( + (candidate) => candidate.id !== archive.id && candidate.source_task_ids.includes(archive.task_id), + ); + if (dependent) throw new Error(`归档仍被任务 #${dependent.task_id} 依赖，无法永久删除`); + archive.state = "delete_queued"; + archive.phase = "等待永久删除"; + archive.progress = 0; + archive.error = undefined; + archive.updated_at = new Date().toISOString(); + setTimeout(() => { + if (archive.state !== "delete_queued") return; + archive.state = "deleting"; + archive.phase = "删除归档包"; + archive.progress = 70; + }, 100); + setTimeout(() => { + const index = mockTaskArchives.indexOf(archive); + if (index >= 0) mockTaskArchives.splice(index, 1); + }, 450); +} + function setMockTaskAssetSource(taskID: string, assetID: number, source: MockTaskAssetSource) { mockTaskAssetSources.set(mockTaskAssetSourceKey(taskID, assetID), source); } @@ -98,6 +298,16 @@ function mockAssetMatchesDSL(asset: Asset, dsl: string): boolean { return query.split(/\s+/).every((term) => haystack.includes(term)); } +function mockFindingMatchesQuery(finding: (typeof mockFindings)[number], query: string | null): boolean { + const needle = query?.trim().toLowerCase(); + if (!needle) return true; + return [finding.name, finding.vulnclass, finding.summary, finding.evidence, finding.report].some((value) => + String(value ?? "") + .toLowerCase() + .includes(needle), + ); +} + function mockTaskCategorySnapshot(): TaskCategory[] { return mockTaskCategories.map((category) => ({ ...category, @@ -303,10 +513,111 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/auth/login" || path === "/auth/init") return { token: "mock-demo" }; if (path === "/auth/change-password") return { ok: true }; + // ── task cold archives ── + if (path === "/task-archives" && m === "GET") { + const page = Math.max(1, Number(q.get("page")) || 1); + const size = Math.min(100, Math.max(1, Number(q.get("size")) || 20)); + const query = (q.get("q") ?? "").trim().toLowerCase(); + const state = (q.get("state") ?? "").trim(); + const filtered = mockTaskArchives.filter((archive) => { + if (state && archive.state !== state) return false; + if (!query) return true; + return [archive.task_id, archive.task_name, archive.task_description].some((value) => + String(value ?? "") + .toLowerCase() + .includes(query), + ); + }); + const offset = (page - 1) * size; + return { + items: filtered.slice(offset, offset + size).map(publicMockTaskArchive), + total: filtered.length, + page, + size, + }; + } + if (seg[0] === "task-archives" && seg.length === 2 && m === "GET") { + const archive = mockTaskArchives.find((item) => item.id === Number(seg[1])); + if (!archive) throw new Error("归档不存在"); + return publicMockTaskArchive(archive); + } + if (seg[0] === "task-archives" && seg[2] === "restore" && seg.length === 3 && m === "POST") { + const archive = mockTaskArchives.find((item) => item.id === Number(seg[1])); + if (!archive) throw new Error("归档不存在"); + mockRestoreArchive(archive); + return publicMockTaskArchive(archive); + } + if (seg[0] === "task-archives" && seg.length === 2 && m === "DELETE") { + const archive = mockTaskArchives.find((item) => item.id === Number(seg[1])); + if (!archive) throw new Error("归档不存在"); + mockDeleteArchive(archive); + return publicMockTaskArchive(archive); + } + if (path === "/task-archives/restore/batch" && m === "POST") { + const ids = bodyIDs(b.archive_ids).map(Number); + const items = ids.map((id) => { + const archive = mockTaskArchives.find((item) => item.id === id); + if (!archive) return { id: String(id), archive_id: id, ok: false, queued: false, error: "归档不存在" }; + try { + mockRestoreArchive(archive); + return { id: String(id), archive_id: id, ok: true, queued: true }; + } catch (error) { + return { id: String(id), archive_id: id, ok: false, queued: false, error: (error as Error).message }; + } + }); + return { items }; + } + if (path === "/task-archives/delete/batch" && m === "POST") { + const ids = bodyIDs(b.archive_ids).map(Number); + const items = ids.map((id) => { + const archive = mockTaskArchives.find((item) => item.id === id); + if (!archive) return { id: String(id), archive_id: id, ok: false, queued: false, error: "归档不存在" }; + try { + mockDeleteArchive(archive); + return { id: String(id), archive_id: id, ok: true, queued: true }; + } catch (error) { + return { id: String(id), archive_id: id, ok: false, queued: false, error: (error as Error).message }; + } + }); + return { items }; + } + if (path === "/tasks/archive/batch" && m === "POST") { + const requested = bodyIDs(b.task_ids); + const selected = new Set(requested); + const ordered: string[] = []; + const visiting = new Set(); + const visited = new Set(); + const visit = (id: string) => { + if (visited.has(id) || visiting.has(id)) return; + visiting.add(id); + for (const candidate of mockTasks) { + if (!selected.has(candidate.id) || !(candidate.source_task_ids ?? []).map(String).includes(id)) continue; + visit(candidate.id); + } + visiting.delete(id); + visited.add(id); + ordered.push(id); + }; + for (const id of requested) visit(id); + const byID = new Map(); + for (const id of ordered) { + try { + const archive = mockArchiveTask(id); + byID.set(id, { id, archive_id: archive.id, ok: true, queued: true }); + } catch (error) { + byID.set(id, { id, ok: false, queued: false, error: (error as Error).message }); + } + } + return { items: requested.map((id) => byID.get(id) ?? { id, ok: false, queued: false, error: "任务不存在" }) }; + } + if (seg[0] === "tasks" && seg[2] === "archive" && seg.length === 3 && m === "POST") { + return publicMockTaskArchive(mockArchiveTask(seg[1])); + } + // ── tasks ── if (path === "/tasks" && m === "GET") { sortMockTasks(); - return { tasks: mockTasks, active: mockActiveTask }; + return { tasks: mockTasks.map(publicMockTask), active: mockActiveTask }; } if (path === "/tasks" && m === "POST") { let suffix = 1; @@ -493,7 +804,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (seg[0] === "tasks" && seg.length === 2 && m === "GET") { const task = mockTasks.find((item) => item.id === seg[1]); if (!task) throw new Error("任务不存在"); - return structuredClone(task); + return publicMockTask(task); } if (seg[0] === "tasks" && seg.length === 2 && m === "PATCH") { const task = mockTasks.find((item) => item.id === seg[1]); @@ -998,6 +1309,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re const filterStatus = q.get("status"); const filterVulnclass = q.get("vulnclass"); const filterTask = q.get("task_id"); + list = list.filter((finding) => mockFindingMatchesQuery(finding, q.get("q"))); if (filterSeverity) list = list.filter((finding) => finding.severity === filterSeverity); if (filterStatus) list = list.filter((finding) => finding.status === filterStatus); if (filterVulnclass) list = list.filter((finding) => finding.vulnclass === filterVulnclass); @@ -1108,6 +1420,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re const fStatus = q.get("status"); const fVuln = q.get("vulnclass"); const fTask = q.get("task_id"); + list = list.filter((finding) => mockFindingMatchesQuery(finding, q.get("q"))); if (fSev) list = list.filter((f) => f.severity === fSev); if (fStatus) list = list.filter((f) => f.status === fStatus); if (fVuln) list = list.filter((f) => f.vulnclass === fVuln); diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 8b0208d..633688e 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -35,6 +35,7 @@ export interface Task { llm_failover_state?: "default" | "ready" | "chain_exhausted" | string; llm_failover_reason?: string; source_task_ids?: string[]; // directly related tasks inherited as read-only context + archive_blocked_by_task_id?: string; // live direct dependent that must be archived first company_ids?: number[]; // associated company scopes; current company assets join the task at creation coverage_enabled?: boolean; // 资产覆盖度功能开关(创建时定,默认开)；false=不计算/不展示覆盖度 } @@ -75,6 +76,74 @@ export interface DeleteTaskResult { cleanup_warning?: string; } +export type TaskArchiveState = + | "archive_queued" + | "archiving" + | "archive_failed" + | "ready" + | "restore_queued" + | "restoring" + | "restore_failed" + | "delete_queued" + | "deleting" + | "delete_failed"; + +export interface TaskArchiveTokenStats { + calls?: number; + input_tokens?: number; + output_tokens?: number; + cache_read_tokens?: number; + cache_write_tokens?: number; +} + +export interface TaskArchive { + id: number; + task_id: number; + state: TaskArchiveState; + phase: string; + progress: number; + error?: string; + warnings?: string[]; + format_version: number; + sha256?: string; + original_size: number; + compressed_size: number; + task_name: string; + task_description: string; + task_goal: string; + original_status: TaskStatus; + category_id?: number; + category_name?: string; + source_task_ids: number[]; + remaining_timeout_seconds: number; + data_counts: Record; + aggregate_stats: { + tokens?: TaskArchiveTokenStats; + skills?: Record; + tools?: Record; + findings?: Record; + }; + archived_at?: string; + requested_at: string; + created_at: string; + updated_at: string; +} + +export interface TaskArchivePage { + items: TaskArchive[]; + total: number; + page: number; + size: number; +} + +export interface ArchiveBatchItem { + id: string; + archive_id?: number; + ok: boolean; + queued: boolean; + error?: string; +} + // ---- Asset graph (global, shared across tasks) ---- export type AssetType = | "company" @@ -477,6 +546,7 @@ export interface FindingQuery { status?: "all" | FindingStatus; vulnclass?: string; task?: string; // 任务 id;"all"/空 = 不按任务筛选 + query?: string; sort?: "severity" | "time"; } From 22132a4d21e1e43a654f608ac4160398914a6dae Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 14:03:07 +0800 Subject: [PATCH 4/7] chore: sync upstream v0.3.6 changes --- CHANGELOG.md | 43 +++++++ agent/blackboard_inheritance_test.go | 2 +- agent/capture_usage_test.go | 11 ++ agent/chat.go | 44 +++++--- agent/goals.go | 7 +- agent/mainagent.go | 32 ++++-- agent/planner.go | 38 ++++--- agent/provider.go | 41 ++++++- agent/tools.go | 57 +++++++++- agent/worker.go | 14 ++- db/commands.go | 69 ++++++++++-- db/config.go | 25 +++-- db/exploration.go | 34 ++++++ db/exploration_sources.go | 47 ++++++++ db/schema.sql | 4 + go.mod | 4 +- go.sum | 4 +- llmpool/pool.go | 38 +++++++ llmpool/pool_test.go | 11 ++ llmrec/llmrec.go | 47 +++++++- llmrec/llmrec_test.go | 36 ++++++ server/commands.go | 40 +++++-- server/dto.go | 4 + server/goals.go | 3 +- server/orchestration.go | 10 +- server/server.go | 50 ++++++++- server/server_mgmt.go | 9 +- server/task_llm.go | 106 +++++++++++++++++- server/task_llm_test.go | 21 ++++ web/src/app/(main)/function/commands/page.tsx | 93 ++++++++++++++- web/src/app/(main)/system/llm/page.tsx | 43 +++++-- web/src/lib/api.ts | 14 ++- web/src/lib/mock/handler.ts | 14 ++- web/src/lib/types.ts | 9 ++ 34 files changed, 897 insertions(+), 127 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a3fbdeb..f38f7b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,49 @@ ## [Unreleased] +## [0.3.6] - 2026-08-27 + +### LLM + +#### 修复的问题 + +- 修复「非流式」开关保存后重新打开又变回「流式」的问题（#69）：配置列表接口的 DTO 漏了 `streaming` 字段，响应从不返回它，前端读到 `undefined` 后一律回落成默认流式；实际写入与 DB 存储本无问题，只是读不回来。DTO 补回该字段（不加 `omitempty`，`false` 也必须出现在响应里）。 +- 测试连接改为校验模型确有回复、并按该配置真实的收发模式来测（#65）：此前只判 HTTP 是否成功，请求通但模型零回复（思考烧光预算 / 正文被安全策略吞掉 / 兼容层丢 `content`）也报「连接成功」，与会话里「不回话」的表现割裂；且测试恒走流式，非流式配置测的其实是另一条通道。现在空回复直接判失败、成功时在提示里展示模型回复，并把当前流式开关一并带入测试，让「测试通过」与「会话跑得通」保持一致。 + +### Agent + +#### 修改的功能 + +- `list_facts` 改为分页 + 关键词过滤，避免事实多时单次调用撑爆上下文（#74）：默认返回最新 20 条，支持 `limit`（上限 100）/ `before` 游标 / `q` 摘要关键词，返回 `{facts, total, has_more, next_before}`；单条 summary 过长按字数截断（全文仍走 `node_detail`）。worker / planner 提示词同步为分页语义。旧库通过一次性迁移把新参数 schema 刷进工具目录表（`SeedTool` 首插入only，否则工具管理页显示「无参数」）。 + +### 工具 + +#### 新增的功能 + +- 「工具执行」页新增工具调用次数统计（#72）：工具栏「统计」按钮打开弹窗，按工具展示调用次数、占比与失败次数（次数降序）；沿用列表的任务 / 关键词筛选，统计整个结果集而非当前页，弹窗打开时才拉取。 + +### 依赖 + +#### 修改的功能 + +- 升级 norma v0.2.9 → v0.3.2：修复 OpenAI 系响应里被静默丢弃的 reasoning/refusal（Responses 补 `reasoning_text`）、reasoning 字段三名去重、空响应（零内容块）有界重试、压缩边界切断工具配对导致的网关 400（同时修复已烤进 transcript 的存量孤儿）。 + +### 贡献者 + +- [@Autumn-27](https://github.com/Autumn-27) + +## [0.3.5] - 2026-08-25 + +### LLM + +#### 新增的功能 + +- 每个 LLM 配置支持流式/非流式切换（默认流式）：开走流式 SSE；关走真·非流式（`stream:false`、一次性返回完整 JSON），可绕开部分网关糟糕的 SSE 实现（空帧、思考字段丢帧），代价是失去运行中的实时进度与实时 Token 计数。worker / planner / mainagent / chat / goals 均按当前激活配置动态取值；`llmpool` / `llmrec` / 任务运行时三层 Provider 包装均兼容非流式；`llm_profiles` 新增 `streaming` 列并 `ALTER` 补旧库（默认 `true`，旧配置无感）。 + +### 贡献者 + +- [@Autumn-27](https://github.com/Autumn-27) + ## [0.3.4] - 2026-08-24 ### UI diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go index ad4fe4b..4549ff5 100644 --- a/agent/blackboard_inheritance_test.go +++ b/agent/blackboard_inheritance_test.go @@ -326,7 +326,7 @@ func TestBlackboardToolsReadDirectSources(t *testing.T) { t.Fatalf("inherited host context missing: %#v", coverage) } - facts := callReadJSON(t, tools.listFacts(), `{}`).([]any) + facts := callReadJSON(t, tools.listFacts(), `{}`).(map[string]any)["facts"].([]any) seenSource, seenGrand := false, false for _, raw := range facts { item := raw.(map[string]any) diff --git a/agent/capture_usage_test.go b/agent/capture_usage_test.go index c76f79c..5b6389c 100644 --- a/agent/capture_usage_test.go +++ b/agent/capture_usage_test.go @@ -19,6 +19,17 @@ func (p captureUsageProvider) Stream(ctx context.Context, _ llm.CompletionReques return func(yield func(llm.StreamEvent, error) bool) { p.stream(ctx, yield) } } +func (p captureUsageProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + acc := llm.NewAccumulator() + for ev, err := range p.Stream(ctx, req) { + if err != nil { + return llm.Message{}, "", llm.Usage{}, err + } + acc.Add(ev) + } + return acc.Message(), acc.StopReason, acc.Usage, nil +} + func TestCaptureRunPersistsUsageOnProviderFailure(t *testing.T) { wantErr := errors.New("provider failed after reporting usage") provider := captureUsageProvider{stream: func(_ context.Context, yield func(llm.StreamEvent, error) bool) { diff --git a/agent/chat.go b/agent/chat.go index 41be93b..0f38582 100644 --- a/agent/chat.go +++ b/agent/chat.go @@ -6,13 +6,13 @@ import ( "path/filepath" "time" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/guard" "github.com/Autumn-27/norma/agentcore" "github.com/Autumn-27/norma/llm" "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" "github.com/Autumn-27/norma/transcript" - "github.com/Autumn-27/artex/db" - "github.com/Autumn-27/artex/guard" ) // ChatAgent is the generic, task-independent conversational runner behind the chat @@ -22,21 +22,28 @@ import ( // LS/Glob/Grep) plus whatever skills/MCP the key is made visible; NO pentest // graph/task context is injected (that stays exclusive to MainAgent). type ChatAgent struct { - prov llm.Provider - model string - workDir string - tx *transcript.Store - window int - proxyAddr string - proxyCACert string - webSearch WebSearchOpts - guard *guard.Guard // optional; nil disables intercept hooks for chat + prov llm.Provider + model string + workDir string + tx *transcript.Store + window int + proxyAddr string + proxyCACert string + webSearch WebSearchOpts + guard *guard.Guard // optional; nil disables intercept hooks for chat + nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming) } func NewChatAgent(prov llm.Provider, model, workDir string, tx *transcript.Store, window int) *ChatAgent { return &ChatAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window} } +// SetNonStreaming wires a resolver deciding whether chat runs use the +// non-streaming model path (true = non-streaming). nil/unset = streaming. +func (c *ChatAgent) SetNonStreaming(fn func() bool) { c.nonStreamingFn = fn } + +func (c *ChatAgent) nonStreaming() bool { return c.nonStreamingFn != nil && c.nonStreamingFn() } + // SetProxy points the chat agent's WebFetch/Bash at the recording proxy plus the // CA cert it trusts (empty addr = direct). Kept for parity with the other agents. func (c *ChatAgent) SetProxy(addr, caCert string) { c.proxyAddr, c.proxyCACert = addr, caCert } @@ -115,18 +122,19 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin BraveSearchAPIKey: ws.BraveKey, TavilySearchAPIKey: ws.TavilyKey, WebSearchProxy: ws.Proxy, - BashEnv: proxyEnv(c.proxyAddr, c.proxyCACert), // Bash 子命令默认走代理+信任 CA - WorkingDir: sessionWorkDir, - MaxTurns: maxTurns, - MaxDuration: maxDuration, - Compaction: compactionConfig(c.window), - Todos: actool.NewTodoStore(), + BashEnv: proxyEnv(c.proxyAddr, c.proxyCACert), // Bash 子命令默认走代理+信任 CA + WorkingDir: sessionWorkDir, + MaxTurns: maxTurns, + MaxDuration: maxDuration, + Compaction: compactionConfig(c.window), + Todos: actool.NewTodoStore(), // large tool output spills to cmd-output/ under the session dir. // 截断上限用 SDK 默认(tool.Capture 的 30000 字符)。 ToolOutputDir: filepath.Join(sessionWorkDir, "cmd-output"), // 命中预算(步数)→ SDK 跑收尾:输出一句总结。Prompt 与收尾轮数按本 agent key 后台可编辑 // (自定义 agent 各自一份;留空/0 用通用默认:10 轮)。 - Settlement: wrapupSettlement(agentKey, nil), + Settlement: wrapupSettlement(agentKey, nil), + NonStreaming: c.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete } if c.guard != nil { opts.Hooks = c.guard.Hooks() diff --git a/agent/goals.go b/agent/goals.go index 388515a..b98e9fe 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -98,13 +98,13 @@ func DecomposeGoals(ctx context.Context, prov llm.Provider, dataDir, goalText, d if prov == nil { return nil } - return DecomposeGoalsWithProvider(ctx, prov, dataDir, goalText, desc, as, ts, taskID, emit) + return DecomposeGoalsWithProvider(ctx, prov, dataDir, goalText, desc, as, ts, taskID, false, 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 { +func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, nonStreaming bool, emit func(db.Activity)) []GoalSpec { if prov == nil { return nil } @@ -145,7 +145,8 @@ func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, PermissionMode: acperm.ModeBypass, DisableBackgroundTasks: true, // 3 步(抽约束 → 登记范围 → 拆目标)各需一次工具调用,给足回合避免收尾前漏调 set_goals。 - MaxTurns: 8, + MaxTurns: 8, + NonStreaming: nonStreaming, // 该 profile 选非流式时走 Provider.Complete }, userMsg, captureEmit) // set_goals persisted the goals directly; read them back so the caller sees what // was written (empty slice ⇒ the LLM produced nothing ⇒ caller falls back). diff --git a/agent/mainagent.go b/agent/mainagent.go index 99531df..9c40eca 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -17,19 +17,26 @@ import ( // (→planner) or direct high-priority intents (→frontier). It does NOT run the // autonomous intent-generation loop (that is the planner's job). type MainAgent struct { - prov llm.Provider - model string - tx *transcript.Store // raw LLM conversation persistence (nil = off) - window int // context window in tokens (for compaction) - windowFn func() int // optional dynamic task-chain minimum - maxTurns int // max agent turns per run (0 = unlimited) - proxyAddr string // recording proxy for WebFetch (empty = direct) - proxyCACert string // recording proxy's CA cert path (HTTPS verify) - webSearch WebSearchOpts // web_search tool backend selection (off by default) - workDir string // shared work dir (surfaced in prompt as artifact-output target) - steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off) + prov llm.Provider + model string + tx *transcript.Store // raw LLM conversation persistence (nil = off) + window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum + maxTurns int // max agent turns per run (0 = unlimited) + proxyAddr string // recording proxy for WebFetch (empty = direct) + proxyCACert string // recording proxy's CA cert path (HTTPS verify) + webSearch WebSearchOpts // web_search tool backend selection (off by default) + workDir string // shared work dir (surfaced in prompt as artifact-output target) + steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off) + nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming) } +// SetNonStreaming wires a resolver deciding whether runs use the non-streaming +// model path (true = non-streaming). nil/unset = streaming (default). +func (m *MainAgent) SetNonStreaming(fn func() bool) { m.nonStreamingFn = fn } + +func (m *MainAgent) nonStreaming() bool { return m.nonStreamingFn != nil && m.nonStreamingFn() } + func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store, window, maxTurns int) *MainAgent { return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns} } @@ -128,7 +135,8 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 // 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。 - Settlement: wrapupSettlement("mainagent", nil), + Settlement: wrapupSettlement("mainagent", nil), + NonStreaming: m.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete } if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task opts.Transcript = m.tx diff --git a/agent/planner.go b/agent/planner.go index dfacc91..2647314 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -20,19 +20,20 @@ import ( // assets, judges whether the task goal is met, and emits 0..N exploration intents // into the frontier. It is the sole intent generator. type Planner struct { - prov llm.Provider - model string - tx *transcript.Store // raw LLM conversation persistence (nil = off) - window int // context window in tokens (for compaction) - windowFn func() int // optional dynamic task-chain minimum - maxTurns int // max agent turns per run (0 = unlimited) - 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) - proxyAddr string // recording proxy for WebFetch (empty = direct) - proxyCACert string // recording proxy's CA cert path (HTTPS verify) - webSearch WebSearchOpts // web_search tool backend selection (off by default) - workDir string // shared work dir (surfaced in prompt as artifact-output target) - injectConstraints func() bool // resolver: inject task operation constraints into system prompt? (nil = yes) + prov llm.Provider + model string + tx *transcript.Store // raw LLM conversation persistence (nil = off) + window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum + maxTurns int // max agent turns per run (0 = unlimited) + 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) + proxyAddr string // recording proxy for WebFetch (empty = direct) + proxyCACert string // recording proxy's CA cert path (HTTPS verify) + webSearch WebSearchOpts // web_search tool backend selection (off by default) + workDir string // shared work dir (surfaced in prompt as artifact-output target) + injectConstraints func() bool // resolver: inject task operation constraints into system prompt? (nil = yes) + nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming) // todos keeps ONE plan-scratchpad per task (keyed by exploration id) so the // planner's multi-step plan survives across wake-ups — each Plan() is a fresh @@ -48,6 +49,12 @@ func NewPlanner(prov llm.Provider, model, workDir string, tx *transcript.Store, func (p *Planner) SetCompactionWindowResolver(fn func() int) { p.windowFn = fn } +// SetNonStreaming wires a resolver deciding whether runs use the non-streaming +// model path (true = non-streaming). nil/unset = streaming (default). +func (p *Planner) SetNonStreaming(fn func() bool) { p.nonStreamingFn = fn } + +func (p *Planner) nonStreaming() bool { return p.nonStreamingFn != nil && p.nonStreamingFn() } + func (p *Planner) compactionWindow() int { if p.windowFn != nil { return p.windowFn() @@ -238,7 +245,7 @@ const plannerDefaultTmpl = `你是一个授权渗透测试系统的"规划者" ⚠️ 最重要的原则：**这一轮生成 0 个意图是完全正常、而且是最常见的结果。** 你不是"每次都要产出意图"的机器。绝大多数唤醒，frontier 里已有的意图已经覆盖了所有已知方向，你应当什么都不加、直接结束。重复/换措辞地生成已经存在的意图是严重错误。 决策流程（每次唤醒）： -1. **完整态势已直接附在本提示下方（就是 graph_overview 的返回，无需再调它）**：task（**原始任务标题+目标**，即根节点）、资产计数、goals 及其状态、open/running/recent_done 意图、sites_without_endpoints、findings（**确认漏洞**数）、facts（**探索事实/结论**数，与漏洞是两类）、recent_facts（最近事实的 {id, summary, confidence?}，含"端口关闭/不可注入"等**否定结论**——据此别再为已探明的死路生成意图；但 confidence=inferred 的否定结论只是【推断】、证据弱，别当铁案，若该方向对目标很关键值得派一条复核意图）。**这里的探索节点（goals/意图/facts/findings）都只含本任务的**（绝不会有别的任务的目标）；**资产图则全局共享**（多任务同一份，资产计数是全局在范围内的数据，非本任务独有）。需要更深的细节时才按需调：list_facts 列全部事实、list_findings 列全部漏洞、**node_detail(id)** 取某条的完整证据/详情（列表/recent_facts 只给摘要）。 +1. **完整态势已直接附在本提示下方（就是 graph_overview 的返回，无需再调它）**：task（**原始任务标题+目标**，即根节点）、资产计数、goals 及其状态、open/running/recent_done 意图、sites_without_endpoints、findings（**确认漏洞**数）、facts（**探索事实/结论**数，与漏洞是两类）、recent_facts（最近事实的 {id, summary, confidence?}，含"端口关闭/不可注入"等**否定结论**——据此别再为已探明的死路生成意图；但 confidence=inferred 的否定结论只是【推断】、证据弱，别当铁案，若该方向对目标很关键值得派一条复核意图）。**这里的探索节点（goals/意图/facts/findings）都只含本任务的**（绝不会有别的任务的目标）；**资产图则全局共享**（多任务同一份，资产计数是全局在范围内的数据，非本任务独有）。需要更深的细节时才按需调：list_facts 分页列事实（最新在前，默认 20 条，可传 q 关键词过滤、before 翻页，返回带 total/has_more）、list_findings 列全部漏洞、**node_detail(id)** 取某条的完整证据/详情（列表/recent_facts 只给摘要）。 - open_intents / running_intents / recent_done_intents——"哪些方向已经有意图在覆盖/已尝试"。每个意图还带 **parents（上游：它派生自哪些事实/意图）和 yields（下游：它产生了哪些事实/发现）**——这就是探索图的**血缘关系**，据此理解"哪些事实来自哪个方向、能否综合成新方向"。recent_facts 里每个事实带 **from_intent**（由哪个意图产生）。 - sites_without_endpoints / findings——"哪些方向【可能】需要探索"。 - 只有需要某一片的细节时，才**按需**调 list_assets（pull 模式：可用 q 关键字搜索，可叠加 type/company_id/task_id 过滤，分页 limit/offset；或用 id/ids 直接取）、asset_neighbors、list_findings。资产图全局共享，别默认拉全量。资产中可能包含非本次任务涉及到的资产，所以需要主要出现非本次任务相关的资产时忽略这些资产。 @@ -353,7 +360,8 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts // 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、 // 能证的 prove_goal、串行链记 TodoWrite),而非停止规划——planner 之后仍会被反复唤醒。 // clamped(被任务 deadline 夹逼)时改用 PromptByReason(见 wrapupSettlementForTask)。 - Settlement: settle, + Settlement: settle, + NonStreaming: p.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete } if p.tx != nil { // persist raw LLM conversation; one accumulating file per task's planner opts.Transcript = p.tx diff --git a/agent/provider.go b/agent/provider.go index 6011ada..c66a1b8 100644 --- a/agent/provider.go +++ b/agent/provider.go @@ -51,6 +51,11 @@ type Config struct { // "" = 不发送(默认); "low"/"medium"/"high"/"xhigh"/"max" = 对应强度. // OpenAI 映射为顶层 reasoning_effort;Anthropic 映射为 output_config.effort. ReasoningEffort string + // Stream 控制该 profile 是否使用流式(SSE)接口。true(默认)= 流式;false = 真· + // 非流式(发 stream:false,一次性拿完整 JSON,走 Provider.Complete)。非流式可绕开 + // 某些网关糟糕的 SSE 实现(空帧、思考字段丢帧),代价是失去运行中的实时进度/实时 + // token 计数。映射为 agentcore.Options.NonStreaming = !Stream。 + Stream bool } // compaction window resolution bounds (in K tokens). Below the floor the @@ -116,6 +121,8 @@ func FromEnv() (Config, bool) { BaseURL: os.Getenv("ARTEX_LLM_BASE_URL"), Model: os.Getenv("ARTEX_LLM_MODEL"), Proxy: strings.TrimSpace(os.Getenv("ARTEX_LLM_PROXY")), + // 默认流式;ARTEX_LLM_STREAM=false/0/off 显式关闭走非流式。 + Stream: !isFalsy(os.Getenv("ARTEX_LLM_STREAM")), } switch prov { case "openai": @@ -147,6 +154,7 @@ func ConfigFrom(provider, model, baseURL, apiKey, proxy string) Config { BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"), APIKey: strings.TrimSpace(apiKey), Proxy: strings.TrimSpace(proxy), + Stream: true, // 默认流式;调用方按 profile 覆盖 } switch strings.TrimSpace(provider) { case "openai": @@ -167,6 +175,16 @@ func ConfigFrom(provider, model, baseURL, apiKey, proxy string) Config { return c } +// isFalsy reports whether an env-var string explicitly requests "off". Empty or +// unrecognized → false (so an unset var keeps the streaming default). +func isFalsy(s string) bool { + switch strings.ToLower(strings.TrimSpace(s)) { + case "0", "false", "off", "no": + return true + } + return false +} + // Provider returns the short provider name ("anthropic"/"openai"). func (c Config) Provider() string { if c.Format == llm.FormatOpenAI { @@ -288,11 +306,12 @@ func quotaAwareHTTPClient(proxy string) (*http.Client, error) { } // 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) { +// endpoint/key actually work. Returns the round-trip latency and the model's +// reply text. +func TestConnection(ctx context.Context, c Config) (time.Duration, string, error) { prov, err := c.NewProvider() if err != nil { - return 0, err + return 0, "", err } ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() @@ -302,12 +321,24 @@ func TestConnection(ctx context.Context, c Config) (time.Duration, error) { // 就撞到输出上限(finish=length)、被截断,连接测试虽仍算通(err=nil)但显示成 // "已中断/length/resume" 一团糟。给足预算让它把 OK 干净吐完(finish=stop)。 // EscalateMaxTokens 保持 false:不因截断而抬额重试,避免 resume 循环空烧。 - _, err = agentcore.Run(ctx, agentcore.Options{ + reply, err := agentcore.Run(ctx, agentcore.Options{ Provider: prov, SystemPrompt: []string{"你是连接测试。直接输出两个字符 OK 即可，不要思考、不要解释、不要别的。"}, PermissionMode: acperm.ModeBypass, MaxTurns: 1, MaxTokens: 8192, + NonStreaming: !c.Stream, // 用该 profile 的真实收发模式做连接测试 }, "ping") - return time.Since(start), err + lat := time.Since(start) + if err != nil { + return lat, "", err + } + // err==nil 还不够：请求通了但模型一个字都不吐的情况真实存在(思考把预算烧光、 + // 正文被安全策略吞掉、兼容层把 content 丢了)。这种配置在会话里就是"不回话", + // 测试却报成功——正是本项要消除的落差。没有可见正文一律判失败。 + reply = strings.TrimSpace(reply) + if reply == "" { + return lat, "", fmt.Errorf("模型无回复内容（请求已通，但未返回任何文本）") + } + return lat, reply, nil } diff --git a/agent/tools.go b/agent/tools.go index f8493ab..ebed6d8 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -898,19 +898,64 @@ func (t *ToolSet) listFindings() actool.CoreTool { }) } +// factsPageSize is the default page size for list_facts. Facts pile up on long +// tasks; returning all of them at once (the old behaviour) could blow up the +// context, so default to the newest page and let the agent page/filter for more. +const factsPageSize = 20 + func (t *ToolSet) listFacts() actool.CoreTool { - 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.ListByKindWithSources(db.KindFact, 500) + return readTool("list_facts", "分页列本任务及直接关联任务的【探索事实/结论】，最新在前(紧凑：id+摘要+状态，摘要过长会截断，全文用 node_detail(id))。参数均可选：limit(默认 20，上限 100)、before(游标，传上一页返回的 next_before 取更旧的一页；省略/0=最新一页)、q(按摘要关键词过滤)。返回 {facts, total, has_more, next_before}：total 是过滤后的总数，has_more=true 时用 next_before 继续翻页。关联任务条目带 source_task_id/inherited=true 且只读。漏洞看 list_findings。", + obj(map[string]any{ + "limit": intp("返回条数，默认 20，上限 100"), + "before": intp("分页游标：只返回 id 小于该值的更旧事实；省略或 0 = 最新一页"), + "q": str("按事实摘要关键词过滤（不区分大小写）；省略 = 不过滤"), + }), + func(_ context.Context, in json.RawMessage) (actool.Result, error) { + var a struct { + Limit int `json:"limit"` + Before int64 `json:"before"` + Q string `json:"q"` + } + _ = json.Unmarshal(in, &a) + limit := a.Limit + if limit <= 0 { + limit = factsPageSize + } + if limit > 100 { + limit = 100 + } + f, hasMore, total, err := t.ts.ListByKindPageWithSources(db.KindFact, a.Before, limit, strings.TrimSpace(a.Q)) + if err != nil { + return actool.Result{}, err + } out := make([]map[string]any, 0, len(f)) for _, n := range f { - out = append(out, compactNode(n)) + out = append(out, compactFact(n)) } - return jsonResult(out) + res := map[string]any{"facts": out, "total": total, "has_more": hasMore} + if hasMore && len(f) > 0 { + res["next_before"] = f[len(f)-1].ID // 传回它取下一页(更旧的) + } + return jsonResult(res) }) } +// factSummaryMax caps a fact summary in list_facts output. Facts carry one-line +// conclusions, but nothing enforces brevity; a runaway summary must not bloat a +// whole page. Full text stays available via node_detail(id). +const factSummaryMax = 160 + +// compactFact is compactNode with the summary rune-capped for list_facts, so a +// page of facts stays bounded regardless of how long any single summary grew. +func compactFact(n *db.Node) map[string]any { + m := compactNode(n) + if s, ok := m["summary"].(string); ok && len([]rune(s)) > factSummaryMax { + m["summary"] = string([]rune(s)[:factSummaryMax]) + "…" + m["summary_truncated"] = true + } + return m +} + func (t *ToolSet) nodeDetail() actool.CoreTool { 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"), diff --git a/agent/worker.go b/agent/worker.go index f2e5ce8..b7a9e7a 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -62,8 +62,19 @@ type Worker struct { // injected into the worker system prompt. Read per run so the settings toggle // takes effect without rebuilding the agent. nil = inject (default). injectConstraints func() bool + // nonStreamingFn resolves whether this run uses the non-streaming (Complete) + // path. Read per run so a profile/task toggle takes effect without rebuilding + // the agent. nil = streaming (default). + nonStreamingFn func() bool } +// SetNonStreaming wires a resolver deciding whether runs use the non-streaming +// model path (true = non-streaming). nil/unset = streaming (default). Read per +// run so a profile or task-chain toggle takes effect without rebuilding. +func (w *Worker) SetNonStreaming(fn func() bool) { w.nonStreamingFn = fn } + +func (w *Worker) nonStreaming() bool { return w.nonStreamingFn != nil && w.nonStreamingFn() } + // SetConstraintInject wires a resolver deciding whether this task's operation // constraints get injected into the worker system prompt. nil = inject (default). func (w *Worker) SetConstraintInject(fn func() bool) { w.injectConstraints = fn } @@ -160,7 +171,7 @@ const workerDefaultTmpl = `你是一个授权渗透测试系统的"执行者"(wo - insert_assets：登记新资产（资产图）。**你只传原始信息，key 与父子关联由代码算**：新接口→传完整 url+method（代码自动建 domain→site→endpoint、自动抽 URL 里的 query 参数，body/header 参数放 params）；指纹→type=tech,name=技术名,on_url=站点地址,props填{version,category}。多个资产放 assets 数组一次批量登记。属性写在资产自己的 props 上，探索结论不要写这里。 - record_fact：把探索【事实/结论】写入探索图并连到意图（传 intent_id）。正向/否定结论、观察、判断用它；**一次探索的多个观察汇总成一条事实**（summary 总结一句话 + detail 放细节），不要一个属性一条。真有多条不同结论才用 facts 数组。**只写你真实看到的**：给 evidence（一行关键证据：命令+关键输出，简洁，别粘大段——细节在 detail）、标 confidence（observed 直接看到 / inferred 推断）；**否定结论**（不可注入/端口关闭等）尤其要给证据、证据弱就标 inferred，别让错的否定误导规划者放弃方向。 - report_finding：确认漏洞 → 记录(含 PoC，传 intent_id=你领到的意图id)。**只有你在本次运行里真实触发过该漏洞、拿到了可复现的证据（请求/响应或命令输出）才用它。** 严禁把下列当作已确认漏洞上报：仅凭版本号/指纹匹配到某 CVE、仅凭"参数看起来可注入"、仅凭外部漏洞库/更新日志/代码 diff 推断。**不要用查 CVE 库或"对比补丁版本"替代实际触发。** 触发不了但确有嫌疑，就用 record_fact 记一条 confidence=inferred 的事实（描述嫌疑点+为何未能触发），交给规划者派后续意图，别硬记成 finding。 -- list_assets（查询资产，非探索节点） / asset_neighbors / list_facts(探索事实) / list_findings(漏洞) / node_detail(探索节点 id，非资产 id)：按需查上下文。 +- list_assets（查询资产，非探索节点） / asset_neighbors / list_facts(探索事实，分页最新在前；默认 20 条，可传 q 关键词过滤、before 翻页) / list_findings(漏洞) / node_detail(探索节点 id，非资产 id)：按需查上下文。 只在授权范围内操作；若系统提示顶部附有【操作约束】，那是最高优先级红线——任何命令/探测在执行前先自检是否违反，违反即不做（哪怕它落在你领到的意图里）。完成本意图后用一句话总结你做了什么、写回了哪些事实。务实、克制、聚焦这一条意图。` @@ -337,6 +348,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. ToolOutputDir: cmdOutDir(runDir), Compaction: compactionConfig(w.compactionWindow()), // long tool-heavy runs stay within the window Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 + NonStreaming: w.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete } if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic) opts.Hooks = hooks diff --git a/db/commands.go b/db/commands.go index 27172da..22fe14e 100644 --- a/db/commands.go +++ b/db/commands.go @@ -19,18 +19,10 @@ type CommandRecord struct { CreatedAt time.Time `json:"created_at"` } -// ListCommands returns tool executions (tool_use + paired tool_result) across all -// explorations, with optional filtering and pagination. Covers every tool, not -// just Bash; q matches the tool name or its input. -func (d *DB) ListCommands(expID *int64, q string, page, size int) ([]CommandRecord, int, error) { - if size <= 0 { - size = 50 - } - if page < 0 { - page = 0 - } - offset := page * size - +// commandFilter builds the WHERE clause shared by the tool-execution list and +// its per-tool tally, so the summary always describes exactly the rows the table +// pages through. Returns the clause, its args, and the next placeholder index. +func commandFilter(expID *int64, q string) (string, []any, int) { where := `WHERE u.kind = 'tool_use'` args := []any{} argN := 1 @@ -45,6 +37,59 @@ func (d *DB) ListCommands(expID *int64, q string, page, size int) ([]CommandReco args = append(args, "%"+q+"%") argN++ } + return where, args, argN +} + +// ToolStat is one tool's execution tally for the usage summary. +type ToolStat struct { + Tool string `json:"tool"` + Total int `json:"total"` + Errors int `json:"errors"` +} + +// ToolStats counts executions grouped by tool under the same filters +// ListCommands takes. Unpaginated on purpose: the tally describes the whole +// filtered set, not the page currently on screen. +func (d *DB) ToolStats(expID *int64, q string) ([]ToolStat, error) { + where, args, _ := commandFilter(expID, q) + + rows, err := d.Query(` +SELECT COALESCE(NULLIF(u.tool,''),'-') AS tool, COUNT(*) AS total, + COUNT(*) FILTER (WHERE COALESCE(r.is_error,false)) AS errors +FROM activity u +LEFT JOIN activity r ON r.tool_use_id = u.tool_use_id AND r.kind = 'tool_result' +`+where+` +GROUP BY 1 +ORDER BY total DESC, tool ASC`, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + out := []ToolStat{} + for rows.Next() { + var s ToolStat + if err := rows.Scan(&s.Tool, &s.Total, &s.Errors); err != nil { + return nil, err + } + out = append(out, s) + } + return out, rows.Err() +} + +// ListCommands returns tool executions (tool_use + paired tool_result) across all +// explorations, with optional filtering and pagination. Covers every tool, not +// just Bash; q matches the tool name or its input. +func (d *DB) ListCommands(expID *int64, q string, page, size int) ([]CommandRecord, int, error) { + if size <= 0 { + size = 50 + } + if page < 0 { + page = 0 + } + offset := page * size + + where, args, argN := commandFilter(expID, q) // count var total int diff --git a/db/config.go b/db/config.go index 40684d6..410994a 100644 --- a/db/config.go +++ b/db/config.go @@ -39,18 +39,23 @@ type LLMProfile struct { // usable when an agent/task binds it explicitly, it just never gets picked up // as a fallback target. PoolExclude bool `json:"pool_exclude"` + // Streaming selects the wire protocol: true (default) = streaming (SSE); + // false = real non-streaming (stream:false, single JSON response via + // Provider.Complete). Non-streaming sidesteps flaky gateway SSE at the cost of + // live in-run progress. Maps to agent.Config.Stream. + Streaming bool `json:"streaming"` } // 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,COALESCE(thinking_type,'')` -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,COALESCE(thinking_type,'')` +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,COALESCE(thinking_type,''),COALESCE(streaming,true)` +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,COALESCE(thinking_type,''),COALESCE(streaming,true)` // 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, &p.ThinkingType) + &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault, &p.Priority, &p.PoolExclude, &p.ThinkingType, &p.Streaming) } func (d *DB) ListProfiles() ([]*LLMProfile, error) { @@ -123,18 +128,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,priority,pool_exclude,thinking_type) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14) 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, p.ThinkingType).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,thinking_type,streaming) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15) 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, p.ThinkingType, p.Streaming).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,priority=$10,pool_exclude=$11,thinking_type=$12 WHERE id=$13`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, 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,thinking_type=$12,streaming=$13 WHERE id=$14`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, 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,priority=$12,pool_exclude=$13,thinking_type=$14 WHERE id=$15`, - 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.ThinkingType, 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,thinking_type=$14,streaming=$15 WHERE id=$16`, + 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.ThinkingType, p.Streaming, p.ID) return p.ID, err } diff --git a/db/exploration.go b/db/exploration.go index f62e280..eb0e677 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -672,6 +672,40 @@ ORDER BY id DESC LIMIT $4`, s.expID, kind, before, limit+1) return nodes, hasMore, nil } +// listByKindPageFiltered is ListByKindPage with an optional summary keyword +// filter (payload->>'summary' ILIKE %q%). Fetches one extra row so the caller can +// probe hasMore. Empty q = no filter. Newest-first by id. +func (s *ExplorationStore) listByKindPageFiltered(kind string, before int64, limit int, q string) ([]*Node, error) { + where := `exploration_id=$1 AND kind=$2 AND ($3 <= 0 OR id < $3)` + args := []any{s.expID, kind, before} + if q != "" { + args = append(args, "%"+q+"%") + where += fmt.Sprintf(` AND payload->>'summary' ILIKE $%d`, len(args)) + } + args = append(args, limit+1) + rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes +WHERE `+where+` ORDER BY id DESC LIMIT $`+fmt.Sprintf("%d", len(args)), args...) + if err != nil { + return nil, err + } + defer rows.Close() + return scanNodes(rows) +} + +// countByKindFiltered counts nodes of a kind under the same summary filter, for a +// page's total. Empty q = count all of that kind. +func (s *ExplorationStore) countByKindFiltered(kind, q string) (int, error) { + where := `exploration_id=$1 AND kind=$2` + args := []any{s.expID, kind} + if q != "" { + args = append(args, "%"+q+"%") + where += fmt.Sprintf(` AND payload->>'summary' ILIKE $%d`, len(args)) + } + var n int + err := s.db.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE `+where, args...).Scan(&n) + return n, err +} + // GetNode returns one node of this exploration by id (nil, nil if not found). func (s *ExplorationStore) GetNode(id int64) (*Node, error) { n, err := scanNode(s.db.QueryRow(`SELECT `+nodeCols+` FROM exploration_nodes WHERE id=$1 AND exploration_id=$2`, id, s.expID)) diff --git a/db/exploration_sources.go b/db/exploration_sources.go index 3ccbb96..c12d5ba 100644 --- a/db/exploration_sources.go +++ b/db/exploration_sources.go @@ -108,6 +108,53 @@ func (s *ExplorationStore) ListByKindWithSources(kind string, limit int) ([]*Nod return out, nil } +// ListByKindPageWithSources is the paginated, keyword-filterable sibling of +// ListByKindWithSources: it returns one newest-first page (id < before, before<=0 +// = newest) spanning this exploration and its direct sources, plus hasMore and the +// filtered total across all of them. Node ids are globally unique, so merging each +// store's own page and re-sorting by id DESC yields the true global page; fetching +// limit+1 per store guarantees the merged top-`limit` is complete. +func (s *ExplorationStore) ListByKindPageWithSources(kind string, before int64, limit int, q string) (nodes []*Node, hasMore bool, total int, err error) { + if limit <= 0 { + limit = 20 + } + own, err := s.listByKindPageFiltered(kind, before, limit, q) + if err != nil { + return nil, false, 0, err + } + merged := own + total, err = s.countByKindFiltered(kind, q) + if err != nil { + return nil, false, 0, err + } + + sources, err := s.DirectSourceStores() + if err != nil { + return nil, false, 0, err + } + for _, source := range sources { + page, err := source.Store.listByKindPageFiltered(kind, before, limit, q) + if err != nil { + return nil, false, 0, err + } + for _, n := range page { + merged = append(merged, markInheritedNode(n, source.Task.TaskID)) + } + cnt, err := source.Store.countByKindFiltered(kind, q) + if err != nil { + return nil, false, 0, err + } + total += cnt + } + + sort.Slice(merged, func(i, j int) bool { return merged[i].ID > merged[j].ID }) + hasMore = len(merged) > limit + if hasMore { + merged = merged[:limit] + } + return merged, hasMore, total, 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. diff --git a/db/schema.sql b/db/schema.sql index de0b7cd..9be7064 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -310,6 +310,8 @@ CREATE TABLE IF NOT EXISTS llm_profiles ( -- pool_exclude true=不作为故障转移目标(仍可被 agent/任务显式绑定使用)。 priority INTEGER NOT NULL DEFAULT 0, pool_exclude BOOLEAN NOT NULL DEFAULT false, + -- streaming=true(默认)走流式 SSE；false 走真·非流式(stream:false，一次性 JSON)。 + streaming BOOLEAN NOT NULL DEFAULT true, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); @@ -320,6 +322,8 @@ CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles -- 轮询顺位/排除标记；补旧库。默认 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; +-- 流式开关；补旧库。默认 true = 保持既有的流式行为，旧配置无感升级。 +ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS streaming BOOLEAN NOT NULL DEFAULT true; -- 思考开关字段 thinking_type，从旧的单一 reasoning_effort 语义一次性拆分而来。 -- schema.sql 每次启动都执行，故迁移必须只跑一次：仅当该列尚不存在时才回填， diff --git a/go.mod b/go.mod index d56340a..b047c4b 100644 --- a/go.mod +++ b/go.mod @@ -3,9 +3,10 @@ module github.com/Autumn-27/artex go 1.26.3 require ( - github.com/Autumn-27/norma v0.2.9 + github.com/Autumn-27/norma v0.3.2 github.com/golang-jwt/jwt/v5 v5.3.1 github.com/jackc/pgx/v5 v5.10.0 + github.com/klauspost/compress v1.18.6 github.com/lqqyt2423/go-mitmproxy v1.9.2 github.com/miekg/dns v1.1.72 github.com/projectdiscovery/dnsx v1.2.3 @@ -29,7 +30,6 @@ require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/klauspost/compress v1.18.6 // indirect github.com/kr/pretty v0.3.1 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/microcosm-cc/bluemonday v1.0.27 // indirect diff --git a/go.sum b/go.sum index 835017c..845ee26 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,5 @@ -github.com/Autumn-27/norma v0.2.9 h1:VTt8Oe7XoiV85Dc+QOPu+zQ7IykQRm2aDTEbLEzbBe4= -github.com/Autumn-27/norma v0.2.9/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw= +github.com/Autumn-27/norma v0.3.2 h1:3gW84SbSrz8luW8kbi2tv3O7CLdRMQvyylCHj1QUdNk= +github.com/Autumn-27/norma v0.3.2/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw= github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro= github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so= diff --git a/llmpool/pool.go b/llmpool/pool.go index 30b3665..be0ddff 100644 --- a/llmpool/pool.go +++ b/llmpool/pool.go @@ -122,6 +122,44 @@ func (p *Pool) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[ } } +// Complete implements llm.Provider with failover for non-streaming calls. A +// non-streaming request is atomic — it never delivers partial output — so every +// failover-eligible failure lands in the safe window and the next member can be +// tried without risk of duplicated output. Mirrors Stream's health-tripping and +// chain-exhaustion behavior. +func (p *Pool) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + order := p.order(req) + var lastErr error + for i, m := range order { + msg, sr, usage, err := m.Prov.Complete(ctx, req) + if err == nil { + p.health.Pass(m.ID) + return msg, sr, usage, nil + } + if !shouldFailover(ctx, err) { + // Non-failover error (e.g. ctx cancel, deterministic 4xx): surface as-is + // without tripping health, matching Stream's non-failover path. + p.health.Pass(m.ID) + return llm.Message{}, "", llm.Usage{}, err + } + lastErr = err + hard := isHardFailure(err) + if p.health.Trip(m.ID, trimErr(err), hard) { + log.Printf("[llmpool] 配置 %q(%s) 已熔断：%s", m.Name, m.Model, trimErr(err)) + } + 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(err)) + } + } + if lastErr == nil { + lastErr = ErrExhausted + } + log.Printf("[llmpool] 轮询链已耗尽(%d 个配置全部失败)，最后错误：%s", len(order), trimErr(lastErr)) + return llm.Message{}, "", llm.Usage{}, 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 diff --git a/llmpool/pool_test.go b/llmpool/pool_test.go index b1e69c8..21b7568 100644 --- a/llmpool/pool_test.go +++ b/llmpool/pool_test.go @@ -55,6 +55,17 @@ func (f *fakeProv) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } } +func (f *fakeProv) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + acc := llm.NewAccumulator() + for ev, err := range f.Stream(ctx, req) { + if err != nil { + return llm.Message{}, "", llm.Usage{}, err + } + acc.Add(ev) + } + return acc.Message(), acc.StopReason, acc.Usage, 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}}}} diff --git a/llmrec/llmrec.go b/llmrec/llmrec.go index 1bbd55c..71faf98 100644 --- a/llmrec/llmrec.go +++ b/llmrec/llmrec.go @@ -39,10 +39,10 @@ func TaskIDFrom(ctx context.Context) string { // Recorder wraps an llm.Provider and records every completion call. type Recorder struct { - inner llm.Provider - pg *db.DB - model string // model name (from config, not in CompletionRequest) - prof string // LLM profile name (from llm_profiles) + inner llm.Provider + pg *db.DB + model string // model name (from config, not in CompletionRequest) + prof string // LLM profile name (from llm_profiles) // thinkingType / reasoningEffort 是配置级思考参数(思考开关 / 思考强度)。它们在 // norma 的 buildBody() 里从 provider 配置注入真正的 HTTP body,不出现在 // CompletionRequest 上,故 Recorder 需在此单独带一份,序列化时写进录制。 @@ -179,6 +179,45 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } } +// Complete implements the atomic completion path while preserving the same +// usage metering and optional body recording semantics as Stream. +func (r *Recorder) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + recordBodies := r.enabled == nil || r.enabled() + start := time.Now() + session := transcript.SessionIDFrom(ctx) + parsedID, worker := parseSession(session) + expID := db.ParseExpID(parsedID) + taskID := TaskIDFrom(ctx) + if taskID == "" { + taskID = parsedID + } + + reqBody := "" + if recordBodies { + reqBody = r.serializeRequest(req) + } + + msg, stopReason, usage, err := r.inner.Complete(ctx, req) + status := "ok" + if err != nil { + status = "error" + } + r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) + if recordBodies { + var textBuf, thinkingBuf strings.Builder + for _, block := range msg.Content { + switch block.Type { + case llm.BlockText: + textBuf.WriteString(block.Text) + case llm.BlockThinking: + thinkingBuf.WriteString(block.Thinking) + } + } + r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err) + } + return msg, stopReason, usage, err +} + // recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips // zero-token calls with no model, which carry nothing worth metering. func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) { diff --git a/llmrec/llmrec_test.go b/llmrec/llmrec_test.go index e998684..19c7d4b 100644 --- a/llmrec/llmrec_test.go +++ b/llmrec/llmrec_test.go @@ -2,9 +2,28 @@ package llmrec import ( "context" + "iter" "testing" + + "github.com/Autumn-27/norma/llm" ) +type completeProvider struct{} + +func (completeProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + return func(func(llm.StreamEvent, error) bool) {} +} + +func (completeProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + return llm.Message{ + Role: llm.RoleAssistant, + Content: []llm.ContentBlock{ + {Type: llm.BlockThinking, Thinking: "reasoning"}, + llm.TextBlock("answer"), + }, + }, "stop", llm.Usage{InputTokens: 7, OutputTokens: 3}, nil +} + func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) { ctx := WithTaskID(context.Background(), " 42 ") if got := TaskIDFrom(ctx); got != "42" { @@ -27,3 +46,20 @@ func TestParseSessionFallbackIsExplorationScoped(t *testing.T) { t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker) } } + +func TestCompleteForwardsAtomicResponse(t *testing.T) { + recorder := Wrap(completeProvider{}, nil, "model", "profile", "", "", func() bool { return false }) + msg, stopReason, usage, err := recorder.Complete(context.Background(), llm.CompletionRequest{}) + if err != nil { + t.Fatalf("Complete() error: %v", err) + } + if got := msg.Text(); got != "answer" { + t.Fatalf("Complete() text=%q want answer", got) + } + if stopReason != "stop" { + t.Fatalf("Complete() stop reason=%q want stop", stopReason) + } + if usage.InputTokens != 7 || usage.OutputTokens != 3 { + t.Fatalf("Complete() usage=%+v", usage) + } +} diff --git a/server/commands.go b/server/commands.go index 12904a7..aa05427 100644 --- a/server/commands.go +++ b/server/commands.go @@ -15,14 +15,7 @@ func (s *Server) pgListCommands(w http.ResponseWriter, r *http.Request) { size := atoiDefault(q.Get("size"), 50) keyword := q.Get("q") - var expID *int64 - if tv := q.Get("task"); tv != "" { - if n, err := strconv.ParseInt(tv, 10, 64); err == nil { - expID = &n - } - } - - records, total, err := s.m.PG().ListCommands(expID, keyword, page, size) + records, total, err := s.m.PG().ListCommands(commandTaskFilter(q.Get("task")), keyword, page, size) if err != nil { writeErr(w, 500, err.Error()) return @@ -30,6 +23,37 @@ func (s *Server) pgListCommands(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"commands": records, "total": total}) } +// pgToolStats returns per-tool call counts for the tool-execution history, under +// the same task/keyword filters the list takes — the summary describes the whole +// filtered set, not the page on screen. +func (s *Server) pgToolStats(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + pg := s.m.PG() + if pg == nil { + writeJSON(w, 200, map[string]any{"stats": []db.ToolStat{}}) + return + } + stats, err := pg.ToolStats(commandTaskFilter(q.Get("task")), q.Get("q")) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"stats": stats}) +} + +// commandTaskFilter parses the ?task= exploration id. Absent or unparsable → nil +// (no filter), matching the list endpoint's long-standing lenient behaviour. +func commandTaskFilter(v string) *int64 { + if v == "" { + return nil + } + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return nil + } + return &n +} + // pgListLLMRecords returns paginated LLM call records. func (s *Server) pgListLLMRecords(w http.ResponseWriter, r *http.Request) { q := r.URL.Query() diff --git a/server/dto.go b/server/dto.go index cf7133d..c681bcf 100644 --- a/server/dto.go +++ b/server/dto.go @@ -580,6 +580,9 @@ type LLMProfileDTO struct { // pool_exclude=true 则不作为故障转移目标，但仍可被 agent/任务显式绑定。 Priority int `json:"priority"` PoolExclude bool `json:"pool_exclude"` + // 收发模式：true=流式(SSE) | false=非流式。没有 omitempty —— false 必须出现在 + // 响应里，否则前端读不到「非流式」，开关会回落成默认的流式。 + Streaming bool `json:"streaming"` } func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO { @@ -599,6 +602,7 @@ func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO { IsDefault: p.IsDefault, Priority: p.Priority, PoolExclude: p.PoolExclude, + Streaming: p.Streaming, } } diff --git a/server/goals.go b/server/goals.go index dbc497d..c8e6d98 100644 --- a/server/goals.go +++ b/server/goals.go @@ -383,7 +383,8 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity } 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) + goalRuntime := s.agentsForTask(t).runtime + decomposed := agent.DecomposeGoalsWithProvider(ctx, goalRuntime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, goalRuntime.nonStreaming(), emit) s.engine.EndLLMCall(t.ID) for _, g := range decomposed { if strings.TrimSpace(g.Text) != "" { diff --git a/server/orchestration.go b/server/orchestration.go index 9057ea7..a48d83b 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -450,9 +450,9 @@ func (s *Server) seedOrchestrationTools() { s.seedPlannerDefaultBindings() s.seedAutoReportFindingBinding() s.unbindGoalMetDefault() - s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性) + s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性) s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性) - s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性) + s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性) // 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把 // list_assets/insert_assets/report_finding/list_findings/list_companies 连同 // pentest 一起 seed 好了（项目尚无旧库，不做迁移）。 @@ -464,7 +464,7 @@ func (s *Server) seedOrchestrationTools() { // reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag. // Bump the flag whenever these tools' schemas/descriptions change in code. func (s *Server) refreshBuiltinToolSchemas() { - const flag = "tool_schema_refresh_v6_insert_assets_related" + const flag = "tool_schema_refresh_v7_list_facts_paging" if v, _, _ := s.m.pg.GetSetting(flag); v == "true" { return } @@ -480,7 +480,9 @@ func (s *Server) refreshBuiltinToolSchemas() { // “结束空轮”的手段、刚开跑就误判整个任务完成。 // - insert_assets：新增 related 入参(标记资产是否与当前任务相关、决定是否入覆盖度)， // SeedTool 首插入only，旧库已 seed 的 schema 否则收不到这个新参数。 - refreshBuiltin := map[string]bool{"goal_met": true, "insert_assets": true} + // - list_facts：改为分页，新增 limit/before/q 入参；旧库已 seed 的空 schema 否则 + // 在工具管理页显示「无参数」，模型也拿不到这几个参数说明。 + refreshBuiltin := map[string]bool{"goal_met": true, "insert_assets": true, "list_facts": true} for _, sd := range agent.BuiltinToolSeeds() { if !refreshBuiltin[sd.Key] { continue diff --git a/server/server.go b/server/server.go index 942c8ca..4968091 100644 --- a/server/server.go +++ b/server/server.go @@ -311,6 +311,7 @@ func (s *Server) loadLLMConfig() (agent.Config, bool) { cfg.ContextWindowK = p.ContextWindowK cfg.ThinkingType = p.ThinkingType cfg.ReasoningEffort = p.ReasoningEffort + cfg.Stream = p.Streaming if cfg.APIKey == "" { return cfg, false } @@ -328,14 +329,15 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error { } 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. + // otherwise this legacy endpoint would silently reset the profile's priority, + // pool_exclude, and streaming to zero values on every save. var priority int var poolExclude bool + streaming := true // 旧库/新建默认流式 if profs, _ := s.m.pg.ListProfiles(); profs != nil { for _, p := range profs { if p.Name == "default" { - id, priority, poolExclude = p.ID, p.Priority, p.PoolExclude + id, priority, poolExclude, streaming = p.ID, p.Priority, p.PoolExclude, p.Streaming break } } @@ -344,7 +346,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error { 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, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, - Priority: priority, PoolExclude: poolExclude, + Priority: priority, PoolExclude: poolExclude, Streaming: streaming, }) if err != nil { return err @@ -394,6 +396,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory"))) wk.SetWebSearch(s.webSearchFor("worker")) wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开;每轮读) + wk.SetNonStreaming(nonStreamingResolver(wCfg)) // 该 profile 选非流式时走 Provider.Complete pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg) pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner")) pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work @@ -401,9 +404,18 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy pl.SetWebSearch(s.webSearchFor("planner")) pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开;每轮读) + pl.SetNonStreaming(nonStreamingResolver(pCfg)) // 该 profile 选非流式时走 Provider.Complete return pl, wk } +// nonStreamingResolver returns a resolver capturing a profile's streaming choice. +// The agents read it per run; a profile change rebuilds the agents (applyLLM), +// so the captured value is always the one in effect for this build. +func nonStreamingResolver(cfg agent.Config) func() bool { + nonStreaming := !cfg.Stream + return func() bool { return nonStreaming } +} + // 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 { @@ -447,12 +459,14 @@ func (s *Server) applyLLM(cfg agent.Config) error { s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy s.mainAgent.SetWebSearch(s.webSearchFor("mainagent")) s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏 + s.mainAgent.SetNonStreaming(nonStreamingResolver(mCfg)) // chat agent serves MANY custom agents by key → it holds the GLOBAL opts // (backend/key) and gates Enabled per-conversation-agent at Chat time. 对话始终用激活配置。 s.chatAgent = agent.NewChatAgent(prov, cfg.Model, s.m.dir, tx, win) // chat page runner s.chatAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) s.chatAgent.SetWebSearch(s.m.WebSearchOpts()) s.chatAgent.SetGuard(s.chatGuard()) + s.chatAgent.SetNonStreaming(nonStreamingResolver(cfg)) s.llmProv = prov s.llmCfg = cfg s.llmOn = true @@ -478,6 +492,7 @@ func (s *Server) loadProfileConfig(id int64) (agent.Config, bool) { cfg.ContextWindowK = p.ContextWindowK cfg.ThinkingType = p.ThinkingType cfg.ReasoningEffort = p.ReasoningEffort + cfg.Stream = p.Streaming if cfg.APIKey == "" { return cfg, false } @@ -629,6 +644,7 @@ func (s *Server) chatAgentForProfile(id int64) *agent.ChatAgent { ca.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) ca.SetWebSearch(s.m.WebSearchOpts()) ca.SetGuard(s.chatGuard()) + ca.SetNonStreaming(nonStreamingResolver(cfg)) s.profMu.Lock() if ex := s.profChatAgents[id]; ex != nil { // lost the race → keep the winner ca = ex @@ -786,6 +802,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/traffic/exchange", s.getTrafficExchange) mux.HandleFunc("GET /api/traffic/blob", s.getTrafficBlob) mux.HandleFunc("GET /api/commands", s.pgListCommands) + mux.HandleFunc("GET /api/commands/stats", s.pgToolStats) // 按工具聚合调用次数 mux.HandleFunc("GET /api/llm/records", s.pgListLLMRecords) mux.HandleFunc("DELETE /api/llm/records", s.pgDeleteLLMRecords) mux.HandleFunc("GET /api/llm/records/tasks", s.pgLLMTasks) @@ -1335,6 +1352,7 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { ThinkingType string `json:"thinking_type"` ReasoningEffort string `json:"reasoning_effort"` ProfileID *int64 `json:"profile_id"` // 测已存 profile 时传入：api_key 为空则用它存的 key + Streaming *bool `json:"streaming"` // 省略=流式，与保存 profile 时同一套默认 } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { writeErr(w, 400, err.Error()) @@ -1345,6 +1363,11 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { // reasoning_effort/thinking field fails the test too (no false "test ok, run 400"). cfg.ThinkingType = req.ThinkingType cfg.ReasoningEffort = req.ReasoningEffort + // 同理，收发模式也照该 profile 的选择来：只支持其中一种通道的端点必须在这里就 + // 暴露，而不是等会话里才发现"测试通过的配置根本跑不动"。 + if req.Streaming != nil { + cfg.Stream = *req.Streaming + } // API Key 解析优先级：表单输入 > 指定 profile 存的 key > 全局配置的 key。 // 已存 profile 的 key 不回传浏览器，所以测试已存配置时表单为空，需从 DB 取。 if cfg.APIKey == "" && req.ProfileID != nil { @@ -1361,12 +1384,27 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"ok": false, "error": "未提供 API Key"}) return } - lat, err := agent.TestConnection(r.Context(), cfg) + lat, reply, err := agent.TestConnection(r.Context(), cfg) if err != nil { writeJSON(w, 200, map[string]any{"ok": false, "error": err.Error()}) return } - writeJSON(w, 200, map[string]any{"ok": true, "latency_ms": lat.Milliseconds(), "model": cfg.Model}) + // 回传模型实际回复，让"测试通过"有据可查：看得见它确实说了话，而不只是 HTTP 200。 + writeJSON(w, 200, map[string]any{ + "ok": true, "latency_ms": lat.Milliseconds(), "model": cfg.Model, "reply": truncateReply(reply), + }) +} + +// truncateReply clips a connection-test reply for display. A model told to answer +// "OK" can still ramble (or think out loud); the UI only needs enough to show it +// really said something. Rune-based so multibyte text never splits mid-character. +func truncateReply(s string) string { + const max = 200 + r := []rune(s) + if len(r) <= max { + return s + } + return string(r[:max]) + "…" } type createTaskReq struct { diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 3b2fad7..02baba2 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -1736,10 +1736,14 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) { return } // APIKey has json:"-" on db.LLMProfile so it is never leaked to the UI on - // read; accept it here via a sibling field for create/update. + // read; accept it here via a sibling field for create/update. Streaming is a + // *bool shadowing the embedded json:"streaming": an absent field must default + // to streaming (true), which a plain bool zero value (false = non-streaming) + // would get wrong — legacy/partial clients that never send it stay streaming. var body struct { db.LLMProfile - APIKey string `json:"api_key"` + APIKey string `json:"api_key"` + Streaming *bool `json:"streaming"` } if err := decode(r, &body); err != nil { writeErr(w, 400, err.Error()) @@ -1747,6 +1751,7 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) { } p := body.LLMProfile p.APIKey = body.APIKey + p.Streaming = body.Streaming == nil || *body.Streaming id, err := pg.SaveProfile(&p) if err != nil { writeErr(w, 500, err.Error()) diff --git a/server/task_llm.go b/server/task_llm.go index 01a94b7..5a22600 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -139,6 +139,36 @@ func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) { return t, 0, pt.LLMChainRevision, prov, nil } +// nonStreaming reports whether the task's currently-active LLM source is set to +// non-streaming. It mirrors current()'s source precedence (agent binding → +// active chain profile → global) but reads only the streaming flag. Read-only +// and best-effort: any resolution error falls back to streaming (false), the +// safe default. If failover switches to a profile with a different streaming +// setting, the change takes effect on the next agent run (a fresh Session is +// built per run in captureRun). +func (r *taskLLMRuntime) nonStreaming() bool { + taskNum, err := parseTaskID(r.taskID) + if err != nil { + return false + } + pt, err := r.s.m.pg.GetTask(taskNum) + if err != nil || pt == nil { + return false + } + if _, cfg, ok := r.s.agentBindingProvider(r.agentKey); ok { + return !cfg.Stream + } + if len(pt.LLMProfileIDs) > 0 && pt.ActiveLLMProfileID != nil { + if _, cfg, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID); ok { + return !cfg.Stream + } + } + if _, cfg, ok := r.s.globalProvider(); ok { + return !cfg.Stream + } + return false +} + func parseTaskID(id string) (int64, error) { n, err := strconv.ParseInt(id, 10, 64) if err != nil || n <= 0 { @@ -147,9 +177,11 @@ func parseTaskID(id string) (int64, error) { 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{ +// streamHooks builds the failover/exhaustion callbacks shared by Stream and +// Complete: how to read the current profile selection, how to mark it quota +// exhausted, and how to emit a failover transition. +func (r *taskLLMRuntime) streamHooks() taskLLMStreamHooks { + return taskLLMStreamHooks{ current: func() (taskLLMSelection, error) { task, profileID, revision, provider, err := r.current() return taskLLMSelection{task: task, profileID: profileID, revision: revision, provider: provider}, err @@ -169,7 +201,70 @@ func (r *taskLLMRuntime) Stream(ctx context.Context, req llm.CompletionRequest) r.s.emitTaskLLMTransition(selection.task, transition, cause) }, } - return streamTaskLLM(ctx, r.taskID, req, hooks) +} + +func (r *taskLLMRuntime) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + ctx = llmrec.WithTaskID(ctx, r.taskID) + return streamTaskLLM(ctx, r.taskID, req, r.streamHooks()) +} + +// Complete is the non-streaming counterpart of Stream. A non-streaming call is +// atomic — it never delivers partial output — so every failure is safe to retry +// on the same provider or fail over to the next profile without risking +// duplicated model output or tool execution (the "committed" bookkeeping the +// streaming path needs is unnecessary here). +func (r *taskLLMRuntime) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + ctx = llmrec.WithTaskID(ctx, r.taskID) + return completeTaskLLM(ctx, r.taskID, req, r.streamHooks()) +} + +func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) (llm.Message, string, llm.Usage, error) { + for { + selection, err := hooks.current() + if err != nil { + return llm.Message{}, "", llm.Usage{}, err + } + var ( + msg llm.Message + sr string + usage llm.Usage + callErr error + ) + // 同 provider 安全窗口重试:非流式调用要么整体成功、要么整体失败,没有 + // 中途已交付输出的问题,所以任何瞬时失败都可原样重试。 + for attempt := 0; ; attempt++ { + msg, sr, usage, callErr = selection.provider.Complete(ctx, req) + if callErr != nil && ctx.Err() == nil && + attempt < sameProviderStreamRetries && isRetryableStreamError(callErr) { + backoff := sameProviderRetryBackoff(attempt) + log.Printf("[task-llm] task %s 非流式调用失败,%v 后同 provider 重试 (%d/%d): %v", + taskID, backoff, attempt+1, sameProviderStreamRetries, callErr) + if sleepCtx(ctx, backoff) { + break // 退避期间 ctx 取消 → 停止重试 + } + continue + } + break + } + if callErr == nil { + return msg, sr, usage, nil + } + // profileID=0: 显式链已被清空;非额度错误:透传。二者都不改任务 failover 状态。 + if selection.profileID == 0 || !isQuotaExhaustedError(callErr) { + return llm.Message{}, "", llm.Usage{}, callErr + } + transition, markErr := hooks.exhaust(selection, callErr) + if markErr != nil { + return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr) + } + if transition.Advanced && !transition.Stale && hooks.transition != nil { + hooks.transition(selection, transition, callErr) + } + if !transition.Stale && transition.NextProfileID == nil { + return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr} + } + // 没有任何输出交付给调用方,换下一个 profile 重放同一逻辑请求是安全的。 + } } func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) iter.Seq2[llm.StreamEvent, error] { @@ -435,6 +530,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { window := workerRuntime.CompactionWindow() wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker")) wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow) + wk.SetNonStreaming(workerRuntime.nonStreaming) // 按任务当前激活 profile 的流式开关(每轮读) wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory"))) @@ -442,6 +538,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开) pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner")) pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow) + pl.SetNonStreaming(plannerRuntime.nonStreaming) pl.SetKillWork(s.engine.KillWork) pl.SetSteerWork(s.engine.SteerWork) pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) @@ -449,6 +546,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开) main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent")) main.SetCompactionWindowResolver(mainRuntime.CompactionWindow) + main.SetNonStreaming(mainRuntime.nonStreaming) main.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) main.SetWebSearch(s.webSearchFor("mainagent")) main.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏 diff --git a/server/task_llm_test.go b/server/task_llm_test.go index 48da7d3..5316b9b 100644 --- a/server/task_llm_test.go +++ b/server/task_llm_test.go @@ -37,6 +37,23 @@ func (p *scriptedLLMProvider) Stream(context.Context, llm.CompletionRequest) ite } } +func (p *scriptedLLMProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + return accumulateStreamForTest(ctx, p.Stream, req) +} + +// accumulateStreamForTest drains a mock's Stream to satisfy the non-streaming +// Complete method. +func accumulateStreamForTest(ctx context.Context, stream func(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error], req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + acc := llm.NewAccumulator() + for ev, err := range stream(ctx, req) { + if err != nil { + return llm.Message{}, "", llm.Usage{}, err + } + acc.Add(ev) + } + return acc.Message(), acc.StopReason, acc.Usage, nil +} + // withZeroRetryBackoff sets the same-provider retry backoff to zero for the // duration of a (serial) test and returns a restore func for defer. func withZeroRetryBackoff() func() { @@ -73,6 +90,10 @@ func (p *flakyThenOKProvider) Stream(context.Context, llm.CompletionRequest) ite } } +func (p *flakyThenOKProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) { + return accumulateStreamForTest(ctx, p.Stream, req) +} + func collectTaskLLMStream(seq iter.Seq2[llm.StreamEvent, error]) ([]llm.StreamEvent, error) { var events []llm.StreamEvent var streamErr error diff --git a/web/src/app/(main)/function/commands/page.tsx b/web/src/app/(main)/function/commands/page.tsx index 7d58b48..6a5954e 100644 --- a/web/src/app/(main)/function/commands/page.tsx +++ b/web/src/app/(main)/function/commands/page.tsx @@ -2,17 +2,18 @@ import * as React from "react"; -import { ChevronLeftIcon, ChevronRightIcon, Loader2Icon, SearchIcon, TerminalIcon } from "lucide-react"; +import { BarChart3Icon, ChevronLeftIcon, ChevronRightIcon, Loader2Icon, SearchIcon, TerminalIcon } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { api } from "@/lib/api"; -import type { CommandRecord } from "@/lib/types"; +import type { CommandRecord, ToolStat } from "@/lib/types"; import { cn } from "@/lib/utils"; function fmtTime(ts: string) { @@ -59,6 +60,11 @@ export default function CommandsPage() { const [total, setTotal] = React.useState(0); const [loading, setLoading] = React.useState(false); + // 各工具调用次数。弹窗打开时才拉取（多一次聚合查询，不必每次翻页都付）。 + const [statsOpen, setStatsOpen] = React.useState(false); + const [stats, setStats] = React.useState([]); + const [statsLoading, setStatsLoading] = React.useState(false); + // Selected execution is rendered in a right-side detail sheet. const [selected, setSelected] = React.useState(null); @@ -96,6 +102,24 @@ export default function CommandsPage() { }; }, [page, size, queryQ, taskFilter]); + // 统计跟随筛选条件走，和表格描述的是同一批记录（但不分页）。 + React.useEffect(() => { + if (!statsOpen) return; + let alive = true; + setStatsLoading(true); + api + .commandStats({ task: taskFilter || undefined, q: queryQ || undefined }) + .then((r) => alive && setStats(r.stats ?? [])) + .catch(() => alive && setStats([])) + .finally(() => alive && setStatsLoading(false)); + return () => { + alive = false; + }; + }, [statsOpen, queryQ, taskFilter]); + + const statsTotal = stats.reduce((n, s) => n + s.total, 0); + const statsErrors = stats.reduce((n, s) => n + s.errors, 0); + const statsMax = stats.reduce((n, s) => Math.max(n, s.total), 0); const totalPages = Math.max(1, Math.ceil(total / size)); const rangeStart = total === 0 ? 0 : page * size + 1; const rangeEnd = page * size + commands.length; @@ -141,6 +165,11 @@ export default function CommandsPage() { + +
{rangeStart}–{rangeEnd} / {total} @@ -243,6 +272,66 @@ export default function CommandsPage() {
+ {/* 工具调用统计：与表格同一批记录（同筛选、不分页） */} + + + + 工具调用统计 + + {taskFilter || queryQ ? "当前筛选条件下的全部记录" : "全部工具执行记录"} + {stats.length > 0 && ( + <> + {" · "} + {stats.length} 个工具 ·{" "} + {statsTotal} 次调用 + {statsErrors > 0 && ( + <> + {" · 失败 "} + {statsErrors} + + )} + + )} + + + + {statsLoading && stats.length === 0 ? ( +
+ +
+ ) : stats.length === 0 ? ( +
暂无统计数据
+ ) : ( +
+ {stats.map((s) => ( +
+
+
+ {s.tool} + {s.errors > 0 && ( + 失败 {s.errors} + )} +
+
+
0 ? (s.total / statsMax) * 100 : 0}%` }} + /> +
+
+
+
{s.total}
+
+ {statsTotal > 0 ? ((s.total / statsTotal) * 100).toFixed(1) : "0.0"}% +
+
+
+ ))} +
+ )} + +
+ !open && setSelected(null)}> {selected && ( diff --git a/web/src/app/(main)/system/llm/page.tsx b/web/src/app/(main)/system/llm/page.tsx index b08c6ef..35c7806 100644 --- a/web/src/app/(main)/system/llm/page.tsx +++ b/web/src/app/(main)/system/llm/page.tsx @@ -309,6 +309,7 @@ function ProfileSheet({ const [effort, setEffort] = React.useState(NONE); const [priority, setPriority] = React.useState("0"); // 轮询顺位;越大越先 const [poolExclude, setPoolExclude] = React.useState(false); + const [streaming, setStreaming] = React.useState(true); // true=流式(默认);false=非流式 const [testing, setTesting] = React.useState(false); const [saving, setSaving] = React.useState(false); const [models, setModels] = React.useState([]); @@ -331,6 +332,7 @@ function ProfileSheet({ setEffort(fromStore(profile?.reasoning_effort)); setPriority(String(profile?.priority ?? 0)); setPoolExclude(profile?.pool_exclude ?? false); + setStreaming(profile?.streaming ?? true); setApiKey(""); setKeyHint(profile?.api_key_hint ?? ""); setModels([]); @@ -365,8 +367,22 @@ function ProfileSheet({ try { // 用配置实际会跑的思考参数来测，这样不支持该字段的模型在这里就失败， // 而不是等到跑任务时才炸。传 profile id：Key 输入框留空时用已存的 Key。 - const r = await api.testLLM(format, model, baseUrl, apiKey, proxy, toStore(thinkingType), toStore(effort), profileId); - if (r.ok) toast.success(`连接成功 · ${r.latency_ms ?? "?"}ms · ${r.model ?? model}`); + const r = await api.testLLM( + format, + model, + baseUrl, + apiKey, + proxy, + toStore(thinkingType), + toStore(effort), + profileId, + streaming, + ); + // 回复内容一并展示：看得见模型确实说了话，才算和会话里跑通是一回事。 + if (r.ok) + toast.success(`连接成功 · ${r.latency_ms ?? "?"}ms · ${r.model ?? model}`, { + description: r.reply ? `回复：${r.reply}` : undefined, + }); else toast.error(`连接失败：${r.error ?? "未知"}`); } catch (e) { toast.error(`测试出错：${(e as Error).message}`); @@ -398,6 +414,7 @@ function ProfileSheet({ reasoning_effort: toStore(effort), priority: Number(priority) || 0, pool_exclude: poolExclude, + streaming, }); if (isNew) toast.success(`已新建：${name.trim()}（在卡片上「设为激活」以启用）`); else toast.success(profile?.is_default ? "已保存，激活配置即时生效，无需重启" : "已保存"); @@ -525,7 +542,8 @@ function ProfileSheet({ onChange={(e) => setProxy(e.target.value)} />

- 仅 LLM 出站请求走此代理，支持 http/https/socks5，可带账号密码（如 socks5://user:pass@host:port，密码含特殊字符需 URL 编码）；留空表示不使用代理（直连）。 + 仅 LLM 出站请求走此代理，支持 http/https/socks5，可带账号密码（如 + socks5://user:pass@host:port，密码含特殊字符需 URL 编码）；留空表示不使用代理（直连）。

@@ -595,6 +613,17 @@ function ProfileSheet({
+
+
+ +

+ 开启（默认）走流式 SSE，有运行中实时进度与实时 token 计数。 关闭则走真·非流式（stream:false， + 一次性返回完整响应）——可绕开部分网关糟糕的 SSE 实现（空帧 / 思考字段丢帧）， + 代价是失去运行中的实时进度。 +

+
+ +
@@ -602,8 +631,8 @@ function ProfileSheet({

- 控制是否发送 thinking 字段。不发送=不带该字段（兼容 MiniMax 等不支持 - 的模型）；关闭=发 disabled；开启=发 enabled。与下面的强度互相独立。 + 控制是否发送 thinking 字段。不发送=不带该字段（兼容 MiniMax 等不支持 的模型）；关闭=发 + disabled；开启=发 enabled。与下面的强度互相独立。

diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index 9a9f3b7..ba197c9 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -74,6 +74,7 @@ import type { TokenTotal, TokenUsage, Tool, + ToolStat, TrafficDetail, TrafficHost, TrafficResp, @@ -677,8 +678,10 @@ export const api = { thinking_type = "", reasoning_effort = "", profile_id?: number, + streaming = true, // 用该配置真实的收发模式来测，别让"流式能通、非流式不通"漏到会话里 ) => - post<{ ok: boolean; error?: string; latency_ms?: number; model?: string }>("/llm/test", { + // reply = 模型实际回复(已截断);一个字都不回的配置后端直接判失败 + post<{ ok: boolean; error?: string; latency_ms?: number; model?: string; reply?: string }>("/llm/test", { provider, model, base_url, @@ -687,6 +690,7 @@ export const api = { thinking_type, reasoning_effort, profile_id, + streaming, }), llmProfiles: () => get<{ profiles: LLMProfile[] }>("/llm/profiles").then((r) => arr(r.profiles)), saveLLMProfile: (p: { @@ -704,6 +708,7 @@ export const api = { reasoning_effort?: string; // ""(不发送)|"low"|"medium"|"high"|"xhigh"|"max" priority?: number; // 轮询顺位，越大越先用 pool_exclude?: boolean; // true=不作为故障转移目标 + streaming?: boolean; // true(默认)=流式 | false=非流式 }) => post<{ id: number }>("/llm/profiles", p), deleteLLMProfile: (id: string) => del<{ deleted: number }>(`/llm/profiles/${id}`), activateLLMProfile: (id: string) => post<{ ok: boolean }>("/llm/profiles/active", { id: Number(id) }), @@ -972,6 +977,13 @@ export const api = { sp.set("size", String(params?.size ?? 50)); return get<{ commands: CommandRecord[]; total: number }>(`/commands?${sp}`); }, + // 各工具调用次数；沿用列表的 task/q 筛选，统计的是整个结果集而非当前页。 + commandStats: (params?: { task?: string; q?: string }) => { + const sp = new URLSearchParams(); + if (params?.task) sp.set("task", params.task); + if (params?.q) sp.set("q", params.q); + return get<{ stats: ToolStat[] }>(`/commands/stats?${sp}`); + }, // ---- LLM records ---- llmRecords: (params?: { model?: string; session?: string; task?: string; page?: number; size?: number }) => { diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 6002dae..407fe4e 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -1493,6 +1493,17 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re // ── 工具执行历史 ── if (path === "/commands" && m === "GET") return { commands: D.commandRecords, total: D.commandRecords.length }; + if (path === "/commands/stats" && m === "GET") { + const tally = new Map(); + for (const c of D.commandRecords) { + const tool = c.tool || "-"; + const s = tally.get(tool) ?? { tool, total: 0, errors: 0 }; + s.total++; + if (c.is_error) s.errors++; + tally.set(tool, s); + } + return { stats: [...tally.values()].sort((a, b) => b.total - a.total || a.tool.localeCompare(b.tool)) }; + } // ── LLM ── if (path === "/llm/records" && m === "GET") return { records: mockLLMRecords, total: mockLLMRecords.length }; @@ -1509,7 +1520,8 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re } if (path === "/llm" && m === "GET") return D.llmConfig; if (path === "/llm" && m === "POST") return { ok: true }; - if (path === "/llm/test") return { ok: true, latency_ms: 128, model: String(b.model ?? "claude-opus-4-8") }; + if (path === "/llm/test") + return { ok: true, latency_ms: 128, model: String(b.model ?? "claude-opus-4-8"), reply: "OK" }; 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 }; diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 633688e..5c672d1 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -869,6 +869,8 @@ export interface LLMProfile { priority?: number; // true = 不作为故障转移目标（仍可被 agent/任务显式绑定使用）。 pool_exclude?: boolean; + // true（默认）= 流式(SSE) | false = 真·非流式(stream:false，一次性返回)。 + streaming?: boolean; } // ---- LLM 轮询（故障转移）---- @@ -1133,6 +1135,13 @@ export interface CommandRecord { created_at: string; } +// 单个工具的调用统计（/commands/stats）；errors 为其中失败的次数。 +export interface ToolStat { + tool: string; + total: number; + errors: number; +} + // ---- LLM recording ---- export interface LLMRecordItem { id: number; From 1fcc4b0588b6b056efdb4734ca03673234e8afd8 Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 14:16:00 +0800 Subject: [PATCH 5/7] fix: stream large task archive snapshots --- db/task_archives.go | 41 +++++++++++++++++++++++++++++++++------- db/task_archives_test.go | 32 +++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/db/task_archives.go b/db/task_archives.go index bef9f3a..c21c89d 100644 --- a/db/task_archives.go +++ b/db/task_archives.go @@ -1,6 +1,7 @@ package db import ( + "bytes" "context" "database/sql" "encoding/json" @@ -423,11 +424,37 @@ func (d *DB) FailTaskArchiveJob(id int64, activeState string, cause error) error } func queryArchiveRows(q interface { - QueryRow(query string, args ...any) *sql.Row -}, inner string, args ...any) (json.RawMessage, error) { - var raw []byte - err := q.QueryRow(`SELECT COALESCE(jsonb_agg(to_jsonb(row_data)), '[]'::jsonb) FROM (`+inner+`) row_data`, args...).Scan(&raw) - return json.RawMessage(raw), err + Query(query string, args ...any) (*sql.Rows, error) +}, inner string, args ...any) (json.RawMessage, int64, error) { + // Do not aggregate the result in PostgreSQL. A jsonb array has a hard limit + // of 256 MiB for its elements, which large LLM request/response histories can + // exceed even though every individual record is valid. Reading row JSON in + // order also avoids building a second copy of the full table in PostgreSQL. + rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + + var output bytes.Buffer + output.WriteByte('[') + var count int64 + for rows.Next() { + var raw []byte + if err := rows.Scan(&raw); err != nil { + return nil, 0, err + } + if count > 0 { + output.WriteByte(',') + } + output.Write(raw) + count++ + } + if err := rows.Err(); err != nil { + return nil, 0, err + } + output.WriteByte(']') + return json.RawMessage(output.Bytes()), count, nil } func rawRowCount(raw json.RawMessage) int64 { @@ -492,12 +519,12 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { } counts := make(map[string]int64, len(queries)) for _, query := range queries { - raw, err := queryArchiveRows(tx, query.query, query.args...) + raw, count, err := queryArchiveRows(tx, query.query, query.args...) if err != nil { return nil, fmt.Errorf("snapshot %s: %w", query.name, err) } tables[query.name] = raw - counts[query.name] = rawRowCount(raw) + counts[query.name] = count } assetIDs, exclusiveAssetIDs, hosts, exclusiveHosts, err := archiveAssetMetadata(tx, taskID, expID, tables["assets"]) diff --git a/db/task_archives_test.go b/db/task_archives_test.go index dbe282f..c9a6d18 100644 --- a/db/task_archives_test.go +++ b/db/task_archives_test.go @@ -8,6 +8,38 @@ import ( "time" ) +func TestQueryArchiveRowsStreamsJSON(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + payload := strings.Repeat("large request/response payload ", 64*1024) + raw, count, err := queryArchiveRows(d, `SELECT * FROM (VALUES +($1::bigint,$2::text),($3::bigint,$4::text)) archived_row(id,body) ORDER BY id`, + int64(1), payload, int64(2), "quoted: \"value\"\nline") + if err != nil { + t.Fatal(err) + } + if count != 2 { + t.Fatalf("row count=%d, want 2", count) + } + var rows []struct { + ID int64 `json:"id"` + Body string `json:"body"` + } + if err := json.Unmarshal(raw, &rows); err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("decoded row count=%d, want 2", len(rows)) + } + if rows[0].ID != 1 || rows[0].Body != payload || rows[1].ID != 2 || rows[1].Body != "quoted: \"value\"\nline" { + t.Fatal("streamed rows were reordered or truncated") + } +} + func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { From 0ec45bb45d5351a8d61127cd012f979d073252ca Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 28 Aug 2026 15:25:20 +0800 Subject: [PATCH 6/7] fix: prevent OOM when archiving large tasks --- db/commands.go | 16 ++- db/db.go | 135 ++++++++++++++++------- db/db_test.go | 47 ++++++++ db/llm_usage.go | 14 ++- db/schema.sql | 3 +- db/task_archives.go | 123 ++++++++++++++++++--- db/task_archives_restore.go | 101 +++++++++++++++-- db/task_archives_test.go | 161 ++++++++++++++++++++++++++-- server/task_archive_package.go | 13 ++- server/task_archive_package_test.go | 16 ++- server/task_archives.go | 63 +++++++++-- 11 files changed, 605 insertions(+), 87 deletions(-) diff --git a/db/commands.go b/db/commands.go index d95f0b8..3edefa6 100644 --- a/db/commands.go +++ b/db/commands.go @@ -189,11 +189,21 @@ ALTER TABLE llm_records ADD COLUMN IF NOT EXISTS raw_response TEXT; // EnsureLLMRecordsTable creates the llm_records table if it does not exist. func (d *DB) EnsureLLMRecordsTable() error { - if _, err := d.Exec(llmRecordsSchema); err != nil { + tx, err := d.Begin() + if err != nil { return err } - _, err := d.Exec(llmRecordsMigrate) - return err + defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return err + } + if _, err := tx.Exec(llmRecordsSchema); err != nil { + return err + } + if _, err := tx.Exec(llmRecordsMigrate); err != nil { + return err + } + return tx.Commit() } // InsertLLMRecord stores one LLM call record. diff --git a/db/db.go b/db/db.go index de0a738..20ec514 100644 --- a/db/db.go +++ b/db/db.go @@ -3,19 +3,84 @@ package db import ( + "context" "database/sql" _ "embed" + "errors" "fmt" "net/url" "strings" + "time" "github.com/Autumn-27/artex/config" + "github.com/jackc/pgx/v5/pgconn" _ "github.com/jackc/pgx/v5/stdlib" // pgx database/sql driver ("pgx") ) //go:embed schema.sql var schemaSQL string +const schemaMigrationLockKey int64 = 7337741001 + +var schemaDeadlockRetryDelays = [...]time.Duration{ + 100 * time.Millisecond, + 250 * time.Millisecond, + 500 * time.Millisecond, + time.Second, +} + +type schemaExecer interface { + ExecContext(context.Context, string, ...any) (sql.Result, error) +} + +func isPostgresDeadlock(err error) bool { + var pgErr *pgconn.PgError + return errors.As(err, &pgErr) && pgErr.Code == "40P01" +} + +func applySchemaWithRetry(ctx context.Context, execer schemaExecer, sleep func(time.Duration)) error { + for attempt := 0; ; attempt++ { + if _, err := execer.ExecContext(ctx, schemaSQL); err != nil { + if !isPostgresDeadlock(err) || attempt >= len(schemaDeadlockRetryDelays) { + return err + } + sleep(schemaDeadlockRetryDelays[attempt]) + continue + } + return nil + } +} + +// withSchemaMigrationLock pins the session-level lock to one checked-out +// connection. Running pg_advisory_lock through *sql.DB is incorrect because a +// later schema or unlock call may use a different pooled PostgreSQL session. +func withSchemaMigrationLock(ctx context.Context, sqlDB *sql.DB, action func(*sql.Conn) error) (err error) { + conn, err := sqlDB.Conn(ctx) + if err != nil { + return err + } + defer conn.Close() + if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, schemaMigrationLockKey); err != nil { + return fmt.Errorf("advisory lock: %w", err) + } + defer func() { + if _, unlockErr := conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock($1)`, schemaMigrationLockKey); unlockErr != nil && err == nil { + err = fmt.Errorf("advisory unlock: %w", unlockErr) + } + }() + return action(conn) +} + +// coordinateWithSchemaMigration makes long, multi-table archive transactions +// mutually exclusive with startup DDL while allowing ordinary runtime queries +// to continue normally. +func coordinateWithSchemaMigration(tx *sql.Tx) error { + if _, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, schemaMigrationLockKey); err != nil { + return fmt.Errorf("coordinate with schema migration: %w", err) + } + return nil +} + // DSN resolves the PostgreSQL connection string and reports where it came from. // Precedence: env ARTEX_PG_DSN > config file (config.json). There is no // built-in default — it errors if neither source is configured. @@ -72,28 +137,22 @@ func Open(dsn string) (*DB, error) { sqlDB.Close() return nil, fmt.Errorf("ping postgres (%s): %w", config.Redact(dsn), err) } - // Serialize schema migration + seeding across concurrent openers (e.g. parallel - // test packages) using a session-level advisory lock. The lock is released - // automatically when the connection is returned to the pool or closed. - if _, err := sqlDB.Exec(`SELECT pg_advisory_lock(7337741001)`); err != nil { - sqlDB.Close() - return nil, fmt.Errorf("advisory lock: %w", err) - } - // pgx runs multi-statement Exec via the simple protocol when there are no args. - if _, err := sqlDB.Exec(schemaSQL); err != nil { - sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck - sqlDB.Close() - return nil, fmt.Errorf("apply schema: %w", err) - } d := &DB{sqlDB} - if err := d.seedBuiltins(); err != nil { - sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck + // pgx runs multi-statement Exec via the simple protocol when there are no args. + // Keep the dedicated lock connection checked out until both DDL and seeding + // finish so concurrent application instances cannot initialize out of order. + err = withSchemaMigrationLock(context.Background(), sqlDB, func(conn *sql.Conn) error { + if err := applySchemaWithRetry(context.Background(), conn, time.Sleep); err != nil { + return fmt.Errorf("apply schema: %w", err) + } + if err := d.seedBuiltins(); err != nil { + return fmt.Errorf("seed builtins: %w", err) + } + return nil + }) + if err != nil { sqlDB.Close() - return nil, fmt.Errorf("seed builtins: %w", err) - } - if _, err := sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`); err != nil { - sqlDB.Close() - return nil, fmt.Errorf("advisory unlock: %w", err) + return nil, err } return d, nil } @@ -375,30 +434,30 @@ func (d *DB) seedDefaultInterceptRules() error { // 2. Python HTTP 客户端 .delete() 方法 // 3. JS/通用脚本里的 method: 'DELETE' / method="DELETE" { - name: "[内置] curl / wget 发送 DELETE 请求", - target: "tool_input", - typ: "regex", - pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`, - action: "deny", - message: "禁止通过 curl/wget 发送 HTTP DELETE 请求，可能删除目标系统数据", + name: "[内置] curl / wget 发送 DELETE 请求", + target: "tool_input", + typ: "regex", + pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`, + action: "deny", + message: "禁止通过 curl/wget 发送 HTTP DELETE 请求，可能删除目标系统数据", priority: 80, }, { - name: "[内置] Python HTTP 客户端 DELETE（requests/httpx/aiohttp）", - target: "tool_input", - typ: "regex", - pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`, - action: "deny", - message: "禁止使用 Python HTTP 客户端发送 DELETE 请求", + name: "[内置] Python HTTP 客户端 DELETE（requests/httpx/aiohttp）", + target: "tool_input", + typ: "regex", + pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`, + action: "deny", + message: "禁止使用 Python HTTP 客户端发送 DELETE 请求", priority: 80, }, { - name: "[内置] 脚本中声明 HTTP DELETE 方法（JS/通用）", - target: "tool_input", - typ: "regex", - pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`, - action: "deny", - message: "禁止在脚本中声明并发送 HTTP DELETE 请求", + name: "[内置] 脚本中声明 HTTP DELETE 方法（JS/通用）", + target: "tool_input", + typ: "regex", + pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`, + action: "deny", + message: "禁止在脚本中声明并发送 HTTP DELETE 请求", priority: 80, }, { diff --git a/db/db_test.go b/db/db_test.go index d549cdd..41c5c05 100644 --- a/db/db_test.go +++ b/db/db_test.go @@ -1,9 +1,56 @@ package db import ( + "context" + "database/sql" + "errors" "testing" + "time" + + "github.com/jackc/pgx/v5/pgconn" ) +type schemaExecFunc func(context.Context, string, ...any) (sql.Result, error) + +func (f schemaExecFunc) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + return f(ctx, query, args...) +} + +func TestApplySchemaRetriesOnlyDeadlocks(t *testing.T) { + deadlockAttempts := 0 + delays := []time.Duration{} + err := applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) { + deadlockAttempts++ + if deadlockAttempts < 3 { + return nil, &pgconn.PgError{Code: "40P01", Message: "deadlock detected"} + } + return nil, nil + }), func(delay time.Duration) { + delays = append(delays, delay) + }) + if err != nil { + t.Fatal(err) + } + if deadlockAttempts != 3 { + t.Fatalf("schema attempts=%d, want 3", deadlockAttempts) + } + if len(delays) != 2 || delays[0] != schemaDeadlockRetryDelays[0] || delays[1] != schemaDeadlockRetryDelays[1] { + t.Fatalf("schema retry delays=%v", delays) + } + + nonRetryable := errors.New("permission denied") + nonRetryableAttempts := 0 + err = applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) { + nonRetryableAttempts++ + return nil, nonRetryable + }), func(time.Duration) { + t.Fatal("non-deadlock error must not be retried") + }) + if !errors.Is(err, nonRetryable) || nonRetryableAttempts != 1 { + t.Fatalf("non-retryable result: attempts=%d err=%v", nonRetryableAttempts, err) + } +} + // testDSN returns the configured DSN, skipping the test when neither the env var // nor a config file supplies one (DSN no longer has a built-in default). func testDSN(t *testing.T) string { diff --git a/db/llm_usage.go b/db/llm_usage.go index 21b3b8c..25d3abe 100644 --- a/db/llm_usage.go +++ b/db/llm_usage.go @@ -49,8 +49,18 @@ CREATE INDEX IF NOT EXISTS idx_llm_usage_exp ON llm_usage(exploration_id); // EnsureLLMUsageTable creates the llm_usage metering table if it does not exist. func (d *DB) EnsureLLMUsageTable() error { - _, err := d.Exec(llmUsageSchema) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return err + } + if _, err := tx.Exec(llmUsageSchema); err != nil { + return err + } + return tx.Commit() } // InsertLLMUsage appends one metering row. Best-effort: callers log and continue on diff --git a/db/schema.sql b/db/schema.sql index 2f7c22f..88deae6 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -452,7 +452,7 @@ CREATE TABLE IF NOT EXISTS task_archives ( progress INTEGER NOT NULL DEFAULT 0 CHECK (progress BETWEEN 0 AND 100), error TEXT NOT NULL DEFAULT '', warnings JSONB NOT NULL DEFAULT '[]', - format_version INTEGER NOT NULL DEFAULT 1, + format_version INTEGER NOT NULL DEFAULT 2, archive_path TEXT NOT NULL DEFAULT '', sha256 TEXT NOT NULL DEFAULT '', original_size BIGINT NOT NULL DEFAULT 0, @@ -472,6 +472,7 @@ CREATE TABLE IF NOT EXISTS task_archives ( created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +ALTER TABLE task_archives ALTER COLUMN format_version SET DEFAULT 2; CREATE INDEX IF NOT EXISTS idx_task_archives_state ON task_archives(state, requested_at, id); CREATE INDEX IF NOT EXISTS idx_task_archives_archived ON task_archives(archived_at DESC, id DESC); CREATE INDEX IF NOT EXISTS idx_task_archives_sources ON task_archives USING GIN(source_task_ids); diff --git a/db/task_archives.go b/db/task_archives.go index c21c89d..e38a4cd 100644 --- a/db/task_archives.go +++ b/db/task_archives.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/url" "sort" "strconv" @@ -14,7 +15,15 @@ import ( "time" ) -const TaskArchiveFormatVersion = 1 +const ( + TaskArchiveFormatVersion = 2 + TaskArchiveLegacyFormatVersion = 1 + TaskArchiveLLMRecordsPath = "database/llm_records.ndjson" +) + +func IsTaskArchiveFormatSupported(version int) bool { + return version >= TaskArchiveLegacyFormatVersion && version <= TaskArchiveFormatVersion +} const ( ArchiveQueued = "archive_queued" @@ -103,8 +112,8 @@ type TaskArchivePage struct { } // TaskArchiveSnapshot is serialized into manifest.json inside the cold package. -// Each value in Tables is a JSON array generated by PostgreSQL row_to_json, which -// preserves exact IDs, timestamps, arrays and JSON payloads for restoration. +// Small tables remain JSON arrays in Tables. Large v2 tables are streamed to +// package files listed in StreamedTables so their size is not bounded by memory. type TaskArchiveSnapshot struct { FormatVersion int `json:"format_version"` CreatedAt time.Time `json:"created_at"` @@ -115,6 +124,7 @@ type TaskArchiveSnapshot struct { ExclusiveHosts []string `json:"exclusive_hosts"` ExclusiveAssetIDs []int64 `json:"exclusive_asset_ids"` Tables map[string]json.RawMessage `json:"tables"` + StreamedTables map[string]string `json:"streamed_tables,omitempty"` DataCounts map[string]int64 `json:"data_counts"` AggregateStats map[string]any `json:"aggregate_stats"` } @@ -279,12 +289,13 @@ LIMIT 1`, taskID).Scan(&dependent) task_description,task_goal,original_status,category_id_snapshot, category_name_snapshot,source_task_ids,remaining_timeout_seconds,requested_at) VALUES ($1,$2,'queued',0,'','[]',$3,$4,$5,$6,$7,$8,$9,$10,$11,now()) -ON CONFLICT (task_id) DO UPDATE SET - state=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.state ELSE task_archives.state END, - phase=CASE WHEN task_archives.state IN ('archive_failed') THEN 'queued' ELSE task_archives.phase END, - progress=CASE WHEN task_archives.state IN ('archive_failed') THEN 0 ELSE task_archives.progress END, - error=CASE WHEN task_archives.state IN ('archive_failed') THEN '' ELSE task_archives.error END, - requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`, + ON CONFLICT (task_id) DO UPDATE SET + state=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.state ELSE task_archives.state END, + phase=CASE WHEN task_archives.state IN ('archive_failed') THEN 'queued' ELSE task_archives.phase END, + progress=CASE WHEN task_archives.state IN ('archive_failed') THEN 0 ELSE task_archives.progress END, + error=CASE WHEN task_archives.state IN ('archive_failed') THEN '' ELSE task_archives.error END, + format_version=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.format_version ELSE task_archives.format_version END, + requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`, taskID, ArchiveQueued, TaskArchiveFormatVersion, name, description, goal, status, categoryID, categoryName, sources, remaining) if err != nil { @@ -341,15 +352,17 @@ WHERE id=$1 RETURNING `+taskArchiveCols, id, DeleteQueued)) return item, tx.Commit() } -// RecoverTaskArchiveJobs makes active operations retryable after an unclean -// shutdown. Idempotent restore/import semantics keep already-completed pieces safe. +// RecoverTaskArchiveJobs keeps restore/delete resumable after an unclean shutdown. +// An interrupted archive requires an explicit retry: automatic startup retries can +// otherwise form a crash loop when the prior process was killed by resource limits. func (d *DB) RecoverTaskArchiveJobs() error { _, err := d.Exec(`UPDATE task_archives SET - state=CASE state WHEN 'archiving' THEN 'archive_queued' - WHEN 'restoring' THEN 'restore_queued' - WHEN 'deleting' THEN 'delete_queued' ELSE state END, - phase='recovered', error='' -WHERE state IN ('archiving','restoring','deleting')`) + state=CASE state WHEN 'archiving' THEN 'archive_failed' + WHEN 'restoring' THEN 'restore_queued' + WHEN 'deleting' THEN 'delete_queued' ELSE state END, + phase='interrupted', + error=CASE WHEN state='archiving' THEN '上次归档进程异常退出，请手动重试' ELSE '' END + WHERE state IN ('archiving','restoring','deleting')`) return err } @@ -435,7 +448,14 @@ func queryArchiveRows(q interface { return nil, 0, err } defer rows.Close() + return encodeArchiveRows(rows) +} +func encodeArchiveRows(rows interface { + Next() bool + Scan(dest ...any) error + Err() error +}) (json.RawMessage, int64, error) { var output bytes.Buffer output.WriteByte('[') var count int64 @@ -457,6 +477,46 @@ func queryArchiveRows(q interface { return json.RawMessage(output.Bytes()), count, nil } +func writeArchiveRows(rows interface { + Next() bool + Scan(dest ...any) error + Err() error +}, writer io.Writer) (int64, error) { + var count int64 + for rows.Next() { + var raw []byte + if err := rows.Scan(&raw); err != nil { + return 0, err + } + if written, err := writer.Write(raw); err != nil { + return 0, err + } else if written != len(raw) { + return 0, io.ErrShortWrite + } + if written, err := io.WriteString(writer, "\n"); err != nil { + return 0, err + } else if written != 1 { + return 0, io.ErrShortWrite + } + count++ + } + if err := rows.Err(); err != nil { + return 0, err + } + return count, nil +} + +func streamArchiveRows(q interface { + Query(query string, args ...any) (*sql.Rows, error) +}, writer io.Writer, inner string, args ...any) (int64, error) { + rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...) + if err != nil { + return 0, err + } + defer rows.Close() + return writeArchiveRows(rows, writer) +} + func rawRowCount(raw json.RawMessage) int64 { var rows []json.RawMessage if json.Unmarshal(raw, &rows) != nil { @@ -480,11 +540,28 @@ UNION SELECT value::bigint FROM findings finding // writes are already quiescent at the server barrier; repeatable-read also keeps // the asset and accounting views mutually consistent during serialization. func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { + return d.snapshotTaskArchive(taskID, nil) +} + +// SnapshotTaskArchiveWithLLMRecords streams the heavyweight record history to +// llmRecords while all other task-owned data is read from the same repeatable +// PostgreSQL snapshot. +func (d *DB) SnapshotTaskArchiveWithLLMRecords(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) { + if llmRecords == nil { + return nil, errors.New("nil LLM record archive writer") + } + return d.snapshotTaskArchive(taskID, llmRecords) +} + +func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) { tx, err := d.BeginTx(context.Background(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true}) if err != nil { return nil, err } defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return nil, err + } var expID int64 if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 AND deleted_at IS NULL`, taskID).Scan(&expID); err != nil { if errors.Is(err, sql.ErrNoRows) { @@ -518,7 +595,18 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { {"assets", `SELECT asset.* FROM assets asset WHERE asset.id IN (` + archiveAssetIDsQuery() + `) ORDER BY asset.id`, []any{taskID, expID}}, } counts := make(map[string]int64, len(queries)) + streamedTables := map[string]string{} for _, query := range queries { + if query.name == "llm_records" && llmRecords != nil { + count, err := streamArchiveRows(tx, llmRecords, query.query, query.args...) + if err != nil { + return nil, fmt.Errorf("snapshot %s: %w", query.name, err) + } + tables[query.name] = json.RawMessage("[]") + counts[query.name] = count + streamedTables[query.name] = TaskArchiveLLMRecordsPath + continue + } raw, count, err := queryArchiveRows(tx, query.query, query.args...) if err != nil { return nil, fmt.Errorf("snapshot %s: %w", query.name, err) @@ -555,7 +643,8 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { snapshot := &TaskArchiveSnapshot{ FormatVersion: TaskArchiveFormatVersion, CreatedAt: time.Now().UTC(), TaskID: taskID, ExplorationID: expID, SourceTaskIDs: sources, Hosts: hosts, ExclusiveHosts: exclusiveHosts, - ExclusiveAssetIDs: exclusiveAssetIDs, Tables: tables, DataCounts: counts, AggregateStats: stats, + ExclusiveAssetIDs: exclusiveAssetIDs, Tables: tables, StreamedTables: streamedTables, + DataCounts: counts, AggregateStats: stats, } return snapshot, tx.Commit() } diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go index d3f15b0..71e916a 100644 --- a/db/task_archives_restore.go +++ b/db/task_archives_restore.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "strconv" "strings" "time" @@ -36,6 +37,9 @@ func (d *DB) CompleteTaskArchive( return err } defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return err + } var taskID, expID int64 var state string if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state @@ -120,10 +124,11 @@ WHERE id=$1`, taskID); err != nil { return err } if _, err := tx.Exec(`UPDATE task_archives SET - state=$2,phase='ready',progress=100,error='',archive_path=$3,sha256=$4, - original_size=$5,compressed_size=$6,data_counts=$7,aggregate_stats=$8, - archived_at=now(),warnings='[]' -WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize, string(countsRaw), string(statsRaw)); err != nil { + state=$2,phase='ready',progress=100,error='',archive_path=$3,sha256=$4, + original_size=$5,compressed_size=$6,data_counts=$7,aggregate_stats=$8, + format_version=$9,archived_at=now(),warnings='[]' +WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize, + string(countsRaw), string(statsRaw), snapshot.FormatVersion); err != nil { return err } return tx.Commit() @@ -133,14 +138,49 @@ WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compre // idempotent for accounting/traffic retry scenarios and returns non-fatal // warnings for global objects that intentionally are not recreated. func (d *DB) RestoreTaskArchive(archiveID int64, snapshot *TaskArchiveSnapshot, remainingTimeoutSeconds int64) ([]string, error) { - if snapshot == nil || snapshot.FormatVersion != TaskArchiveFormatVersion { + return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, nil) +} + +// RestoreTaskArchiveWithLLMRecords restores a v2 package whose heavyweight LLM +// record history is stored as a sequence of JSON objects outside manifest.json. +func (d *DB) RestoreTaskArchiveWithLLMRecords( + archiveID int64, + snapshot *TaskArchiveSnapshot, + remainingTimeoutSeconds int64, + llmRecords io.Reader, +) ([]string, error) { + if llmRecords == nil { + return nil, errors.New("nil streamed LLM record reader") + } + return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, llmRecords) +} + +func (d *DB) restoreTaskArchive( + archiveID int64, + snapshot *TaskArchiveSnapshot, + remainingTimeoutSeconds int64, + llmRecords io.Reader, +) ([]string, error) { + if snapshot == nil || !IsTaskArchiveFormatSupported(snapshot.FormatVersion) { return nil, ErrTaskArchiveFormatMismatch } + streamedLLMRecords := snapshot.StreamedTables["llm_records"] + if (len(snapshot.StreamedTables) > 0 && snapshot.FormatVersion < 2) || + len(snapshot.StreamedTables) > 1 || + (len(snapshot.StreamedTables) == 1 && streamedLLMRecords != TaskArchiveLLMRecordsPath) { + return nil, fmt.Errorf("%w: unsupported streamed table metadata", ErrTaskArchiveFormatMismatch) + } + if streamedLLMRecords != "" && llmRecords == nil { + return nil, fmt.Errorf("%w: streamed LLM records are missing", ErrTaskArchiveFormatMismatch) + } tx, err := d.Begin() if err != nil { return nil, err } defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return nil, err + } var taskID, expID int64 var state string if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state @@ -215,7 +255,23 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID } else { warnings = append(warnings, warning...) } - for _, table := range []string{"task_asset_links", "findings", "llm_records", "llm_usage", "skill_usage", "tool_usage"} { + for _, table := range []string{"task_asset_links", "findings"} { + if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil { + return nil, fmt.Errorf("restore %s: %w", table, err) + } + } + if streamedLLMRecords != "" { + count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords) + if err != nil { + return nil, fmt.Errorf("restore llm_records: %w", err) + } + if expected := snapshot.DataCounts["llm_records"]; count != expected { + return nil, fmt.Errorf("restore llm_records: row count %d does not match manifest %d", count, expected) + } + } else if err := insertArchiveRows(tx, "llm_records", remappedTables["llm_records"]); err != nil { + return nil, fmt.Errorf("restore llm_records: %w", err) + } + for _, table := range []string{"llm_usage", "skill_usage", "tool_usage"} { if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil { return nil, fmt.Errorf("restore %s: %w", table, err) } @@ -306,6 +362,36 @@ func insertArchiveRows(tx *sql.Tx, table string, raw json.RawMessage) error { return err } +func insertArchiveJSONSequenceRows(tx *sql.Tx, table string, reader io.Reader) (int64, error) { + if table != "llm_records" { + return 0, fmt.Errorf("archive streamed restore table %q is not allowed", table) + } + statement, err := tx.Prepare(`INSERT INTO llm_records SELECT * FROM json_populate_record(NULL::llm_records,$1::json)`) + if err != nil { + return 0, err + } + defer statement.Close() + decoder := json.NewDecoder(reader) + var count int64 + for { + var raw json.RawMessage + if err := decoder.Decode(&raw); errors.Is(err, io.EOF) { + break + } else if err != nil { + return 0, err + } + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || trimmed[0] != '{' { + return 0, errors.New("streamed archive row must be a JSON object") + } + if _, err := statement.Exec(string(trimmed)); err != nil { + return 0, err + } + count++ + } + return count, nil +} + func restoreArchiveAssets(tx *sql.Tx, taskID int64, raw json.RawMessage) (map[int64]int64, []string, error) { rows, err := decodeArchiveRows(raw) if err != nil { @@ -598,6 +684,9 @@ func (d *DB) DeleteTaskArchiveStub(archiveID int64) error { return err } defer tx.Rollback() //nolint:errcheck + if err := coordinateWithSchemaMigration(tx); err != nil { + return err + } var taskID, expID int64 var state string if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state diff --git a/db/task_archives_test.go b/db/task_archives_test.go index c9a6d18..498fa97 100644 --- a/db/task_archives_test.go +++ b/db/task_archives_test.go @@ -1,24 +1,55 @@ package db import ( + "bytes" "encoding/json" + "errors" "fmt" "strings" "testing" "time" ) -func TestQueryArchiveRowsStreamsJSON(t *testing.T) { - d, err := Open(testDSN(t)) - if err != nil { - t.Skipf("postgres unavailable (%v) — skipping", err) - } - defer d.Close() +type archiveJSONTestRows struct { + values [][]byte + next int +} +func (r *archiveJSONTestRows) Next() bool { + if r.next >= len(r.values) { + return false + } + r.next++ + return true +} + +func (r *archiveJSONTestRows) Scan(dest ...any) error { + if len(dest) != 1 || r.next == 0 || r.next > len(r.values) { + return fmt.Errorf("invalid archive test row scan") + } + target, ok := dest[0].(*[]byte) + if !ok { + return fmt.Errorf("archive test row destination is %T", dest[0]) + } + *target = append((*target)[:0], r.values[r.next-1]...) + return nil +} + +func (r *archiveJSONTestRows) Err() error { return nil } + +func TestQueryArchiveRowsStreamsJSON(t *testing.T) { payload := strings.Repeat("large request/response payload ", 64*1024) - raw, count, err := queryArchiveRows(d, `SELECT * FROM (VALUES -($1::bigint,$2::text),($3::bigint,$4::text)) archived_row(id,body) ORDER BY id`, - int64(1), payload, int64(2), "quoted: \"value\"\nline") + values := make([][]byte, 2) + var err error + values[0], err = json.Marshal(map[string]any{"id": int64(1), "body": payload}) + if err != nil { + t.Fatal(err) + } + values[1], err = json.Marshal(map[string]any{"id": int64(2), "body": "quoted: \"value\"\nline"}) + if err != nil { + t.Fatal(err) + } + raw, count, err := encodeArchiveRows(&archiveJSONTestRows{values: values}) if err != nil { t.Fatal(err) } @@ -40,6 +71,55 @@ func TestQueryArchiveRowsStreamsJSON(t *testing.T) { } } +func TestWriteArchiveRowsProducesJSONSequence(t *testing.T) { + values := [][]byte{ + json.RawMessage(`{"id":1,"body":"first"}`), + json.RawMessage(`{"id":2,"body":"second\\nline"}`), + } + var output bytes.Buffer + count, err := writeArchiveRows(&archiveJSONTestRows{values: values}, &output) + if err != nil { + t.Fatal(err) + } + if count != 2 { + t.Fatalf("streamed row count=%d, want 2", count) + } + decoder := json.NewDecoder(&output) + for wantID := int64(1); wantID <= 2; wantID++ { + var row struct { + ID int64 `json:"id"` + } + if err := decoder.Decode(&row); err != nil { + t.Fatal(err) + } + if row.ID != wantID { + t.Fatalf("streamed row id=%d, want %d", row.ID, wantID) + } + } +} + +func TestTaskArchiveFormatCompatibility(t *testing.T) { + for _, version := range []int{TaskArchiveLegacyFormatVersion, TaskArchiveFormatVersion} { + if !IsTaskArchiveFormatSupported(version) { + t.Fatalf("archive format %d should be supported", version) + } + } + for _, version := range []int{0, TaskArchiveFormatVersion + 1} { + if IsTaskArchiveFormatSupported(version) { + t.Fatalf("archive format %d should be rejected", version) + } + } + invalidSnapshots := []*TaskArchiveSnapshot{ + {FormatVersion: TaskArchiveLegacyFormatVersion, StreamedTables: map[string]string{"llm_records": TaskArchiveLLMRecordsPath}}, + {FormatVersion: TaskArchiveFormatVersion, StreamedTables: map[string]string{"unknown": "database/unknown.ndjson"}}, + } + for _, snapshot := range invalidSnapshots { + if _, err := (&DB{}).RestoreTaskArchive(1, snapshot, 0); !errors.Is(err, ErrTaskArchiveFormatMismatch) { + t.Fatalf("invalid streamed table metadata returned %v", err) + } + } +} + func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { @@ -111,6 +191,12 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro if err := d.InsertToolUsage(&ToolUsage{ToolKey: toolName, AgentKey: "worker", TaskID: task.ID, ExplorationID: task.ExplorationID}); err != nil { t.Fatal(err) } + if err := d.InsertLLMRecord(&LLMRecord{ + TaskID: fmt.Sprint(task.ID), Model: "archive-model", SessionID: "archive-session", + Status: "ok", RawRequest: strings.Repeat("request", 1024), RawResponse: strings.Repeat("response", 1024), + }); err != nil { + t.Fatal(err) + } if _, err := d.AddFinding(task.ID, nodeID, vulnclass, "archive finding", SeverityCritical, "summary", "evidence", "worker", []int64{assetID}); err != nil { t.Fatal(err) } @@ -122,10 +208,17 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro if err != nil || claimed == nil || claimed.ID != archive.ID || claimed.State != Archiving { t.Fatalf("claim = %+v, %v", claimed, err) } - snapshot, err := d.SnapshotTaskArchive(task.ID) + var llmRecords bytes.Buffer + snapshot, err := d.SnapshotTaskArchiveWithLLMRecords(task.ID, &llmRecords) if err != nil { t.Fatal(err) } + if snapshot.StreamedTables["llm_records"] != TaskArchiveLLMRecordsPath || snapshot.DataCounts["llm_records"] != 1 { + t.Fatalf("unexpected streamed LLM metadata: paths=%v counts=%v", snapshot.StreamedTables, snapshot.DataCounts) + } + if rawRowCount(snapshot.Tables["llm_records"]) != 0 { + t.Fatal("streamed LLM records were also retained in manifest memory") + } if snapshot.DataCounts["assets"] != 1 || snapshot.DataCounts["exploration_nodes"] < 2 { t.Fatalf("unexpected snapshot counts: %#v", snapshot.DataCounts) } @@ -170,7 +263,9 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro if err != nil || claimed == nil || claimed.State != Restoring { t.Fatalf("restore claim = %+v, %v", claimed, err) } - warnings, err := d.RestoreTaskArchive(archive.ID, snapshot, int64((10 * time.Minute).Seconds())) + warnings, err := d.RestoreTaskArchiveWithLLMRecords( + archive.ID, snapshot, int64((10 * time.Minute).Seconds()), bytes.NewReader(llmRecords.Bytes()), + ) if err != nil { t.Fatal(err) } @@ -202,6 +297,14 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro if nodes < 2 || assets != 1 || usage != 1 { t.Fatalf("restore incomplete nodes=%d assets=%d usage=%d", nodes, assets, usage) } + var restoredRawRequest, restoredRawResponse string + if err := d.QueryRow(`SELECT COALESCE(raw_request,''),COALESCE(raw_response,'') FROM llm_records WHERE task_id=$1`, fmt.Sprint(task.ID)). + Scan(&restoredRawRequest, &restoredRawResponse); err != nil { + t.Fatal(err) + } + if restoredRawRequest != strings.Repeat("request", 1024) || restoredRawResponse != strings.Repeat("response", 1024) { + t.Fatal("restored streamed LLM record body was truncated") + } var restoredCompanyID *int64 var restoredStatus, restoredError string var restoredExhaustedAt, restoredCreatedAt time.Time @@ -274,6 +377,42 @@ func TestTaskArchiveBlockersIgnoreQueuedDependents(t *testing.T) { } } +func TestRecoverInterruptedArchiveRequiresManualRetry(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + task, err := d.CreateTaskWithOptions("interrupted archive", "must not restart automatically", TaskCreateOptions{}) + if err != nil { + t.Fatal(err) + } + defer func() { + _, _ = d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID) + _ = d.DeleteTask(task.ID) + }() + if err := d.SetPaused(task.ID, true); err != nil { + t.Fatal(err) + } + queued, err := d.QueueTaskArchive(task.ID) + if err != nil { + t.Fatal(err) + } + if _, err := d.Exec(`UPDATE task_archives SET state=$2,phase='snapshot_database',progress=10 WHERE id=$1`, queued.ID, Archiving); err != nil { + t.Fatal(err) + } + if err := d.RecoverTaskArchiveJobs(); err != nil { + t.Fatal(err) + } + recovered, err := d.GetTaskArchive(queued.ID) + if err != nil { + t.Fatal(err) + } + if recovered == nil || recovered.State != ArchiveFailed || recovered.Phase != "interrupted" || recovered.Error == "" { + t.Fatalf("recovered archive=%+v, want explicit-retry failure", recovered) + } +} + func assertArchiveGlobalStats(t *testing.T, d *DB, profileName, skillName, toolName, vulnclass string) { t.Helper() profiles, err := d.UsageByProfile() diff --git a/server/task_archive_package.go b/server/task_archive_package.go index ce4494d..bc55617 100644 --- a/server/task_archive_package.go +++ b/server/task_archive_package.go @@ -480,11 +480,20 @@ func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchive if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil { return 0, 0, "", err } - manifest, err := json.Marshal(snapshot) + manifestFile, err := os.OpenFile( + filepath.Join(payloadDir, "manifest.json"), + os.O_CREATE|os.O_TRUNC|os.O_WRONLY, + archiveFileMode, + ) if err != nil { return 0, 0, "", err } - if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil { + encodeErr := json.NewEncoder(manifestFile).Encode(snapshot) + var syncErr error + if encodeErr == nil { + syncErr = manifestFile.Sync() + } + if err := errors.Join(encodeErr, syncErr, manifestFile.Close()); err != nil { return 0, 0, "", err } temporary := path + ".partial" diff --git a/server/task_archive_package_test.go b/server/task_archive_package_test.go index 74e86ad..9846673 100644 --- a/server/task_archive_package_test.go +++ b/server/task_archive_package_test.go @@ -47,7 +47,18 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) { } archivePath := taskArchivePath(dataDir, 1, taskID) - snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID} + streamPath := filepath.Join(stage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath)) + if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil { + t.Fatal(err) + } + streamBody := []byte("{\"id\":1,\"raw_request\":\"large\"}\n") + if err := os.WriteFile(streamPath, streamBody, archiveFileMode); err != nil { + t.Fatal(err) + } + snapshot := &pgdb.TaskArchiveSnapshot{ + FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID, + StreamedTables: map[string]string{"llm_records": pgdb.TaskArchiveLLMRecordsPath}, + } original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot) if err != nil { t.Fatal(err) @@ -62,6 +73,9 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) { if err := extractTaskArchivePackage(archivePath, checksum, extracted); err != nil { t.Fatal(err) } + if got, err := os.ReadFile(filepath.Join(extracted, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))); err != nil || !bytes.Equal(got, streamBody) { + t.Fatalf("streamed LLM archive payload=%q err=%v", got, err) + } installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1) if err != nil { t.Fatal(err) diff --git a/server/task_archives.go b/server/task_archives.go index c953a1a..30cea69 100644 --- a/server/task_archives.go +++ b/server/task_archives.go @@ -131,13 +131,15 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) { if err := s.waitTaskQuiescent(drainCtx, taskID); err != nil { return errors.New("任务仍有运行中的 Agent，请先暂停后重试归档") } - _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 10) - snapshot, err := s.m.pg.SnapshotTaskArchive(job.TaskID) + task, err := s.m.pg.GetTask(job.TaskID) if err != nil { return err } - _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 25) - fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, snapshot.ExplorationID) + if task == nil { + return pgdb.ErrTaskArchiveNotFound + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 10) + fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, task.ExplorationID) if err != nil { return err } @@ -146,6 +148,27 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) { runErr = errors.Join(runErr, fileStage.rollback()) } }() + streamPath := filepath.Join(fileStage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath)) + if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil { + return err + } + streamFile, err := os.OpenFile(streamPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode) + if err != nil { + return err + } + _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 20) + snapshot, snapshotErr := s.m.pg.SnapshotTaskArchiveWithLLMRecords(job.TaskID, streamFile) + var syncErr error + if snapshotErr == nil { + syncErr = streamFile.Sync() + } + closeErr := streamFile.Close() + if err := errors.Join(snapshotErr, syncErr, closeErr); err != nil { + return err + } + if snapshot.ExplorationID != task.ExplorationID { + return errors.New("任务探索记录在归档快照期间发生变化") + } if s.m.traffic != nil && len(snapshot.Hosts) > 0 { _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_traffic", 38) trafficCount, err := s.m.traffic.ExportHosts(snapshot.Hosts, filepath.Join(fileStage.payload, "traffic")) @@ -244,9 +267,27 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) { if err := json.Unmarshal(manifestRaw, &snapshot); err != nil { return err } - if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion { + if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) { return pgdb.ErrTaskArchiveFormatMismatch } + var llmRecordsFile *os.File + relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"] + if len(snapshot.StreamedTables) > 0 && + (snapshot.FormatVersion < 2 || !hasStreamedLLMRecords || len(snapshot.StreamedTables) != 1 || + relative != pgdb.TaskArchiveLLMRecordsPath) { + return pgdb.ErrTaskArchiveFormatMismatch + } + if hasStreamedLLMRecords { + llmRecordsFile, err = os.Open(filepath.Join(extracted, filepath.FromSlash(relative))) + if err != nil { + return err + } + defer func() { + if llmRecordsFile != nil { + _ = llmRecordsFile.Close() + } + }() + } if s.m.traffic != nil { _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_traffic", 35) if _, err := s.m.traffic.ImportArchive(filepath.Join(extracted, "traffic")); err != nil { @@ -265,7 +306,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) { } }() _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70) - warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds) + var warnings []string + if llmRecordsFile != nil { + warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords( + job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile, + ) + closeErr := llmRecordsFile.Close() + llmRecordsFile = nil + err = errors.Join(err, closeErr) + } else { + warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds) + } if err != nil { return err } From 21656cc2e31000803782bd532836f7b463bd304c Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:37:46 -0400 Subject: [PATCH 7/7] fix: restore changes dropped by the upstream sync merge The manual v0.3.6 sync (22132a4) was incomplete, and the follow-up merge (d803ae2) resolved conflicts in favour of that copy, silently reverting three shipped changes. main is already an ancestor of this branch, so re-merging cannot bring them back; restore them explicitly. - schema.sql: re-add the idempotent llm_profiles_format_check migration. CREATE TABLE IF NOT EXISTS leaves an existing table alone, so databases created before openai-responses kept the old two-value CHECK and rejected every openai-responses profile. - orchestration.go: re-add the reseedPlannerPrompt/reseedWorkerPrompt calls, dropped by a comment-alignment hunk while the functions stayed behind as dead code. Without them existing installs never receive the quantified acceptance check or the negative-conclusion evidence threshold. - CHANGELOG.md: restore the 0.3.6 entries, the @neouks credit and the correct norma version range. Also close three recovery gaps in the new cold-archive path: - traffic: write the staging journal even when no legacy host directory exists. Modern installs keep traffic in SQLite only, so a crash between the PostgreSQL commit and the SQLite commit left the archived task's exclusive traffic in the hot store with nothing for recovery to act on. - recoverTaskArchiveStages: treat a missing journal.json as a discardable directory instead of an error, which used to set recoveryFailed and stop the archive worker from starting for the rest of the process lifetime. - stageTaskArchiveFiles: replay the journal's renames when resuming an existing staging root. The journal is written before the first rename, so its presence did not guarantee a complete payload, and a resumed round could package an archive missing transcripts or workspace files. --- CHANGELOG.md | 39 +++++++++++++++++++- db/schema.sql | 5 +++ server/orchestration.go | 2 ++ server/task_archive_package.go | 65 ++++++++++++++++++++++++++++------ traffic/traffic.go | 9 +++-- 5 files changed, 107 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d96ad57..b6e5616 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,7 +29,7 @@ #### 修改的功能 -- 升级 norma v0.2.9 → v0.3.2：修复 OpenAI 系响应里被静默丢弃的 reasoning/refusal（Responses 补 `reasoning_text`）、reasoning 字段三名去重、空响应（零内容块）有界重试、压缩边界切断工具配对导致的网关 400（同时修复已烤进 transcript 的存量孤儿）。 +- 升级 norma v0.3.1 → v0.3.2：修复 OpenAI 系响应里被静默丢弃的 reasoning/refusal（Responses 补 `reasoning_text`）、reasoning 字段三名去重、空响应（零内容块）有界重试、压缩边界切断工具配对导致的网关 400（同时修复已烤进 transcript 的存量孤儿）。 ### 贡献者 @@ -42,10 +42,47 @@ #### 新增的功能 - 每个 LLM 配置支持流式/非流式切换（默认流式）：开走流式 SSE；关走真·非流式（`stream:false`、一次性返回完整 JSON），可绕开部分网关糟糕的 SSE 实现（空帧、思考字段丢帧），代价是失去运行中的实时进度与实时 Token 计数。worker / planner / mainagent / chat / goals 均按当前激活配置动态取值；`llmpool` / `llmrec` / 任务运行时三层 Provider 包装均兼容非流式；`llm_profiles` 新增 `streaming` 列并 `ALTER` 补旧库（默认 `true`，旧配置无感）。 +- 支持 OpenAI Responses API 格式的 LLM 配置：每个配置新增第三种格式 `openai-responses`（打 `POST /v1/responses`），与 Chat Completions / Anthropic 并列；`BaseURL` 归一化、默认模型 `gpt-5`；`llm_profiles` 的 `format` 约束加入 `openai-responses` 并幂等迁移补旧库；前端格式下拉新增「OpenAI (Responses API)」。依赖升级 norma v0.3.1（含 `reasoning_content` 回传修复）。 +- 录制 LLM 请求/响应 HTTP 原文：在 HTTP transport 层捕获真实 wire body，保留归一化视图看不到的工具 schema、`tool_use` 块与原始 SSE 帧；norma 内部重试的多次尝试逐次保留；`llm_records` 新增 `raw_request` / `raw_response` 两列并 `ALTER` 补旧库；录制页详情面板增加「原文」视图切换与请求/响应复制按钮（兼容非安全上下文的 `execCommand` 回退）。 + +### 对话 + +#### 修改的功能 + +- 会话列表多选改为「多选」模式开关：默认列表不再常驻每行勾选框（观感更干净），表头改为「共 N 个 + 多选」按钮；点「多选」进入选择模式（勾选框、全选、批量删除），「完成」退出并清空选择，批量删除全部成功后自动退出（有失败则留在选择模式便于重试）；单条重命名 / 置顶 / 删除仍走每行 ⋯ 菜单。 + +### UI + +#### 修改的功能 + +- 任务管理、会话操作与流量查看体验增强（#57）：任务列表列排序偏好可持久化记忆、行内重命名改用图标直接触发（不再经菜单）、抽屉（sheet）交互与流量查看细节优化。 + +#### 修复的问题 + +- Worker 资产标签只显示域名与 IP，并正确处理空标签的场景。 + +### Agent + +#### 新增的功能 + +- planner 判目标时新增「量化验收核对」：目标含可量化验收条件（资产测试覆盖度达到 X%、拿到 N 个 flag、获得某权限）时，`prove_goal` 前必须核对 `graph_overview` 的实测值（`coverage.pct`、findings 计数等），实测未达标则禁止 `prove_goal`、改为派意图补足差距，不得以「主要部分已完成/大体达成」为由提前标 met。修复了目标要求覆盖度 100%、实测仅 40% 却被判完成的问题。 + +#### 修改的功能 + +- 重写 planner 提示词中「本轮 0 意图」的判定依据：原文将 0 意图描述为「最常见、最重要的原则」，会让规划者在目标未达成、范围内仍有未测面时过早收手。现改为只在两种具体情况下才应产出 0 意图——① 想到的方向都已被 open/running/recent_done 意图覆盖；② 下一步依赖当前正在运行的 work 的产出、而产出尚未出现（应等其跑完、图更新后的下一次唤醒再规划）。并补充反向约束：确有未被覆盖且不依赖在跑 work 的新方向、或目标未达成且仍有未测面时，不要因「0 意图常见」而收手。 +- 强化 worker 提示词中「否定结论的证据门槛」：对「不可注入/端口关闭/无登录入口」等可能让规划者放弃一整条方向的否定结论，要求下结论前先穷尽该意图内的合理手段（换编码/参数/路径/方法）；手段未走完或证据偏弱一律标 `confidence=inferred`，避免用轻率的 `observed` 否定把整条路线焊死（任务早期的错误否定尤其会把方向带偏且难以自愈）。 +- 上述 planner / worker 默认提示词变更通过一次性迁移（`reseedPlannerPrompt` / `reseedWorkerPrompt`，各带 settings flag 守卫）追加为新版本并切为当前版本；旧版本保留在版本历史中，自定义过提示词的用户可在版本记录里恢复。 + +### 运维 + +#### 新增的功能 + +- 新增 `reset-password.sh` 重置管理员（用户名固定 `ARTEX`）密码：支持 local / docker 两种部署，连接信息可显式指定或自动从 `--dsn`/`$ARTEX_PG_DSN`/`config.json` 读取；在库内用 `pgcrypto` 生成与后端登录兼容的 bcrypt 哈希并写回 `settings.auth.password_hash`，重置后无需重启服务。密码经环境变量传入、不进入进程 argv，并做转义防注入。 ### 贡献者 - [@Autumn-27](https://github.com/Autumn-27) +- [@neouks](https://github.com/neouks) ## [0.3.4] - 2026-08-24 diff --git a/db/schema.sql b/db/schema.sql index 88deae6..fe6924e 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -324,6 +324,11 @@ ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS priority INTEGER NOT NULL ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS pool_exclude BOOLEAN NOT NULL DEFAULT false; -- 流式开关；补旧库。默认 true = 保持既有的流式行为，旧配置无感升级。 ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS streaming BOOLEAN NOT NULL DEFAULT true; +-- 放开 format 约束以容纳 openai-responses(OpenAI Responses API)；补旧库。 +-- 每次启动执行,幂等:先删旧 CHECK 再建含三值的新 CHECK。 +ALTER TABLE llm_profiles DROP CONSTRAINT IF EXISTS llm_profiles_format_check; +ALTER TABLE llm_profiles ADD CONSTRAINT llm_profiles_format_check + CHECK (format IN ('openai','anthropic','openai-responses')); -- 思考开关字段 thinking_type，从旧的单一 reasoning_effort 语义一次性拆分而来。 -- schema.sql 每次启动都执行，故迁移必须只跑一次：仅当该列尚不存在时才回填， diff --git a/server/orchestration.go b/server/orchestration.go index 4288d08..191bde2 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -452,6 +452,8 @@ func (s *Server) seedOrchestrationTools() { s.unbindGoalMetDefault() s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性) s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性) + s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性) + s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性) s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性) // 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把 // list_assets/insert_assets/report_finding/list_findings/list_companies 连同 diff --git a/server/task_archive_package.go b/server/task_archive_package.go index bc55617..a07631f 100644 --- a/server/task_archive_package.go +++ b/server/task_archive_package.go @@ -96,6 +96,17 @@ func stageTaskArchiveFiles(dataDir string, archiveID int64, taskID string, explo if err := json.Unmarshal(raw, &stage.journal); err != nil { return nil, fmt.Errorf("read task archive staging journal: %w", err) } + // The journal is written before the first rename, so its presence does not mean + // the payload is complete: a previous round may have failed mid-loop with a + // rollback that itself errored, which leaves the root in place. Replay the moves + // rather than packaging a payload that is missing transcripts or workspace files. + if err := os.MkdirAll(filepath.Join(stage.payload, "files"), archiveDirMode); err != nil { + return nil, err + } + if err := stage.applyMoves(); err != nil { + _ = stage.rollback() + return nil, err + } return stage, nil } else if !os.IsNotExist(err) { return nil, err @@ -133,20 +144,42 @@ func stageTaskArchiveFiles(dataDir string, archiveID int64, taskID string, explo _ = stage.rollback() return nil, err } - for _, move := range targets { - destination := filepath.Join(stage.payload, move.Relative) - if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil { - _ = stage.rollback() - return nil, err - } - if err := os.Rename(move.Source, destination); err != nil { - _ = stage.rollback() - return nil, fmt.Errorf("stage task archive path %s: %w", move.Source, err) - } + if err := stage.applyMoves(); err != nil { + _ = stage.rollback() + return nil, err } return stage, nil } +// applyMoves performs the renames recorded in the journal. Entries already sitting +// in the payload are skipped, so an interrupted staging round can be resumed in +// place without moving anything twice. +func (s *taskArchiveFileStage) applyMoves() error { + for _, move := range s.journal.Moves { + destination := filepath.Join(s.payload, move.Relative) + if _, err := os.Lstat(destination); err == nil { + continue + } else if !os.IsNotExist(err) { + return err + } + if _, err := os.Lstat(move.Source); err != nil { + if os.IsNotExist(err) { + // Neither side exists: the path was removed outside the archive flow + // after the journal was written. Nothing can be staged for it. + continue + } + return err + } + if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil { + return err + } + if err := os.Rename(move.Source, destination); err != nil { + return fmt.Errorf("stage task archive path %s: %w", move.Source, err) + } + } + return nil +} + func (s *taskArchiveFileStage) rollback() error { if s == nil || s.done { return nil @@ -368,6 +401,18 @@ func recoverTaskArchiveStages(dataDir string, pg *pgdb.DB) error { } root := filepath.Join(parent, entry.Name()) raw, err := os.ReadFile(filepath.Join(root, "journal.json")) + if os.IsNotExist(err) { + // Either staging died between creating the directory tree and writing the + // journal — no rename had run, so payload holds nothing — or commit/rollback + // failed part-way through removing the root, in which case payload only holds + // copies already inside the package. Without a journal there is nothing to + // roll back, and reporting an error here keeps the archive worker from ever + // starting, so discard the directory instead. + if err := os.RemoveAll(root); err != nil { + errs = append(errs, err) + } + continue + } if err != nil { errs = append(errs, err) continue diff --git a/traffic/traffic.go b/traffic/traffic.go index 6dfbce3..d7a2b7f 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -911,7 +911,12 @@ func (t *Traffic) stageTreesForArchive(dirs, hosts []string, archiveID, taskID i } planned = append(planned, stagedTrafficPath{source: source}) } - if len(planned) == 0 { + // 归档路径即使一个历史 host 目录都没有(新装机的流量只落在 SQLite + _blobs) + // 也必须留下 journal：崩溃点若落在 PostgreSQL 提交与 SQLite 提交之间，重启后 + // SQLite 事务被回滚，只有这份 journal 能让恢复流程补做 host 行的删除。少了它， + // 已转冷任务的独占流量会永久留在热库里。 + uniqueHosts := uniqueArchiveHosts(hosts) + if len(planned) == 0 && len(uniqueHosts) == 0 { return "", nil, nil } parent := filepath.Join(t.dir, "_delete_staging") @@ -921,7 +926,7 @@ func (t *Traffic) stageTreesForArchive(dirs, hosts []string, archiveID, taskID i if stageDir, err = os.MkdirTemp(parent, "hosts-"); err != nil { return "", nil, fmt.Errorf("创建流量暂存目录: %w", err) } - journal := hostDeleteStageJournal{Version: 1, ArchiveID: archiveID, TaskID: taskID, Hosts: uniqueArchiveHosts(hosts)} + journal := hostDeleteStageJournal{Version: 1, ArchiveID: archiveID, TaskID: taskID, Hosts: uniqueHosts} for i := range planned { planned[i].staged = filepath.Join(stageDir, fmt.Sprintf("%d-%s", i, filepath.Base(planned[i].source))) journal.Moves = append(journal.Moves, hostDeleteStageMove{Source: planned[i].source, Staged: planned[i].staged})