admin/ARTEX-chixuweihu
1
0 Fork 0
mirror of https://github.com/RuoJi6/ARTEX.git synced 2026-10-27 17:03:49 +08:00
Code Issues 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
}
server/dto.go
+14
View File
@@ -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"`
server/findings_groups.go
+1
View File
@@ -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)
server/findings_groups_test.go
+22
View File
@@ -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()
server/manager.go
+19
View File
@@ -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())
server/server.go
+24 -3
View File
@@ -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:
server/task_archive_package.go
+655
View File
@@ -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
}
server/task_archive_package_test.go
+209
View File
@@ -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)
}
}
server/task_archives.go
+656
View File
@@ -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())
}
server/task_archives_test.go
+88
View File
@@ -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())
}
}
server/task_llm_test.go
+17
View File
@@ -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"}
traffic/archive.go
+237
View File
@@ -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
}
traffic/archive_test.go
+158
View File
@@ -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()
}
traffic/traffic.go
+222 -12
View File
@@ -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 {
web/src/app/(main)/function/findings/page.tsx
+80 -5
View File
@@ -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<string>("all");
const [task, setTask] = React.useState<string>("all");
const [sort, setSort] = React.useState<"severity" | "time">("severity");
const [search, setSearch] = React.useState("");
const [query, setQuery] = React.useState("");
const [expanded, setExpanded] = React.useState<string | null>(null);
const [groups, setGroups] = React.useState<FindingGroup[]>([]);
const [groupTotal, setGroupTotal] = React.useState(0);
@@ -132,15 +138,54 @@ export default function FindingsPage() {
const [groupFindings, setGroupFindings] = React.useState<Record<string, GroupFindingsState>>({});
const [total, setTotal] = React.useState(0);
const [stats, setStats] = React.useState<FindingStats>(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<Finding | null>(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<string, GroupFindingsState> = {};
@@ -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() {
</div>
<div className="flex flex-wrap items-center gap-2">
<InputGroup className="w-full sm:w-72">
<InputGroupInput
type="search"
value={search}
onChange={(event) => setSearch(event.target.value)}
placeholder="检索漏洞内容"
aria-label="检索漏洞内容"
/>
<InputGroupAddon>
<SearchIcon aria-hidden="true" />
</InputGroupAddon>
</InputGroup>
<ToggleGroup
type="single"
value={severity}
web/src/app/(main)/function/tasks/detail/page.tsx
+70 -2
View File
@@ -3,14 +3,25 @@
import * as React from "react";
import Link from "next/link";
import { useSearchParams } from "next/navigation";
import { useRouter, useSearchParams } from "next/navigation";
import { ArrowLeftIcon, BrainIcon, CheckIcon, CircleAlertIcon, PauseIcon, PlayIcon } from "lucide-react";
import { ArchiveIcon, 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 {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
AlertDialogTrigger,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
Popover,
@@ -24,6 +35,7 @@ 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 { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { api } from "@/lib/api";
import type { LLMProfile, Task } from "@/lib/types";
@@ -190,6 +202,7 @@ function TaskLLMControl({ task, profiles, onUpdated }: { task: Task; profiles: L
}
function TaskDetailInner() {
const router = useRouter();
const searchParams = useSearchParams();
const id = searchParams.get("id") ?? "";
const [task, setTask] = React.useState<Task | null>(null);
@@ -198,6 +211,7 @@ function TaskDetailInner() {
const [tab, setTab] = React.useState("sessions");
const [interceptPendingCount, setInterceptPendingCount] = React.useState(0);
const [profiles, setProfiles] = React.useState<LLMProfile[]>([]);
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 (
<div className="flex flex-1 flex-col items-center justify-center gap-3 p-10 text-center">
@@ -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 = <PauseIcon data-icon="inline-start" />;
@@ -302,6 +335,16 @@ function TaskDetailInner() {
controlIcon = <PlayIcon data-icon="inline-start" />;
controlLabel = "恢复";
}
const archiveTrigger = (
<Button
size="icon-sm"
variant="ghost"
disabled={!canArchive || archiving}
aria-label={canArchive ? "归档任务" : archiveDisabledReason}
>
{archiving ? <Spinner /> : <ArchiveIcon />}
</Button>
);
return (
<Tabs value={tab} onValueChange={setTab} className="flex flex-1 flex-col gap-0">
@@ -323,6 +366,31 @@ function TaskDetailInner() {
<Separator orientation="vertical" className="mx-1 hidden h-4 sm:block" />
<TaskLLMControl task={task} profiles={profiles} onUpdated={load} />
<StatusBadge domain={terminal ? "task" : "engine"} value={terminal ? task.status : engineMode} dot />
{canArchive ? (
<AlertDialog>
<AlertDialogTrigger asChild>{archiveTrigger}</AlertDialogTrigger>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>归档任务 #{task.id}？</AlertDialogTitle>
<AlertDialogDescription>
任务图谱、关联记录、独占资产与流量、工作文件和 LLM
历史将压缩到冷存储。归档完成后可在任务列表的“已归档”页还原。
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>取消</AlertDialogCancel>
<AlertDialogAction onClick={() => void archiveTask()}>确认归档</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
) : (
<Tooltip>
<TooltipTrigger asChild>
<span className="inline-flex">{archiveTrigger}</span>
</TooltipTrigger>
<TooltipContent>{archiveDisabledReason}</TooltipContent>
</Tooltip>
)}
<Button size="sm" variant={controlVariant} onClick={togglePause} disabled={terminal}>
{controlIcon}
{controlLabel}
web/src/app/(main)/function/tasks/page.tsx
+869 -220
View File
File diff suppressed because it is too large. Load diff
web/src/app/(main)/function/traffic/page.tsx
+45 -12
View File
@@ -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<TrafficHost[]>([]); // target picker
const [selectedHosts, setSelectedHosts] = React.useState<string[]>([]); // checked in picker
const [pickerOpen, setPickerOpen] = React.useState(false);
const [hostCountSortDirection, setHostCountSortDirection] = React.useState<HostCountSortDirection>("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() {
</Button>
</PopoverTrigger>
<PopoverContent
className="w-80 p-0 data-open:animate-none data-closed:animate-none"
className="w-[calc(100vw-2rem)] p-0 data-open:animate-none data-closed:animate-none sm:w-80"
align="start"
collisionPadding={16}
>
<div className="flex items-center justify-between border-b px-3 py-2">
<span className="text-xs font-medium text-muted-foreground">按目标批量删除</span>
{hosts.length > 0 && (
<Button
variant="ghost"
size="sm"
className="h-6 px-2 text-xs"
onClick={() => setSelectedHosts(allSelected ? [] : hosts.map((h) => h.host))}
>
{allSelected ? "取消全选" : "全选"}
</Button>
)}
<div className="flex items-center gap-1">
{hosts.length > 0 && (
<Tooltip>
<TooltipTrigger asChild>
<Button
variant="ghost"
size="icon-xs"
onClick={() => setHostCountSortDirection((current) => (current === "desc" ? "asc" : "desc"))}
aria-label={
hostCountSortDirection === "desc"
? "数据包数量当前倒序，点击切换为正序"
: "数据包数量当前正序，点击切换为倒序"
}
>
{hostCountSortDirection === "desc" ? <ArrowDownWideNarrowIcon /> : <ArrowUpNarrowWideIcon />}
</Button>
</TooltipTrigger>
<TooltipContent>按数据包数量{hostCountSortDirection === "desc" ? "倒序" : "正序"}</TooltipContent>
</Tooltip>
)}
{hosts.length > 0 && (
<Button
variant="ghost"
size="sm"
className="h-6 px-2 text-xs"
onClick={() => setSelectedHosts(allSelected ? [] : hosts.map((h) => h.host))}
>
{allSelected ? "取消全选" : "全选"}
</Button>
)}
</div>
</div>
<div className="max-h-64 overflow-y-auto">
{hosts.length === 0 ? (
<div className="px-3 py-6 text-center text-xs text-muted-foreground">暂无流量记录</div>
) : (
hosts.map((h, index) => (
sortedHosts.map((h, index) => (
<label
key={h.host}
htmlFor={`traffic-host-${index}`}
web/src/components/transcript.tsx
+24 -1
View File
@@ -129,6 +129,27 @@ function groupSteps(steps: Activity[], chat: boolean): Group[] {
const kindLabel = (k: string) => (k === "thinking" ? "推理" : k === "result" ? "总结" : "说明");
function ActivityTime({ ts }: { ts: string }) {
const date = new Date(ts);
if (!ts || Number.isNaN(date.getTime())) return null;
return (
<time
dateTime={date.toISOString()}
title={date.toLocaleString("zh-CN")}
className="text-[10px] text-muted-foreground tabular-nums"
suppressHydrationWarning
>
{date.toLocaleString("zh-CN", {
month: "2-digit",
day: "2-digit",
hour: "2-digit",
minute: "2-digit",
hour12: false,
})}
</time>
);
}
// toolInputText renders a tool_use input for display. For Bash it pulls the shell
// command out of the raw input JSON ({"command":…,"description":…}) so the UI shows
// the command itself instead of JSON; other tools fall back to the raw text.
@@ -520,6 +541,7 @@ function UserRow({ step, intent, getDetail }: { step: Activity; intent?: boolean
{text}
</div>
)}
<ActivityTime ts={step.ts} />
</div>
<div className="mt-0.5 flex size-6 shrink-0 items-center justify-center rounded-full bg-primary/10">
<Icon className="size-3.5 text-primary" />
@@ -549,7 +571,7 @@ function AnswerBlock({ step, getDetail }: { step: Activity; getDetail: (seq: num
};
}, [inView, step.seq, getDetail, step.summary]);
return (
<div ref={ref} className="mb-2 mt-1 flex">
<div ref={ref} className="mb-2 mt-1 flex min-w-0 flex-col gap-1">
<div
className={
"min-w-0 flex-1 break-words rounded-lg bg-muted px-3 py-2 " +
@@ -562,6 +584,7 @@ function AnswerBlock({ step, getDetail }: { step: Activity; getDetail: (seq: num
<Markdown text={full ?? step.summary} />
)}
</div>
<ActivityTime ts={step.ts} />
</div>
);
}
web/src/lib/api.ts
+27
View File
@@ -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<TaskArchivePage>(`/task-archives?${query.toString()}`).then((response) => ({
...response,
items: arr(response.items),
}));
},
taskArchive: (id: number) => get<TaskArchive>(`/task-archives/${id}`),
archiveTask: (id: string) => post<TaskArchive>(`/tasks/${id}/archive`),
archiveTasks: (taskIds: string[]) =>
post<{ items: ArchiveBatchItem[] }>("/tasks/archive/batch", { task_ids: taskIds }),
restoreTaskArchive: (id: number) => post<TaskArchive>(`/task-archives/${id}/restore`),
restoreTaskArchives: (archiveIds: number[]) =>
post<{ items: ArchiveBatchItem[] }>("/task-archives/restore/batch", { archive_ids: archiveIds }),
deleteTaskArchive: (id: number) => del<TaskArchive>(`/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<FindingsPage>(`/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<FindingGroupsPage>(`/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();
web/src/lib/mock/handler.ts
+315 -2
View File
@@ -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<string, TaskScopeRow[]>();
type MockTaskAssetSource = Pick<Asset, "task_source" | "task_source_summary" | "task_source_node_id">;
@@ -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<Record<string, number>>((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<ArchiveBatchItem>((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<ArchiveBatchItem>((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<string>();
const visited = new Set<string>();
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<string, ArchiveBatchItem>();
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);
web/src/lib/types.ts
+70
View File
@@ -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<string, number>;
aggregate_stats: {
tokens?: TaskArchiveTokenStats;
skills?: Record<string, number>;
tools?: Record<string, number>;
findings?: Record<string, number>;
};
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";
}