admin/ARTEX
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

归档

Browse Source
This commit is contained in:
pphua committed 2026-08-28 12:56:29 +08:00
1 parent 594436fddd
commit 9da521ee4e
31 files changed
+5768 -267
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

db/findings.go
+47 -2
View File
@@ -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
db/findings_test.go
+36
View File
@@ -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 {
db/llm_usage.go
+80 -3
View File
@@ -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
db/schema.sql
+46
View File
@@ -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.
db/skill_usage.go
+21 -4
View File
@@ -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).
db/task_archive_aggregate_stats.go
+91
View File
@@ -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
}
db/task_archives.go
+667
View File
@@ -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
}
}
db/task_archives_restore.go
+648
View File
@@ -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()
}
db/task_archives_test.go
+297
View File
@@ -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)
}
}
db/tool_usage.go
+13 -1
View File
@@ -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
}