Merge pull request #79 from neouks/codex/sync-v0.3.6

feat: 增加任务冷归档与大任务性能优化
This commit is contained in:
Autumn authored and GitHub committed 2026-08-29 13:52:05 +08:00
commit 6d4408276c
52 files changed
+7366 -511

No files matched your search

+104 -1
View File
@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"strings"
"testing"
"time"
@@ -51,6 +52,12 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 {
t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors)
}
assets := d.Assets()
assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: domain})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) })
task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
@@ -62,7 +69,14 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
tools := NewToolSet(d.Exploration(task.ExplorationID), "planner")
tools.SetTaskID(task.ID)
tools.SetAssetStore(d.Assets(), companies)
tools.SetAssetStore(assets, companies)
linkedAssets, err := assets.QueryByTask(task.ID, "root_domain", 10, 0)
if err != nil || len(linkedAssets) != 1 || linkedAssets[0].ID != assetID {
t.Fatalf("company asset was not linked to task: assets=%+v err=%v", linkedAssets, err)
}
if linkedAssets[0].TaskSource != "company" || !strings.Contains(linkedAssets[0].TaskSourceSummary, companiesName(t, companies, companyID)) {
t.Fatalf("company asset provenance missing: %+v", linkedAssets[0])
}
overview := tools.graphOverviewData()
coverage, ok := overview["coverage"].(map[string]any)
if !ok {
@@ -89,6 +103,95 @@ func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
if !ok || len(keywords) != 1 || keywords[0] != keyword {
t.Fatalf("company keywords missing: %#v", scopeRows[0]["company_keywords"])
}
hosts, ok := coverage["hosts"].([]string)
if !ok || !containsString(hosts, domain) {
t.Fatalf("company asset host missing from agent context: %#v", coverage["hosts"])
}
untested := callReadJSON(t, tools.listUntestedAssets(), `{"type":"root_domain","page":1,"page_size":10}`)
if !strings.Contains(fmt.Sprint(untested), domain) {
t.Fatalf("company asset missing from untested backlog: %#v", untested)
}
byTask := callReadJSON(t, tools.listAssets(), fmt.Sprintf(
`{"dsl":"task_id==%d","type":"root_domain","limit":10}`, task.ID,
))
if !strings.Contains(fmt.Sprint(byTask), domain) {
t.Fatalf("company asset missing from task-scoped agent query: %#v", byTask)
}
byCompany := callReadJSON(t, tools.listAssets(), fmt.Sprintf(
`{"dsl":"company_id==%d","type":"root_domain","limit":10}`, companyID,
))
if !strings.Contains(fmt.Sprint(byCompany), domain) {
t.Fatalf("company asset missing from company-scoped agent query: %#v", byCompany)
}
assetJSON, _ := json.Marshal(assetID)
intentID, err := tools.addOneIntent(intentItem{
Summary: "test associated company asset",
AssetIDs: []json.RawMessage{assetJSON},
})
if err != nil {
t.Fatal(err)
}
workerAssets, err := assets.IntentAssets(task.ID)
if err != nil || len(workerAssets) != 1 || workerAssets[0].IntentID != intentID || workerAssets[0].AssetID != assetID {
t.Fatalf("worker target did not retain company asset: assets=%+v err=%v", workerAssets, err)
}
coverageDisabled := false
disabledTask, err := d.CreateTaskWithOptions("company context without coverage", "still expose company assets", db.TaskCreateOptions{
CompanyIDs: []int64{companyID}, CoverageEnabled: &coverageDisabled,
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(disabledTask.ID) })
disabledTools := NewToolSet(d.Exploration(disabledTask.ExplorationID), "planner")
disabledTools.SetTaskID(disabledTask.ID)
disabledTools.SetCoverageEnabled(false)
disabledTools.SetAssetStore(assets, companies)
disabledOverview := disabledTools.graphOverviewData()
disabledCoverage, ok := disabledOverview["coverage"].(map[string]any)
if !ok {
t.Fatalf("coverage-disabled task lost asset context: %#v", disabledOverview["coverage"])
}
disabledHosts, ok := disabledCoverage["hosts"].([]string)
if !ok || !containsString(disabledHosts, domain) {
t.Fatalf("coverage-disabled task lost company asset host: %#v", disabledCoverage["hosts"])
}
if _, exists := disabledCoverage["denominator"]; exists {
t.Fatalf("coverage-disabled task unexpectedly exposed metrics: %#v", disabledCoverage)
}
if linked, err := assets.QueryByTask(disabledTask.ID, "root_domain", 10, 0); err != nil || len(linked) != 1 || linked[0].ID != assetID {
t.Fatalf("coverage-disabled task asset link=%+v err=%v", linked, err)
}
disabledIntentID, err := disabledTools.addOneIntent(intentItem{
Summary: "test associated company asset without coverage metrics",
AssetIDs: []json.RawMessage{assetJSON},
})
if err != nil {
t.Fatal(err)
}
disabledWorkerAssets, err := assets.IntentAssets(disabledTask.ID)
if err != nil || len(disabledWorkerAssets) != 1 || disabledWorkerAssets[0].IntentID != disabledIntentID || disabledWorkerAssets[0].AssetID != assetID {
t.Fatalf("coverage-disabled worker target=%+v err=%v", disabledWorkerAssets, err)
}
}
func containsString(items []string, target string) bool {
for _, item := range items {
if item == target {
return true
}
}
return false
}
func companiesName(t *testing.T, companies *db.CompanyStore, companyID int64) string {
t.Helper()
company, err := companies.GetCompany(companyID)
if err != nil || company == nil {
t.Fatalf("company %d: company=%+v err=%v", companyID, company, err)
}
return company.Name
}
func TestBlackboardToolsReadDirectSources(t *testing.T) {
+1 -1
View File
@@ -199,7 +199,7 @@ func isFalsy(s string) bool {
return false
}
// Provider returns the short provider name ("anthropic"/"openai"/"openai-responses").
// Provider returns the short provider name ("anthropic"/"openai").
func (c Config) Provider() string {
switch c.Format {
case llm.FormatOpenAI:
+13 -3
View File
@@ -189,11 +189,21 @@ ALTER TABLE llm_records ADD COLUMN IF NOT EXISTS raw_response TEXT;
// EnsureLLMRecordsTable creates the llm_records table if it does not exist.
func (d *DB) EnsureLLMRecordsTable() error {
if _, err := d.Exec(llmRecordsSchema); err != nil {
tx, err := d.Begin()
if err != nil {
return err
}
_, err := d.Exec(llmRecordsMigrate)
return err
defer tx.Rollback() //nolint:errcheck
if err := coordinateWithSchemaMigration(tx); err != nil {
return err
}
if _, err := tx.Exec(llmRecordsSchema); err != nil {
return err
}
if _, err := tx.Exec(llmRecordsMigrate); err != nil {
return err
}
return tx.Commit()
}
// InsertLLMRecord stores one LLM call record.
+7 -3
View File
@@ -177,9 +177,13 @@ func TestCompanyScopeInvalid(t *testing.T) {
}
defer cleanupCompany(d, id)
// TLD-only domains and overly broad CIDRs should be rejected
lines := []string{"com", "1.2.3.4/8"} // rejected
added, _, invalid, _ := cs.AddScope(id, lines, "test")
// Explicitly typed TLD-only domains and overly broad CIDRs should be
// rejected. Untyped plain text is intentionally classified as a keyword.
inputs := []ScopeInput{
{Kind: "domain", Value: "com"},
{Kind: "cidr", Value: "1.2.3.4/8"},
}
added, _, invalid, _ := cs.AddScopeInputs(id, inputs, "test")
if added != 0 {
t.Errorf("want 0 added for invalid lines, got %d", added)
}
+97 -38
View File
@@ -3,19 +3,84 @@
package db
import (
"context"
"database/sql"
_ "embed"
"errors"
"fmt"
"net/url"
"strings"
"time"
"github.com/Autumn-27/artex/config"
"github.com/jackc/pgx/v5/pgconn"
_ "github.com/jackc/pgx/v5/stdlib" // pgx database/sql driver ("pgx")
)
//go:embed schema.sql
var schemaSQL string
const schemaMigrationLockKey int64 = 7337741001
var schemaDeadlockRetryDelays = [...]time.Duration{
100 * time.Millisecond,
250 * time.Millisecond,
500 * time.Millisecond,
time.Second,
}
type schemaExecer interface {
ExecContext(context.Context, string, ...any) (sql.Result, error)
}
func isPostgresDeadlock(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "40P01"
}
func applySchemaWithRetry(ctx context.Context, execer schemaExecer, sleep func(time.Duration)) error {
for attempt := 0; ; attempt++ {
if _, err := execer.ExecContext(ctx, schemaSQL); err != nil {
if !isPostgresDeadlock(err) || attempt >= len(schemaDeadlockRetryDelays) {
return err
}
sleep(schemaDeadlockRetryDelays[attempt])
continue
}
return nil
}
}
// withSchemaMigrationLock pins the session-level lock to one checked-out
// connection. Running pg_advisory_lock through *sql.DB is incorrect because a
// later schema or unlock call may use a different pooled PostgreSQL session.
func withSchemaMigrationLock(ctx context.Context, sqlDB *sql.DB, action func(*sql.Conn) error) (err error) {
conn, err := sqlDB.Conn(ctx)
if err != nil {
return err
}
defer conn.Close()
if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, schemaMigrationLockKey); err != nil {
return fmt.Errorf("advisory lock: %w", err)
}
defer func() {
if _, unlockErr := conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock($1)`, schemaMigrationLockKey); unlockErr != nil && err == nil {
err = fmt.Errorf("advisory unlock: %w", unlockErr)
}
}()
return action(conn)
}
// coordinateWithSchemaMigration makes long, multi-table archive transactions
// mutually exclusive with startup DDL while allowing ordinary runtime queries
// to continue normally.
func coordinateWithSchemaMigration(tx *sql.Tx) error {
if _, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, schemaMigrationLockKey); err != nil {
return fmt.Errorf("coordinate with schema migration: %w", err)
}
return nil
}
// DSN resolves the PostgreSQL connection string and reports where it came from.
// Precedence: env ARTEX_PG_DSN > config file (config.json). There is no
// built-in default — it errors if neither source is configured.
@@ -72,28 +137,22 @@ func Open(dsn string) (*DB, error) {
sqlDB.Close()
return nil, fmt.Errorf("ping postgres (%s): %w", config.Redact(dsn), err)
}
// Serialize schema migration + seeding across concurrent openers (e.g. parallel
// test packages) using a session-level advisory lock. The lock is released
// automatically when the connection is returned to the pool or closed.
if _, err := sqlDB.Exec(`SELECT pg_advisory_lock(7337741001)`); err != nil {
sqlDB.Close()
return nil, fmt.Errorf("advisory lock: %w", err)
}
// pgx runs multi-statement Exec via the simple protocol when there are no args.
if _, err := sqlDB.Exec(schemaSQL); err != nil {
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
sqlDB.Close()
return nil, fmt.Errorf("apply schema: %w", err)
}
d := &DB{sqlDB}
if err := d.seedBuiltins(); err != nil {
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
// pgx runs multi-statement Exec via the simple protocol when there are no args.
// Keep the dedicated lock connection checked out until both DDL and seeding
// finish so concurrent application instances cannot initialize out of order.
err = withSchemaMigrationLock(context.Background(), sqlDB, func(conn *sql.Conn) error {
if err := applySchemaWithRetry(context.Background(), conn, time.Sleep); err != nil {
return fmt.Errorf("apply schema: %w", err)
}
if err := d.seedBuiltins(); err != nil {
return fmt.Errorf("seed builtins: %w", err)
}
return nil
})
if err != nil {
sqlDB.Close()
return nil, fmt.Errorf("seed builtins: %w", err)
}
if _, err := sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`); err != nil {
sqlDB.Close()
return nil, fmt.Errorf("advisory unlock: %w", err)
return nil, err
}
return d, nil
}
@@ -375,30 +434,30 @@ func (d *DB) seedDefaultInterceptRules() error {
// 2. Python HTTP 客户端 .delete() 方法
// 3. JS/通用脚本里的 method: 'DELETE' / method="DELETE"
{
name: "[内置] curl / wget 发送 DELETE 请求",
target: "tool_input",
typ: "regex",
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
action: "deny",
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求,可能删除目标系统数据",
name: "[内置] curl / wget 发送 DELETE 请求",
target: "tool_input",
typ: "regex",
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
action: "deny",
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求,可能删除目标系统数据",
priority: 80,
},
{
name: "[内置] Python HTTP 客户端 DELETE(requests/httpx/aiohttp)",
target: "tool_input",
typ: "regex",
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
action: "deny",
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
name: "[内置] Python HTTP 客户端 DELETE(requests/httpx/aiohttp)",
target: "tool_input",
typ: "regex",
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
action: "deny",
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
priority: 80,
},
{
name: "[内置] 脚本中声明 HTTP DELETE 方法(JS/通用)",
target: "tool_input",
typ: "regex",
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
action: "deny",
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
name: "[内置] 脚本中声明 HTTP DELETE 方法(JS/通用)",
target: "tool_input",
typ: "regex",
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
action: "deny",
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
priority: 80,
},
{
+47
View File
@@ -1,9 +1,56 @@
package db
import (
"context"
"database/sql"
"errors"
"testing"
"time"
"github.com/jackc/pgx/v5/pgconn"
)
type schemaExecFunc func(context.Context, string, ...any) (sql.Result, error)
func (f schemaExecFunc) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
return f(ctx, query, args...)
}
func TestApplySchemaRetriesOnlyDeadlocks(t *testing.T) {
deadlockAttempts := 0
delays := []time.Duration{}
err := applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) {
deadlockAttempts++
if deadlockAttempts < 3 {
return nil, &pgconn.PgError{Code: "40P01", Message: "deadlock detected"}
}
return nil, nil
}), func(delay time.Duration) {
delays = append(delays, delay)
})
if err != nil {
t.Fatal(err)
}
if deadlockAttempts != 3 {
t.Fatalf("schema attempts=%d, want 3", deadlockAttempts)
}
if len(delays) != 2 || delays[0] != schemaDeadlockRetryDelays[0] || delays[1] != schemaDeadlockRetryDelays[1] {
t.Fatalf("schema retry delays=%v", delays)
}
nonRetryable := errors.New("permission denied")
nonRetryableAttempts := 0
err = applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) {
nonRetryableAttempts++
return nil, nonRetryable
}), func(time.Duration) {
t.Fatal("non-deadlock error must not be retried")
})
if !errors.Is(err, nonRetryable) || nonRetryableAttempts != 1 {
t.Fatalf("non-retryable result: attempts=%d err=%v", nonRetryableAttempts, err)
}
}
// testDSN returns the configured DSN, skipping the test when neither the env var
// nor a config file supplies one (DSN no longer has a built-in default).
func testDSN(t *testing.T) string {
+68
View File
@@ -1025,6 +1025,74 @@ func (d *DB) LastActivityAll() (map[int64]int64, error) {
// GoalCounts is the goal summary for one exploration.
type GoalCounts struct{ Total, Met int }
// TaskListMetrics contains the aggregates rendered in task lists.
type TaskListMetrics struct {
Tokens TokenUsage
LastActivity int64
Goals GoalCounts
}
// TaskListMetricsAll returns list aggregates for every live task in one query.
// The lateral lookups use the per-exploration indexes instead of grouping the
// complete activity history on every poll.
func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
rows, err := d.Query(`
SELECT task.exploration_id,
COALESCE(token_metrics.input_tokens,0),
COALESCE(token_metrics.output_tokens,0),
COALESCE(token_metrics.cache_read_tokens,0),
COALESCE(token_metrics.cache_write_tokens,0),
COALESCE(latest_activity.created_at,0),
COALESCE(goal_metrics.total,0),
COALESCE(goal_metrics.met,0)
FROM tasks task
LEFT JOIN LATERAL (
SELECT SUM(input_tokens) AS input_tokens,
SUM(output_tokens) AS output_tokens,
SUM(cache_read_tokens) AS cache_read_tokens,
SUM(cache_write_tokens) AS cache_write_tokens
FROM activity
WHERE exploration_id=task.exploration_id AND kind='result'
) token_metrics ON true
LEFT JOIN LATERAL (
SELECT EXTRACT(EPOCH FROM created_at)::bigint AS created_at
FROM activity
WHERE exploration_id=task.exploration_id
ORDER BY created_at DESC
LIMIT 1
) latest_activity ON true
LEFT JOIN LATERAL (
SELECT COUNT(*) AS total,
COUNT(*) FILTER (WHERE state='met') AS met
FROM exploration_nodes
WHERE exploration_id=task.exploration_id AND kind='goal'
) goal_metrics ON true
WHERE task.deleted_at IS NULL`)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]TaskListMetrics{}
for rows.Next() {
var explorationID int64
var metrics TaskListMetrics
if err := rows.Scan(
&explorationID,
&metrics.Tokens.InputTokens,
&metrics.Tokens.OutputTokens,
&metrics.Tokens.CacheReadTokens,
&metrics.Tokens.CacheWriteTokens,
&metrics.LastActivity,
&metrics.Goals.Total,
&metrics.Goals.Met,
); err != nil {
return nil, err
}
out[explorationID] = metrics
}
return out, rows.Err()
}
// GoalCountsAll returns goal totals (total / met) per exploration id, one query
// for all tasks — used by listTasks so the task list shows progress for every task.
func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) {
+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
+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 {
+92 -5
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
@@ -45,8 +49,18 @@ CREATE INDEX IF NOT EXISTS idx_llm_usage_exp ON llm_usage(exploration_id);
// EnsureLLMUsageTable creates the llm_usage metering table if it does not exist.
func (d *DB) EnsureLLMUsageTable() error {
_, err := d.Exec(llmUsageSchema)
return err
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
if err := coordinateWithSchemaMigration(tx); err != nil {
return err
}
if _, err := tx.Exec(llmUsageSchema); err != nil {
return err
}
return tx.Commit()
}
// InsertLLMUsage appends one metering row. Best-effort: callers log and continue on
@@ -129,7 +143,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 +219,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
+54
View File
@@ -271,6 +271,13 @@ CREATE INDEX IF NOT EXISTS idx_act_since ON activity(exploration_id, id);
-- Main/Plan history pages filter by worker (both carry NULL node_id, so idx_act_node
-- can't distinguish them); this covers reverse pagination of those sessions.
CREATE INDEX IF NOT EXISTS idx_act_worker ON activity(exploration_id, worker, id);
-- Task-list polls aggregate result usage and find the latest event repeatedly.
-- Cover the token columns for index-only aggregation and the timestamp order for
-- per-exploration latest-activity lookups.
CREATE INDEX IF NOT EXISTS idx_act_result_usage ON activity(exploration_id)
INCLUDE (input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
WHERE kind='result';
CREATE INDEX IF NOT EXISTS idx_act_latest ON activity(exploration_id, created_at DESC);
-- =====================================================================
-- C. LLM profiles
@@ -397,6 +404,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(),
@@ -426,10 +434,56 @@ 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 2,
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()
);
ALTER TABLE task_archives ALTER COLUMN format_version SET DEFAULT 2;
CREATE INDEX IF NOT EXISTS idx_task_archives_state ON task_archives(state, requested_at, id);
CREATE INDEX IF NOT EXISTS idx_task_archives_archived ON task_archives(archived_at DESC, id DESC);
CREATE INDEX IF NOT EXISTS idx_task_archives_sources ON task_archives USING GIN(source_task_ids);
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.
+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).
+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
}
+783
View File
@@ -0,0 +1,783 @@
package db
import (
"bytes"
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"io"
"net/url"
"sort"
"strconv"
"strings"
"time"
)
const (
TaskArchiveFormatVersion = 2
TaskArchiveLegacyFormatVersion = 1
TaskArchiveLLMRecordsPath = "database/llm_records.ndjson"
)
func IsTaskArchiveFormatSupported(version int) bool {
return version >= TaskArchiveLegacyFormatVersion && version <= TaskArchiveFormatVersion
}
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.
// Small tables remain JSON arrays in Tables. Large v2 tables are streamed to
// package files listed in StreamedTables so their size is not bounded by memory.
type TaskArchiveSnapshot struct {
FormatVersion int `json:"format_version"`
CreatedAt time.Time `json:"created_at"`
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"`
StreamedTables map[string]string `json:"streamed_tables,omitempty"`
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,
format_version=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.format_version ELSE task_archives.format_version END,
requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`,
taskID, ArchiveQueued, TaskArchiveFormatVersion, name, description, goal, status,
categoryID, categoryName, sources, remaining)
if err != nil {
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 keeps restore/delete resumable after an unclean shutdown.
// An interrupted archive requires an explicit retry: automatic startup retries can
// otherwise form a crash loop when the prior process was killed by resource limits.
func (d *DB) RecoverTaskArchiveJobs() error {
_, err := d.Exec(`UPDATE task_archives SET
state=CASE state WHEN 'archiving' THEN 'archive_failed'
WHEN 'restoring' THEN 'restore_queued'
WHEN 'deleting' THEN 'delete_queued' ELSE state END,
phase='interrupted',
error=CASE WHEN state='archiving' THEN '上次归档进程异常退出,请手动重试' ELSE '' END
WHERE state IN ('archiving','restoring','deleting')`)
return err
}
// 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 {
Query(query string, args ...any) (*sql.Rows, error)
}, inner string, args ...any) (json.RawMessage, int64, error) {
// Do not aggregate the result in PostgreSQL. A jsonb array has a hard limit
// of 256 MiB for its elements, which large LLM request/response histories can
// exceed even though every individual record is valid. Reading row JSON in
// order also avoids building a second copy of the full table in PostgreSQL.
rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...)
if err != nil {
return nil, 0, err
}
defer rows.Close()
return encodeArchiveRows(rows)
}
func encodeArchiveRows(rows interface {
Next() bool
Scan(dest ...any) error
Err() error
}) (json.RawMessage, int64, error) {
var output bytes.Buffer
output.WriteByte('[')
var count int64
for rows.Next() {
var raw []byte
if err := rows.Scan(&raw); err != nil {
return nil, 0, err
}
if count > 0 {
output.WriteByte(',')
}
output.Write(raw)
count++
}
if err := rows.Err(); err != nil {
return nil, 0, err
}
output.WriteByte(']')
return json.RawMessage(output.Bytes()), count, nil
}
func writeArchiveRows(rows interface {
Next() bool
Scan(dest ...any) error
Err() error
}, writer io.Writer) (int64, error) {
var count int64
for rows.Next() {
var raw []byte
if err := rows.Scan(&raw); err != nil {
return 0, err
}
if written, err := writer.Write(raw); err != nil {
return 0, err
} else if written != len(raw) {
return 0, io.ErrShortWrite
}
if written, err := io.WriteString(writer, "\n"); err != nil {
return 0, err
} else if written != 1 {
return 0, io.ErrShortWrite
}
count++
}
if err := rows.Err(); err != nil {
return 0, err
}
return count, nil
}
func streamArchiveRows(q interface {
Query(query string, args ...any) (*sql.Rows, error)
}, writer io.Writer, inner string, args ...any) (int64, error) {
rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...)
if err != nil {
return 0, err
}
defer rows.Close()
return writeArchiveRows(rows, writer)
}
func rawRowCount(raw json.RawMessage) int64 {
var rows []json.RawMessage
if json.Unmarshal(raw, &rows) != nil {
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) {
return d.snapshotTaskArchive(taskID, nil)
}
// SnapshotTaskArchiveWithLLMRecords streams the heavyweight record history to
// llmRecords while all other task-owned data is read from the same repeatable
// PostgreSQL snapshot.
func (d *DB) SnapshotTaskArchiveWithLLMRecords(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) {
if llmRecords == nil {
return nil, errors.New("nil LLM record archive writer")
}
return d.snapshotTaskArchive(taskID, llmRecords)
}
func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) {
tx, err := d.BeginTx(context.Background(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true})
if err != nil {
return nil, err
}
defer tx.Rollback() //nolint:errcheck
if err := coordinateWithSchemaMigration(tx); err != nil {
return nil, err
}
var expID int64
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 AND deleted_at IS NULL`, taskID).Scan(&expID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
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))
streamedTables := map[string]string{}
for _, query := range queries {
if query.name == "llm_records" && llmRecords != nil {
count, err := streamArchiveRows(tx, llmRecords, query.query, query.args...)
if err != nil {
return nil, fmt.Errorf("snapshot %s: %w", query.name, err)
}
tables[query.name] = json.RawMessage("[]")
counts[query.name] = count
streamedTables[query.name] = TaskArchiveLLMRecordsPath
continue
}
raw, count, err := queryArchiveRows(tx, query.query, query.args...)
if err != nil {
return nil, fmt.Errorf("snapshot %s: %w", query.name, err)
}
tables[query.name] = raw
counts[query.name] = count
}
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, StreamedTables: streamedTables,
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
}
}
+737
View File
@@ -0,0 +1,737 @@
package db
import (
"bytes"
"database/sql"
"encoding/json"
"errors"
"fmt"
"io"
"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
if err := coordinateWithSchemaMigration(tx); err != nil {
return err
}
var taskID, expID int64
var state string
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
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,
format_version=$9,archived_at=now(),warnings='[]'
WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize,
string(countsRaw), string(statsRaw), snapshot.FormatVersion); err != nil {
return err
}
return tx.Commit()
}
// 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) {
return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, nil)
}
// RestoreTaskArchiveWithLLMRecords restores a v2 package whose heavyweight LLM
// record history is stored as a sequence of JSON objects outside manifest.json.
func (d *DB) RestoreTaskArchiveWithLLMRecords(
archiveID int64,
snapshot *TaskArchiveSnapshot,
remainingTimeoutSeconds int64,
llmRecords io.Reader,
) ([]string, error) {
if llmRecords == nil {
return nil, errors.New("nil streamed LLM record reader")
}
return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, llmRecords)
}
func (d *DB) restoreTaskArchive(
archiveID int64,
snapshot *TaskArchiveSnapshot,
remainingTimeoutSeconds int64,
llmRecords io.Reader,
) ([]string, error) {
if snapshot == nil || !IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
return nil, ErrTaskArchiveFormatMismatch
}
streamedLLMRecords := snapshot.StreamedTables["llm_records"]
if (len(snapshot.StreamedTables) > 0 && snapshot.FormatVersion < 2) ||
len(snapshot.StreamedTables) > 1 ||
(len(snapshot.StreamedTables) == 1 && streamedLLMRecords != TaskArchiveLLMRecordsPath) {
return nil, fmt.Errorf("%w: unsupported streamed table metadata", ErrTaskArchiveFormatMismatch)
}
if streamedLLMRecords != "" && llmRecords == nil {
return nil, fmt.Errorf("%w: streamed LLM records are missing", ErrTaskArchiveFormatMismatch)
}
tx, err := d.Begin()
if err != nil {
return nil, err
}
defer tx.Rollback() //nolint:errcheck
if err := coordinateWithSchemaMigration(tx); err != nil {
return nil, err
}
var taskID, expID int64
var state string
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
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"} {
if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil {
return nil, fmt.Errorf("restore %s: %w", table, err)
}
}
if streamedLLMRecords != "" {
count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords)
if err != nil {
return nil, fmt.Errorf("restore llm_records: %w", err)
}
if expected := snapshot.DataCounts["llm_records"]; count != expected {
return nil, fmt.Errorf("restore llm_records: row count %d does not match manifest %d", count, expected)
}
} else if err := insertArchiveRows(tx, "llm_records", remappedTables["llm_records"]); err != nil {
return nil, fmt.Errorf("restore llm_records: %w", err)
}
for _, table := range []string{"llm_usage", "skill_usage", "tool_usage"} {
if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil {
return nil, fmt.Errorf("restore %s: %w", table, err)
}
}
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 insertArchiveJSONSequenceRows(tx *sql.Tx, table string, reader io.Reader) (int64, error) {
if table != "llm_records" {
return 0, fmt.Errorf("archive streamed restore table %q is not allowed", table)
}
statement, err := tx.Prepare(`INSERT INTO llm_records SELECT * FROM json_populate_record(NULL::llm_records,$1::json)`)
if err != nil {
return 0, err
}
defer statement.Close()
decoder := json.NewDecoder(reader)
var count int64
for {
var raw json.RawMessage
if err := decoder.Decode(&raw); errors.Is(err, io.EOF) {
break
} else if err != nil {
return 0, err
}
trimmed := bytes.TrimSpace(raw)
if len(trimmed) == 0 || trimmed[0] != '{' {
return 0, errors.New("streamed archive row must be a JSON object")
}
if _, err := statement.Exec(string(trimmed)); err != nil {
return 0, err
}
count++
}
return count, nil
}
func restoreArchiveAssets(tx *sql.Tx, taskID int64, raw json.RawMessage) (map[int64]int64, []string, error) {
rows, err := decodeArchiveRows(raw)
if err != nil {
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
if err := coordinateWithSchemaMigration(tx); err != nil {
return err
}
var taskID, expID int64
var state string
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
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()
}
+468
View File
@@ -0,0 +1,468 @@
package db
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"strings"
"testing"
"time"
)
type archiveJSONTestRows struct {
values [][]byte
next int
}
func (r *archiveJSONTestRows) Next() bool {
if r.next >= len(r.values) {
return false
}
r.next++
return true
}
func (r *archiveJSONTestRows) Scan(dest ...any) error {
if len(dest) != 1 || r.next == 0 || r.next > len(r.values) {
return fmt.Errorf("invalid archive test row scan")
}
target, ok := dest[0].(*[]byte)
if !ok {
return fmt.Errorf("archive test row destination is %T", dest[0])
}
*target = append((*target)[:0], r.values[r.next-1]...)
return nil
}
func (r *archiveJSONTestRows) Err() error { return nil }
func TestQueryArchiveRowsStreamsJSON(t *testing.T) {
payload := strings.Repeat("large request/response payload ", 64*1024)
values := make([][]byte, 2)
var err error
values[0], err = json.Marshal(map[string]any{"id": int64(1), "body": payload})
if err != nil {
t.Fatal(err)
}
values[1], err = json.Marshal(map[string]any{"id": int64(2), "body": "quoted: \"value\"\nline"})
if err != nil {
t.Fatal(err)
}
raw, count, err := encodeArchiveRows(&archiveJSONTestRows{values: values})
if err != nil {
t.Fatal(err)
}
if count != 2 {
t.Fatalf("row count=%d, want 2", count)
}
var rows []struct {
ID int64 `json:"id"`
Body string `json:"body"`
}
if err := json.Unmarshal(raw, &rows); err != nil {
t.Fatal(err)
}
if len(rows) != 2 {
t.Fatalf("decoded row count=%d, want 2", len(rows))
}
if rows[0].ID != 1 || rows[0].Body != payload || rows[1].ID != 2 || rows[1].Body != "quoted: \"value\"\nline" {
t.Fatal("streamed rows were reordered or truncated")
}
}
func TestWriteArchiveRowsProducesJSONSequence(t *testing.T) {
values := [][]byte{
json.RawMessage(`{"id":1,"body":"first"}`),
json.RawMessage(`{"id":2,"body":"second\\nline"}`),
}
var output bytes.Buffer
count, err := writeArchiveRows(&archiveJSONTestRows{values: values}, &output)
if err != nil {
t.Fatal(err)
}
if count != 2 {
t.Fatalf("streamed row count=%d, want 2", count)
}
decoder := json.NewDecoder(&output)
for wantID := int64(1); wantID <= 2; wantID++ {
var row struct {
ID int64 `json:"id"`
}
if err := decoder.Decode(&row); err != nil {
t.Fatal(err)
}
if row.ID != wantID {
t.Fatalf("streamed row id=%d, want %d", row.ID, wantID)
}
}
}
func TestTaskArchiveFormatCompatibility(t *testing.T) {
for _, version := range []int{TaskArchiveLegacyFormatVersion, TaskArchiveFormatVersion} {
if !IsTaskArchiveFormatSupported(version) {
t.Fatalf("archive format %d should be supported", version)
}
}
for _, version := range []int{0, TaskArchiveFormatVersion + 1} {
if IsTaskArchiveFormatSupported(version) {
t.Fatalf("archive format %d should be rejected", version)
}
}
invalidSnapshots := []*TaskArchiveSnapshot{
{FormatVersion: TaskArchiveLegacyFormatVersion, StreamedTables: map[string]string{"llm_records": TaskArchiveLLMRecordsPath}},
{FormatVersion: TaskArchiveFormatVersion, StreamedTables: map[string]string{"unknown": "database/unknown.ndjson"}},
}
for _, snapshot := range invalidSnapshots {
if _, err := (&DB{}).RestoreTaskArchive(1, snapshot, 0); !errors.Is(err, ErrTaskArchiveFormatMismatch) {
t.Fatalf("invalid streamed table metadata returned %v", err)
}
}
}
func TestTaskArchiveDatabaseRoundTrip(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
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.InsertLLMRecord(&LLMRecord{
TaskID: fmt.Sprint(task.ID), Model: "archive-model", SessionID: "archive-session",
Status: "ok", RawRequest: strings.Repeat("request", 1024), RawResponse: strings.Repeat("response", 1024),
}); err != nil {
t.Fatal(err)
}
if _, err := d.AddFinding(task.ID, nodeID, vulnclass, "archive finding", SeverityCritical, "summary", "evidence", "worker", []int64{assetID}); err != nil {
t.Fatal(err)
}
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)
}
var llmRecords bytes.Buffer
snapshot, err := d.SnapshotTaskArchiveWithLLMRecords(task.ID, &llmRecords)
if err != nil {
t.Fatal(err)
}
if snapshot.StreamedTables["llm_records"] != TaskArchiveLLMRecordsPath || snapshot.DataCounts["llm_records"] != 1 {
t.Fatalf("unexpected streamed LLM metadata: paths=%v counts=%v", snapshot.StreamedTables, snapshot.DataCounts)
}
if rawRowCount(snapshot.Tables["llm_records"]) != 0 {
t.Fatal("streamed LLM records were also retained in manifest memory")
}
if snapshot.DataCounts["assets"] != 1 || snapshot.DataCounts["exploration_nodes"] < 2 {
t.Fatalf("unexpected snapshot counts: %#v", snapshot.DataCounts)
}
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.RestoreTaskArchiveWithLLMRecords(
archive.ID, snapshot, int64((10 * time.Minute).Seconds()), bytes.NewReader(llmRecords.Bytes()),
)
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 restoredRawRequest, restoredRawResponse string
if err := d.QueryRow(`SELECT COALESCE(raw_request,''),COALESCE(raw_response,'') FROM llm_records WHERE task_id=$1`, fmt.Sprint(task.ID)).
Scan(&restoredRawRequest, &restoredRawResponse); err != nil {
t.Fatal(err)
}
if restoredRawRequest != strings.Repeat("request", 1024) || restoredRawResponse != strings.Repeat("response", 1024) {
t.Fatal("restored streamed LLM record body was truncated")
}
var restoredCompanyID *int64
var restoredStatus, restoredError string
var restoredExhaustedAt, restoredCreatedAt time.Time
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 TestRecoverInterruptedArchiveRequiresManualRetry(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
task, err := d.CreateTaskWithOptions("interrupted archive", "must not restart automatically", TaskCreateOptions{})
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID)
_ = d.DeleteTask(task.ID)
}()
if err := d.SetPaused(task.ID, true); err != nil {
t.Fatal(err)
}
queued, err := d.QueueTaskArchive(task.ID)
if err != nil {
t.Fatal(err)
}
if _, err := d.Exec(`UPDATE task_archives SET state=$2,phase='snapshot_database',progress=10 WHERE id=$1`, queued.ID, Archiving); err != nil {
t.Fatal(err)
}
if err := d.RecoverTaskArchiveJobs(); err != nil {
t.Fatal(err)
}
recovered, err := d.GetTaskArchive(queued.ID)
if err != nil {
t.Fatal(err)
}
if recovered == nil || recovered.State != ArchiveFailed || recovered.Phase != "interrupted" || recovered.Error == "" {
t.Fatalf("recovered archive=%+v, want explicit-retry failure", recovered)
}
}
func assertArchiveGlobalStats(t *testing.T, d *DB, profileName, skillName, toolName, vulnclass string) {
t.Helper()
profiles, err := d.UsageByProfile()
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)
}
}
+163 -3
View File
@@ -52,9 +52,6 @@ func (d *DB) hydrateTaskContext(t *Task) error {
if err != nil {
return err
}
t.LLMProfileID = legacyProfileID
t.ActiveLLMProfileID = activeProfileID
t.LLMChainRevision = revision
sources, err := d.TaskSourceIDs(t.ID)
if err != nil {
return err
@@ -65,8 +62,22 @@ func (d *DB) hydrateTaskContext(t *Task) error {
return err
}
t.CompanyIDs = companies
applyTaskLLMContext(t, legacyProfileID, activeProfileID, revision, chain)
return nil
}
func applyTaskLLMContext(
t *Task,
legacyProfileID, activeProfileID *int64,
revision int64,
chain []TaskLLMProfile,
) {
t.LLMProfileID = legacyProfileID
t.ActiveLLMProfileID = activeProfileID
t.LLMChainRevision = revision
t.LLMProfileIDs = make([]int64, 0, len(chain))
t.LLMFailoverState = "default"
t.LLMFailoverReason = ""
var latest *TaskLLMProfile
activeReady := false
for i := range chain {
@@ -90,6 +101,155 @@ func (d *DB) hydrateTaskContext(t *Task) error {
if latest != nil {
t.LLMFailoverReason = latest.LastError
}
}
type taskBatchContext struct {
legacyProfileID *int64
activeProfileID *int64
revision int64
chain []TaskLLMProfile
sourceIDs []int64
companyIDs []int64
}
// hydrateTasksContext loads every task's LLM chain, source tasks, and company
// scopes with three bulk queries. ListTasks used to issue these queries once per
// task, making startup and task-list hydration grow as 3N+1 database round trips.
func (d *DB) hydrateTasksContext(tasks []*Task) error {
if len(tasks) == 0 {
return nil
}
ids := make([]int64, 0, len(tasks))
contexts := make(map[int64]*taskBatchContext, len(tasks))
byID := make(map[int64]*Task, len(tasks))
for _, task := range tasks {
if task == nil {
continue
}
ids = append(ids, task.ID)
contexts[task.ID] = &taskBatchContext{}
byID[task.ID] = task
}
if len(ids) == 0 {
return nil
}
rows, err := d.Query(`
SELECT t.id, t.llm_profile_id, t.active_llm_profile_id, t.llm_chain_revision,
p.profile_id, p.position, p.status, COALESCE(p.last_error,''), p.exhausted_at
FROM tasks t
LEFT JOIN task_llm_profiles p ON p.task_id=t.id
WHERE t.id=ANY($1::bigint[])
ORDER BY t.id, p.position NULLS LAST`, ids)
if err != nil {
return err
}
for rows.Next() {
var taskID, revision int64
var legacy, active, profileID, position sql.NullInt64
var status, lastError sql.NullString
var exhaustedAt sql.NullTime
if err := rows.Scan(
&taskID, &legacy, &active, &revision,
&profileID, &position, &status, &lastError, &exhaustedAt,
); err != nil {
rows.Close()
return err
}
context := contexts[taskID]
if context == nil {
continue
}
context.revision = revision
if legacy.Valid {
id := legacy.Int64
context.legacyProfileID = &id
}
if active.Valid {
id := active.Int64
context.activeProfileID = &id
}
if profileID.Valid {
entry := TaskLLMProfile{
ProfileID: profileID.Int64,
Position: int(position.Int64),
Status: status.String,
LastError: lastError.String,
}
if exhaustedAt.Valid {
ts := exhaustedAt.Time
entry.ExhaustedAt = &ts
}
context.chain = append(context.chain, entry)
}
}
if err := rows.Err(); err != nil {
rows.Close()
return err
}
rows.Close()
rows, err = d.Query(`
SELECT task_id, source_task_id
FROM task_relations
WHERE task_id=ANY($1::bigint[])
ORDER BY task_id, created_at, source_task_id`, ids)
if err != nil {
return err
}
for rows.Next() {
var taskID, sourceID int64
if err := rows.Scan(&taskID, &sourceID); err != nil {
rows.Close()
return err
}
if context := contexts[taskID]; context != nil {
context.sourceIDs = append(context.sourceIDs, sourceID)
}
}
if err := rows.Err(); err != nil {
rows.Close()
return err
}
rows.Close()
rows, err = d.Query(`
SELECT task_id, company_id
FROM task_scope
WHERE task_id=ANY($1::bigint[]) AND kind='company' AND company_id IS NOT NULL
ORDER BY task_id, id`, ids)
if err != nil {
return err
}
for rows.Next() {
var taskID, companyID int64
if err := rows.Scan(&taskID, &companyID); err != nil {
rows.Close()
return err
}
if context := contexts[taskID]; context != nil {
context.companyIDs = append(context.companyIDs, companyID)
}
}
if err := rows.Err(); err != nil {
rows.Close()
return err
}
rows.Close()
for id, context := range contexts {
task := byID[id]
task.SourceTaskIDs = context.sourceIDs
task.CompanyIDs = context.companyIDs
applyTaskLLMContext(
task,
context.legacyProfileID,
context.activeProfileID,
context.revision,
context.chain,
)
}
return nil
}
+117
View File
@@ -0,0 +1,117 @@
package db
import (
"fmt"
"testing"
"time"
)
func TestListTasksBulkHydratesTaskContext(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
stamp := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("bulk-list-profile-%d", stamp), Format: "openai", Model: "test", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteProfile(profileID) })
source, err := d.CreateTask(fmt.Sprintf("bulk-list-source-%d", stamp), "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(source.ID) })
companyID, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Bulk List Company %d", stamp), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.Companies().DeleteCompany(companyID) })
child, err := d.CreateTaskWithOptions("bulk-list-child", "goal", TaskCreateOptions{
LLMProfileIDs: []int64{profileID},
SourceTaskIDs: []int64{source.ID},
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(child.ID) })
tasks, err := d.ListTasks()
if err != nil {
t.Fatal(err)
}
var got *Task
for _, task := range tasks {
if task.ID == child.ID {
got = task
break
}
}
if got == nil {
t.Fatalf("task %d missing from list", child.ID)
}
if len(got.LLMProfileIDs) != 1 || got.LLMProfileIDs[0] != profileID ||
got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileID || got.LLMFailoverState != "ready" {
t.Fatalf("LLM context was not bulk hydrated: %+v", got)
}
if len(got.SourceTaskIDs) != 1 || got.SourceTaskIDs[0] != source.ID {
t.Fatalf("source task context was not bulk hydrated: %v", got.SourceTaskIDs)
}
if len(got.CompanyIDs) != 1 || got.CompanyIDs[0] != companyID {
t.Fatalf("company context was not bulk hydrated: %v", got.CompanyIDs)
}
}
func TestTaskListMetricsAll(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("task-list-metrics", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
resultAt := time.Now().Add(-time.Second).Truncate(time.Second)
latestAt := resultAt.Add(time.Second)
if _, err := d.Exec(`
INSERT INTO activity(
exploration_id, kind, input_tokens, output_tokens,
cache_read_tokens, cache_write_tokens, created_at
) VALUES ($1,'result',11,7,3,2,$2), ($1,'tool_result',999,999,999,999,$3)`,
task.ExplorationID, resultAt, latestAt); err != nil {
t.Fatal(err)
}
if _, err := d.Exec(`
INSERT INTO exploration_nodes(exploration_id, kind, payload, state)
VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExplorationID); err != nil {
t.Fatal(err)
}
all, err := d.TaskListMetricsAll()
if err != nil {
t.Fatal(err)
}
metrics, ok := all[task.ExplorationID]
if !ok {
t.Fatalf("metrics for exploration %d missing", task.ExplorationID)
}
if metrics.Tokens.InputTokens != 11 || metrics.Tokens.OutputTokens != 7 ||
metrics.Tokens.CacheReadTokens != 3 || metrics.Tokens.CacheWriteTokens != 2 {
t.Fatalf("unexpected token metrics: %+v", metrics.Tokens)
}
if metrics.LastActivity != latestAt.Unix() {
t.Fatalf("last activity=%d, want %d", metrics.LastActivity, latestAt.Unix())
}
if metrics.Goals.Total != 2 || metrics.Goals.Met != 1 {
t.Fatalf("unexpected goal metrics: %+v", metrics.Goals)
}
}
+37 -6
View File
@@ -89,10 +89,12 @@ const MinPlanHeartbeatSeconds = 600
const MaxTaskSourceCount = 8
// MaxTaskCompanyCount bounds the number of company asset scopes attached to a
// task. Company scopes are prompt context, so an unbounded list would inflate
// every planner and worker call.
// task. Company scopes are prompt context, and their currently attributed
// assets are snapshotted into the task at creation time.
const MaxTaskCompanyCount = 32
const taskCompanyAssetSource = "company"
var (
ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids")
ErrTaskCompanyNotFound = errors.New("task company not found")
@@ -248,6 +250,11 @@ func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
if len(companyIDs) == 0 {
return nil
}
// Company scope edits rebuild assets.company_id. Serialize the creation-time
// snapshot with those edits so the task sees one committed attribution state.
if err := lockCompanyScopeMutation(tx); err != nil {
return err
}
var inserted int
err := tx.QueryRow(`
WITH requested(company_id, position) AS (
@@ -268,6 +275,32 @@ SELECT count(*) FROM inserted`, taskID, companyIDs).Scan(&inserted)
if inserted != len(companyIDs) {
return fmt.Errorf("%w: one or more companies do not exist", ErrTaskCompanyNotFound)
}
// Updating task_ids fires sync_task_asset_links, which first creates generic
// source rows. The provenance upsert must therefore run afterwards so the
// company name and creation-time reason remain visible to operators.
if _, err := tx.Exec(`
UPDATE assets
SET task_ids=CASE
WHEN $1=ANY(task_ids) THEN task_ids
ELSE array_append(task_ids, $1)
END
WHERE company_id=ANY($2::bigint[])`, taskID, companyIDs); err != nil {
return err
}
if _, err := tx.Exec(`
INSERT INTO task_asset_links(task_id, asset_id, source, source_summary)
SELECT $1, asset.id, $3, '任务创建时关联企业:' || company.name
FROM assets asset
JOIN companies company ON company.id=asset.company_id
WHERE asset.company_id=ANY($2::bigint[])
AND $1=ANY(asset.task_ids)
ON CONFLICT (task_id, asset_id) DO UPDATE
SET source=EXCLUDED.source,
source_summary=EXCLUDED.source_summary,
source_node_id=NULL`, taskID, companyIDs, taskCompanyAssetSource); err != nil {
return err
}
return nil
}
@@ -344,10 +377,8 @@ ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, id DESC`)
return nil, err
}
rows.Close()
for _, task := range out {
if err := d.hydrateTaskContext(task); err != nil {
return nil, err
}
if err := d.hydrateTasksContext(out); err != nil {
return nil, err
}
return out, nil
}
+78 -1
View File
@@ -509,10 +509,36 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) {
if err != nil {
t.Fatal(err)
}
emptyCompany, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Empty Company %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2,$3)`, companyA, companyB, emptyCompany)
})
existingTask, err := d.CreateTask("existing company asset owner", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(existingTask.ID) })
var companyAssetA, companyAssetB int64
domainA := fmt.Sprintf("task-company-a-%d.example.test", suffix)
domainB := fmt.Sprintf("task-company-b-%d.example.test", suffix)
if err := d.QueryRow(`
INSERT INTO assets(type, domain, root_domain, company_id, company_source, task_ids)
VALUES ('root_domain',$1,$1,$2,'explicit',ARRAY[$3]::bigint[])
RETURNING id`, domainA, companyA, existingTask.ID).Scan(&companyAssetA); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`
INSERT INTO assets(type, domain, root_domain, company_id, company_source)
VALUES ('root_domain',$1,$1,$2,'explicit')
RETURNING id`, domainB, companyB).Scan(&companyAssetB); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{companyAssetA, companyAssetB}) })
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyB},
})
@@ -537,6 +563,57 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) {
if scopeCount != 2 {
t.Fatalf("company task scope rows=%d want 2", scopeCount)
}
assets, err := d.Assets().QueryByTask(task.ID, "root_domain", 10, 0)
if err != nil {
t.Fatal(err)
}
if len(assets) != 2 {
t.Fatalf("company task assets=%d want 2: %+v", len(assets), assets)
}
assetByID := make(map[int64]*Asset, len(assets))
for _, asset := range assets {
assetByID[asset.ID] = asset
}
for assetID, companyName := range map[int64]string{
companyAssetA: fmt.Sprintf("Task Company A %d", suffix),
companyAssetB: fmt.Sprintf("Task Company B %d", suffix),
} {
asset := assetByID[assetID]
if asset == nil {
t.Errorf("company asset %d missing from task", assetID)
continue
}
if asset.TaskSource != taskCompanyAssetSource || asset.TaskSourceSummary != "任务创建时关联企业:"+companyName {
t.Errorf("asset %d provenance=%q/%q", assetID, asset.TaskSource, asset.TaskSourceSummary)
}
}
var existingAssociation bool
if err := d.QueryRow(`SELECT $1=ANY(task_ids) FROM assets WHERE id=$2`, existingTask.ID, companyAssetA).Scan(&existingAssociation); err != nil {
t.Fatal(err)
}
if !existingAssociation {
t.Fatal("company snapshot removed the asset's existing task association")
}
var existingSource string
if err := d.QueryRow(`SELECT source FROM task_asset_links WHERE task_id=$1 AND asset_id=$2`, existingTask.ID, companyAssetA).Scan(&existingSource); err != nil {
t.Fatal(err)
}
if existingSource == taskCompanyAssetSource {
t.Fatalf("company snapshot overwrote another task's provenance: %q", existingSource)
}
emptyTask, err := d.CreateTaskWithOptions("empty company scope", "no current assets", TaskCreateOptions{
CompanyIDs: []int64{emptyCompany},
})
if err != nil {
t.Fatal(err)
}
if assets, err := d.Assets().QueryByTask(emptyTask.ID, "", 10, 0); err != nil || len(assets) != 0 {
t.Fatalf("empty company task assets=%+v err=%v", assets, err)
}
if err := d.DeleteTask(emptyTask.ID); err != nil {
t.Fatal(err)
}
duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyA, companyB, companyA},
+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
}
+1 -1
View File
@@ -6,6 +6,7 @@ require (
github.com/Autumn-27/norma v0.3.2
github.com/golang-jwt/jwt/v5 v5.3.1
github.com/jackc/pgx/v5 v5.10.0
github.com/klauspost/compress v1.18.6
github.com/lqqyt2423/go-mitmproxy v1.9.2
github.com/miekg/dns v1.1.72
github.com/projectdiscovery/dnsx v1.2.3
@@ -29,7 +30,6 @@ require (
github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/jackc/puddle/v2 v2.2.2 // indirect
github.com/klauspost/compress v1.18.6 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/microcosm-cc/bluemonday v1.0.27 // indirect
-2
View File
@@ -1,5 +1,3 @@
github.com/Autumn-27/norma v0.3.1 h1:cSRojFJjivm4l3ck/KJXym7VyCnM1GLY7iq1pRYk4kk=
github.com/Autumn-27/norma v0.3.1/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
github.com/Autumn-27/norma v0.3.2 h1:3gW84SbSrz8luW8kbi2tv3O7CLdRMQvyylCHj1QUdNk=
github.com/Autumn-27/norma v0.3.2/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
+6 -16
View File
@@ -185,12 +185,8 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
}
}
// Complete implements llm.Provider for non-streaming calls, recording the same
// metering + body trace as Stream. Metering (token usage) is always written;
// the heavy request/response body trace is gated by the llm_record setting. The
// raw wire bodies are captured by the shared HTTP transport (via NewCapture on
// ctx), identical to the streaming path — a non-streaming response is a single
// JSON body the transport tees just the same.
// Complete implements the atomic completion path while preserving the same
// usage metering, normalized body recording and raw wire capture as Stream.
func (r *Recorder) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
recordBodies := r.enabled == nil || r.enabled()
start := time.Now()
@@ -210,28 +206,22 @@ func (r *Recorder) Complete(ctx context.Context, req llm.CompletionRequest) (llm
}
msg, stopReason, usage, err := r.inner.Complete(ctx, req)
status := "ok"
if err != nil {
status = "error"
}
// Lightweight metering row — always written.
r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status)
// Heavy trace row — only when body recording is on.
if recordBodies {
text, thinking := msg.Text(), thinkingText(msg)
r.record(req, session, taskID, worker, reqBody, capt, start, text, thinking, usage, stopReason, err)
r.record(req, session, taskID, worker, reqBody, capt, start, msg.Text(), thinkingText(msg), usage, stopReason, err)
}
return msg, stopReason, usage, err
}
// thinkingText concatenates the thinking blocks of a message (the non-streaming
// counterpart of the streaming path's accumulated thinkingBuf).
func thinkingText(msg llm.Message) string {
var b strings.Builder
for _, bl := range msg.Content {
if bl.Type == llm.BlockThinking && bl.Thinking != "" {
b.WriteString(bl.Thinking)
for _, block := range msg.Content {
if block.Type == llm.BlockThinking && block.Thinking != "" {
b.WriteString(block.Thinking)
}
}
return b.String()
+36
View File
@@ -2,9 +2,28 @@ package llmrec
import (
"context"
"iter"
"testing"
"github.com/Autumn-27/norma/llm"
)
type completeProvider struct{}
func (completeProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(func(llm.StreamEvent, error) bool) {}
}
func (completeProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
return llm.Message{
Role: llm.RoleAssistant,
Content: []llm.ContentBlock{
{Type: llm.BlockThinking, Thinking: "reasoning"},
llm.TextBlock("answer"),
},
}, "stop", llm.Usage{InputTokens: 7, OutputTokens: 3}, nil
}
func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) {
ctx := WithTaskID(context.Background(), " 42 ")
if got := TaskIDFrom(ctx); got != "42" {
@@ -27,3 +46,20 @@ func TestParseSessionFallbackIsExplorationScoped(t *testing.T) {
t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker)
}
}
func TestCompleteForwardsAtomicResponse(t *testing.T) {
recorder := Wrap(completeProvider{}, nil, "model", "profile", "", "", func() bool { return false })
msg, stopReason, usage, err := recorder.Complete(context.Background(), llm.CompletionRequest{})
if err != nil {
t.Fatalf("Complete() error: %v", err)
}
if got := msg.Text(); got != "answer" {
t.Fatalf("Complete() text=%q want answer", got)
}
if stopReason != "stop" {
t.Fatalf("Complete() stop reason=%q want stop", stopReason)
}
if usage.InputTokens != 7 || usage.OutputTokens != 3 {
t.Fatalf("Complete() usage=%+v", usage)
}
}
+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"`
+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)
+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()
+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())
+4 -4
View File
@@ -450,11 +450,11 @@ func (s *Server) seedOrchestrationTools() {
s.seedPlannerDefaultBindings()
s.seedAutoReportFindingBinding()
s.unbindGoalMetDefault()
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
// 注:pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
// pentest 一起 seed 好了(项目尚无旧库,不做迁移)。
+31 -12
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
}
@@ -705,6 +711,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)
@@ -1011,22 +1025,22 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
active = t.ID
}
list := s.m.List()
toks, _ := s.m.PG().TokenTotalsAll() // whole-task token totals, one query for all tasks
lastAct, _ := s.m.PG().LastActivityAll() // persisted last-activity per task, one query
goalCounts, _ := s.m.PG().GoalCountsAll() // goal progress per exploration, one query
metrics, _ := s.m.PG().TaskListMetricsAll()
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
dtos := make([]TaskDTO, 0, len(list))
for _, t := range list {
dto := taskDTO(t, s.resolvedTaskStatus(t))
dto.Tokens = tokenTotalDTO(toks[t.ExpID])
applyTaskArchiveBlocker(&dto, archiveBlockers)
metric := metrics[t.ExpID]
dto.Tokens = tokenTotalDTO(metric.Tokens)
// prefer the live in-memory heartbeat (fresher) and fall back to the
// persisted max activity time (survives restarts) for run-duration display.
dto.LastActivity = lastAct[t.ExpID]
dto.LastActivity = metric.LastActivity
if live := s.engine.LastActivity(t.ID); live > dto.LastActivity {
dto.LastActivity = live
}
gc := goalCounts[t.ExpID]
dto.GoalsTotal = gc.Total
dto.GoalsMet = gc.Met
dto.GoalsTotal = metric.Goals.Total
dto.GoalsMet = metric.Goals.Met
dtos = append(dtos, dto)
}
writeJSON(w, 200, map[string]any{"tasks": dtos, "active": active})
@@ -1400,7 +1414,7 @@ type createTaskReq struct {
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效
SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务
CompanyIDs []int64 `json:"company_ids,omitempty"` // 关联企业资产范围;不复制资产或强制生成意图
CompanyIDs []int64 `json:"company_ids,omitempty"` // 关联企业范围并快照关联当前企业资产;不复制资产或强制生成意图
TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600
SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。
@@ -1679,7 +1693,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).
@@ -1881,6 +1898,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)
@@ -2050,7 +2068,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")
@@ -2084,6 +2102,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:
+709
View File
@@ -0,0 +1,709 @@
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)
}
// The journal is written before the first rename, so its presence does not mean
// the payload is complete: a previous round may have failed mid-loop with a
// rollback that itself errored, which leaves the root in place. Replay the moves
// rather than packaging a payload that is missing transcripts or workspace files.
if err := os.MkdirAll(filepath.Join(stage.payload, "files"), archiveDirMode); err != nil {
return nil, err
}
if err := stage.applyMoves(); err != nil {
_ = stage.rollback()
return nil, err
}
return stage, nil
} else if !os.IsNotExist(err) {
return nil, err
}
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
}
if err := stage.applyMoves(); err != nil {
_ = stage.rollback()
return nil, err
}
return stage, nil
}
// applyMoves performs the renames recorded in the journal. Entries already sitting
// in the payload are skipped, so an interrupted staging round can be resumed in
// place without moving anything twice.
func (s *taskArchiveFileStage) applyMoves() error {
for _, move := range s.journal.Moves {
destination := filepath.Join(s.payload, move.Relative)
if _, err := os.Lstat(destination); err == nil {
continue
} else if !os.IsNotExist(err) {
return err
}
if _, err := os.Lstat(move.Source); err != nil {
if os.IsNotExist(err) {
// Neither side exists: the path was removed outside the archive flow
// after the journal was written. Nothing can be staged for it.
continue
}
return err
}
if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil {
return err
}
if err := os.Rename(move.Source, destination); err != nil {
return fmt.Errorf("stage task archive path %s: %w", move.Source, err)
}
}
return nil
}
func (s *taskArchiveFileStage) rollback() error {
if s == nil || s.done {
return nil
}
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 os.IsNotExist(err) {
// Either staging died between creating the directory tree and writing the
// journal — no rename had run, so payload holds nothing — or commit/rollback
// failed part-way through removing the root, in which case payload only holds
// copies already inside the package. Without a journal there is nothing to
// roll back, and reporting an error here keeps the archive worker from ever
// starting, so discard the directory instead.
if err := os.RemoveAll(root); err != nil {
errs = append(errs, err)
}
continue
}
if err != nil {
errs = append(errs, err)
continue
}
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
}
manifestFile, err := os.OpenFile(
filepath.Join(payloadDir, "manifest.json"),
os.O_CREATE|os.O_TRUNC|os.O_WRONLY,
archiveFileMode,
)
if err != nil {
return 0, 0, "", err
}
encodeErr := json.NewEncoder(manifestFile).Encode(snapshot)
var syncErr error
if encodeErr == nil {
syncErr = manifestFile.Sync()
}
if err := errors.Join(encodeErr, syncErr, manifestFile.Close()); err != nil {
return 0, 0, "", err
}
temporary := path + ".partial"
_ = 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
}
+223
View File
@@ -0,0 +1,223 @@
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)
streamPath := filepath.Join(stage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))
if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil {
t.Fatal(err)
}
streamBody := []byte("{\"id\":1,\"raw_request\":\"large\"}\n")
if err := os.WriteFile(streamPath, streamBody, archiveFileMode); err != nil {
t.Fatal(err)
}
snapshot := &pgdb.TaskArchiveSnapshot{
FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID,
StreamedTables: map[string]string{"llm_records": pgdb.TaskArchiveLLMRecordsPath},
}
original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot)
if err != nil {
t.Fatal(err)
}
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)
}
if got, err := os.ReadFile(filepath.Join(extracted, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))); err != nil || !bytes.Equal(got, streamBody) {
t.Fatalf("streamed LLM archive payload=%q err=%v", got, err)
}
installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1)
if err != nil {
t.Fatal(err)
}
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)
}
}
+707
View File
@@ -0,0 +1,707 @@
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,请先暂停后重试归档")
}
task, err := s.m.pg.GetTask(job.TaskID)
if err != nil {
return err
}
if task == nil {
return pgdb.ErrTaskArchiveNotFound
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 10)
fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, task.ExplorationID)
if err != nil {
return err
}
defer func() {
if !committed {
runErr = errors.Join(runErr, fileStage.rollback())
}
}()
streamPath := filepath.Join(fileStage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))
if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil {
return err
}
streamFile, err := os.OpenFile(streamPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode)
if err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 20)
snapshot, snapshotErr := s.m.pg.SnapshotTaskArchiveWithLLMRecords(job.TaskID, streamFile)
var syncErr error
if snapshotErr == nil {
syncErr = streamFile.Sync()
}
closeErr := streamFile.Close()
if err := errors.Join(snapshotErr, syncErr, closeErr); err != nil {
return err
}
if snapshot.ExplorationID != task.ExplorationID {
return errors.New("任务探索记录在归档快照期间发生变化")
}
if s.m.traffic != nil && len(snapshot.Hosts) > 0 {
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_traffic", 38)
trafficCount, err := s.m.traffic.ExportHosts(snapshot.Hosts, filepath.Join(fileStage.payload, "traffic"))
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 || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
return pgdb.ErrTaskArchiveFormatMismatch
}
var llmRecordsFile *os.File
relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"]
if len(snapshot.StreamedTables) > 0 &&
(snapshot.FormatVersion < 2 || !hasStreamedLLMRecords || len(snapshot.StreamedTables) != 1 ||
relative != pgdb.TaskArchiveLLMRecordsPath) {
return pgdb.ErrTaskArchiveFormatMismatch
}
if hasStreamedLLMRecords {
llmRecordsFile, err = os.Open(filepath.Join(extracted, filepath.FromSlash(relative)))
if err != nil {
return err
}
defer func() {
if llmRecordsFile != nil {
_ = llmRecordsFile.Close()
}
}()
}
if s.m.traffic != nil {
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_traffic", 35)
if _, err := s.m.traffic.ImportArchive(filepath.Join(extracted, "traffic")); err != nil {
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)
var warnings []string
if llmRecordsFile != nil {
warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords(
job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
)
closeErr := llmRecordsFile.Close()
llmRecordsFile = nil
err = errors.Join(err, closeErr)
} else {
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
}
if err != nil {
return err
}
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())
}
+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())
}
}
+17
View File
@@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"iter"
"strings"
"sync"
"testing"
"time"
@@ -169,6 +170,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"}
+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
}
+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()
}
+227 -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,34 @@ 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})
}
// 归档路径即使一个历史 host 目录都没有(新装机的流量只落在 SQLite + _blobs)
// 也必须留下 journal:崩溃点若落在 PostgreSQL 提交与 SQLite 提交之间,重启后
// SQLite 事务被回滚,只有这份 journal 能让恢复流程补做 host 行的删除。少了它,
// 已转冷任务的独占流量会永久留在热库里。
uniqueHosts := uniqueArchiveHosts(hosts)
if len(planned) == 0 && len(uniqueHosts) == 0 {
return "", nil, nil
}
parent := filepath.Join(t.dir, "_delete_staging")
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: uniqueHosts}
for i := range planned {
planned[i].staged = filepath.Join(stageDir, fmt.Sprintf("%d-%s", i, filepath.Base(planned[i].source)))
journal.Moves = append(journal.Moves, hostDeleteStageMove{Source: planned[i].source, Staged: planned[i].staged})
}
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 +945,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 +1050,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 +1109,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 {
+94 -67
View File
@@ -81,6 +81,7 @@ function fmtDuration(ms: number): string {
// then one more page each time the user scrolls to the top. Kept modest so a long
// thread stays snappy (only ~a page of rows is in the DOM until you scroll up).
const HISTORY_PAGE = 200;
const CONVERSATION_LIST_PAGE = 100;
function conversationIsPinned(conversation: Conversation): boolean {
return conversation.pinned ?? Boolean(conversation.pinned_at);
@@ -750,7 +751,7 @@ function ChatView({
// ConversationItem is one row in the left rail: title, agent subtitle, inline
// rename, pin marker, and a compact action menu.
function ConversationItem({
const ConversationItem = React.memo(function ConversationItem({
conv,
agent,
active,
@@ -772,16 +773,16 @@ function ConversationItem({
active: boolean;
renaming: boolean;
renameText: string;
onSelect: () => void;
onStartRename: () => void;
onSelect: (id: number) => void;
onStartRename: (conversation: Conversation) => void;
onRenameText: (v: string) => void;
onCommitRename: () => void;
onCommitRename: (id: number, title: string) => void;
onCancelRename: () => void;
onTogglePinned: () => void;
onDelete: () => void;
onTogglePinned: (conversation: Conversation) => void;
onDelete: (id: number) => void;
selectionMode: boolean;
selectedForDelete: boolean;
onSelectedForDeleteChange: (checked: boolean) => void;
onSelectedForDeleteChange: (id: number, checked: boolean) => void;
}) {
const [deleteOpen, setDeleteOpen] = React.useState(false);
const renameInputRef = React.useRef<HTMLInputElement>(null);
@@ -805,7 +806,7 @@ function ConversationItem({
{selectionMode && (
<Checkbox
checked={selectedForDelete}
onCheckedChange={(checked) => onSelectedForDeleteChange(checked === true)}
onCheckedChange={(checked) => onSelectedForDeleteChange(conv.id, checked === true)}
aria-label={`选择对话「${conv.title || "新对话"}」`}
className="ml-1 shrink-0"
/>
@@ -816,10 +817,10 @@ function ConversationItem({
value={renameText}
onChange={(e) => onRenameText(e.target.value)}
onBlur={() => {
if (!cancelRenameRef.current) onCommitRename();
if (!cancelRenameRef.current) onCommitRename(conv.id, renameText);
}}
onKeyDown={(e) => {
if (e.key === "Enter") onCommitRename();
if (e.key === "Enter") onCommitRename(conv.id, renameText);
if (e.key === "Escape") {
e.preventDefault();
cancelRenameRef.current = true;
@@ -831,8 +832,8 @@ function ConversationItem({
) : (
<button
type="button"
onClick={onSelect}
onDoubleClick={onStartRename}
onClick={() => onSelect(conv.id)}
onDoubleClick={() => onStartRename(conv)}
title="双击重命名"
className="min-w-0 flex-1 rounded-md px-2 py-1.5 text-left"
>
@@ -869,11 +870,11 @@ function ConversationItem({
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuGroup>
<DropdownMenuItem onSelect={onStartRename}>
<DropdownMenuItem onSelect={() => onStartRename(conv)}>
<PencilIcon />
重命名
</DropdownMenuItem>
<DropdownMenuItem onSelect={onTogglePinned}>
<DropdownMenuItem onSelect={() => onTogglePinned(conv)}>
{pinned ? <PinOffIcon /> : <PinIcon />}
{pinned ? "取消置顶" : "置顶"}
</DropdownMenuItem>
@@ -895,13 +896,13 @@ function ConversationItem({
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>取消</AlertDialogCancel>
<AlertDialogAction onClick={onDelete}>删除</AlertDialogAction>
<AlertDialogAction onClick={() => onDelete(conv.id)}>删除</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}
});
export default function ChatPage() {
const [agents, setAgents] = React.useState<Agent[]>([]);
@@ -917,6 +918,7 @@ export default function ChatPage() {
const [selectionMode, setSelectionMode] = React.useState(false);
const [bulkDeleteOpen, setBulkDeleteOpen] = React.useState(false);
const [bulkDeleting, setBulkDeleting] = React.useState(false);
const [visibleConversationCount, setVisibleConversationCount] = React.useState(CONVERSATION_LIST_PAGE);
// Pending text + uploaded attachments handed off from a draft that created a
// conversation via the paperclip, keyed by the new conversation id (consumed once
// by ChatView on mount; ids never repeat, so leftover entries are harmless).
@@ -974,11 +976,16 @@ export default function ChatPage() {
window.history.replaceState(null, "", url);
}, [selectedId]);
const selected = convs.find((c) => c.id === selectedId) ?? null;
const selected = React.useMemo(() => convs.find((c) => c.id === selectedId) ?? null, [convs, selectedId]);
const agentByKey = React.useMemo(() => new Map(agents.map((agent) => [agent.key, agent])), [agents]);
const visibleConversations = React.useMemo(
() => convs.slice(0, visibleConversationCount),
[convs, visibleConversationCount],
);
// conversation agents: custom agents + conversational built-ins (role=assistant,
// e.g. Auto / 渗透测试). The orchestration built-ins (goals/planner/mainagent/worker)
// are task-specific and stay hidden from the chat page.
const chatAgents = agents.filter((a) => !a.builtin || a.role === "assistant");
const chatAgents = React.useMemo(() => agents.filter((a) => !a.builtin || a.role === "assistant"), [agents]);
const selectedConversationCount = selectedConversationIds.size;
const allConversationsSelected = convs.length > 0 && selectedConversationCount === convs.length;
@@ -987,14 +994,14 @@ export default function ChatPage() {
if (allConversationsSelected) conversationHeaderChecked = true;
else if (someConversationsSelected) conversationHeaderChecked = "indeterminate";
function toggleConversationSelected(id: number, checked: boolean) {
const toggleConversationSelected = React.useCallback((id: number, checked: boolean) => {
setSelectedConversationIds((current) => {
const next = new Set(current);
if (checked) next.add(id);
else next.delete(id);
return next;
});
}
}, []);
function toggleAllConversations(checked: boolean) {
setSelectedConversationIds(checked ? new Set(convs.map((conversation) => conversation.id)) : new Set());
@@ -1005,21 +1012,24 @@ export default function ChatPage() {
setSelectedConversationIds(new Set());
}
async function del(id: number) {
try {
await api.deleteConversation(id);
if (selectedId === id) setSelectedId(null);
setSelectedConversationIds((current) => {
if (!current.has(id)) return current;
const next = new Set(current);
next.delete(id);
return next;
});
reloadConvs();
} catch (e) {
toast.error("删除失败:" + (e as Error).message);
}
}
const del = React.useCallback(
async (id: number) => {
try {
await api.deleteConversation(id);
setSelectedId((current) => (current === id ? null : current));
setSelectedConversationIds((current) => {
if (!current.has(id)) return current;
const next = new Set(current);
next.delete(id);
return next;
});
reloadConvs();
} catch (e) {
toast.error("删除失败:" + (e as Error).message);
}
},
[reloadConvs],
);
async function deleteSelectedConversations() {
const ids = [...selectedConversationIds];
@@ -1062,32 +1072,38 @@ export default function ChatPage() {
}
}
async function togglePinned(conversation: Conversation) {
const pinned = conversationIsPinned(conversation);
try {
await api.pinConversation(conversation.id, !pinned);
reloadConvs();
} catch (e) {
toast.error(`${pinned ? "取消置顶" : "置顶"}失败:${(e as Error).message}`);
}
}
const togglePinned = React.useCallback(
async (conversation: Conversation) => {
const pinned = conversationIsPinned(conversation);
try {
await api.pinConversation(conversation.id, !pinned);
reloadConvs();
} catch (e) {
toast.error(`${pinned ? "取消置顶" : "置顶"}失败:${(e as Error).message}`);
}
},
[reloadConvs],
);
function startRename(c: Conversation) {
const startRename = React.useCallback((c: Conversation) => {
setRenamingId(c.id);
setRenameText(c.title || "");
}
async function commitRename() {
const id = renamingId;
const title = renameText.trim();
setRenamingId(null);
if (!id || !title) return;
try {
await api.renameConversation(id, title);
reloadConvs();
} catch (e) {
toast.error("重命名失败:" + (e as Error).message);
}
}
}, []);
const commitRename = React.useCallback(
async (id: number, value: string) => {
const title = value.trim();
setRenamingId(null);
if (!title) return;
try {
await api.renameConversation(id, title);
reloadConvs();
} catch (e) {
toast.error("重命名失败:" + (e as Error).message);
}
},
[reloadConvs],
);
const cancelRename = React.useCallback(() => setRenamingId(null), []);
return (
<div
@@ -1147,26 +1163,37 @@ export default function ChatPage() {
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<div className="flex min-w-0 flex-col gap-0.5 p-2">
{convs.length === 0 && <p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>}
{convs.map((c) => (
{visibleConversations.map((c) => (
<ConversationItem
key={c.id}
conv={c}
agent={agents.find((a) => a.key === c.agent_key)}
agent={agentByKey.get(c.agent_key)}
active={selectedId === c.id}
renaming={renamingId === c.id}
renameText={renameText}
onSelect={() => setSelectedId(c.id)}
onStartRename={() => startRename(c)}
renameText={renamingId === c.id ? renameText : ""}
onSelect={setSelectedId}
onStartRename={startRename}
onRenameText={setRenameText}
onCommitRename={commitRename}
onCancelRename={() => setRenamingId(null)}
onTogglePinned={() => void togglePinned(c)}
onDelete={() => del(c.id)}
onCancelRename={cancelRename}
onTogglePinned={togglePinned}
onDelete={del}
selectionMode={selectionMode}
selectedForDelete={selectedConversationIds.has(c.id)}
onSelectedForDeleteChange={(checked) => toggleConversationSelected(c.id, checked)}
onSelectedForDeleteChange={toggleConversationSelected}
/>
))}
{visibleConversationCount < convs.length && (
<Button
type="button"
variant="ghost"
size="sm"
className="mt-1 w-full"
onClick={() => setVisibleConversationCount((count) => count + CONVERSATION_LIST_PAGE)}
>
加载更多
</Button>
)}
</div>
</ScrollArea>
</div>
+48 -14
View File
@@ -148,13 +148,46 @@ export default function DashboardPage() {
const [tools, setTools] = React.useState<Tool[]>([]);
const [llmProfiles, setLLMProfiles] = React.useState<LLMProfile[]>([]);
// fast poll: tasks, findings, stats, pending, activity (every 5s)
// The task list is the most expensive dashboard source. Poll it independently
// so a large history cannot hold back every other dashboard panel.
React.useEffect(() => {
let alive = true;
let loading = false;
let signature = "";
const load = async () => {
if (loading) return;
loading = true;
try {
const [tr, fr, sr, sets, pr, act, tok, ctok] = await Promise.all([
api.tasks(),
const tr = await api.tasks();
if (!alive) return;
const nextSignature = JSON.stringify(tr.tasks);
if (nextSignature !== signature) {
signature = nextSignature;
setTasks(tr.tasks);
}
} catch {
/* transient errors — next poll retries */
} finally {
loading = false;
}
};
void load();
const t = setInterval(load, 5000);
return () => {
alive = false;
clearInterval(t);
};
}, []);
// The remaining fast sources stay batched so one poll causes a single render.
React.useEffect(() => {
let alive = true;
let loading = false;
const load = async () => {
if (loading) return;
loading = true;
try {
const [findings, stats, settings, pending, activity, tokens, conversationTokens] = await Promise.all([
api.findings(),
api.stats(),
api.settings(),
@@ -164,23 +197,24 @@ export default function DashboardPage() {
api.conversationTokens(),
]);
if (!alive) return;
setTasks(tr.tasks);
setFindings(fr);
setStats(sr);
setSettings(sets);
setPending(pr);
setActivity(act.items);
setTokens(tok.total ?? null);
setConvTokens(ctok);
setFindings(findings);
setStats(stats);
setSettings(settings);
setPending(pending);
setActivity(activity.items);
setTokens(tokens.total ?? null);
setConvTokens(conversationTokens);
} catch {
/* transient errors — next poll retries */
// Preserve stale data and retry on the next interval.
} finally {
loading = false;
}
};
void load();
const t = setInterval(load, 5000);
const timer = setInterval(load, 5000);
return () => {
alive = false;
clearInterval(t);
clearInterval(timer);
};
}, []);
+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}
@@ -56,7 +56,7 @@ function scopeValue(row: TaskScopeRow): string {
if (row.value) return row.value;
if (row.domain) return row.domain;
if (row.net) return row.net;
if (row.company_id) return row.company_name || `企业 #${row.company_id}`;
if (row.company_id) return row.company_name?.trim() ? row.company_name : `企业 #${row.company_id}`;
return "—";
}
@@ -346,17 +346,32 @@ export function OverviewTab({ taskId }: { taskId: string }) {
React.useEffect(() => {
let cancelled = false;
let loading = false;
const load = async () => {
if (loading) return;
loading = true;
try {
const [tasksResp, statsResp, intentsResp, findingsResp] = await Promise.all([
api.tasks(),
const [taskResp, statsResp, intentsResp, findingsResp] = await Promise.all([
api.task(taskId),
api.stats(taskId),
api.intents(taskId),
api.findings(taskId),
]);
if (cancelled) return;
setTask(tasksResp.tasks.find((t) => t.id === taskId) ?? null);
const activeTask = statsResp.active_task;
setTask(
activeTask
? {
...taskResp,
in_flight: activeTask.in_flight,
goals_total: activeTask.goals_total,
goals_met: activeTask.goals_met,
engine_mode: statsResp.engine_mode ?? activeTask.engine_mode,
paused: activeTask.paused,
}
: taskResp,
);
setStats(statsResp);
setIntents(intentsResp);
setFindings(findingsResp);
@@ -372,10 +387,12 @@ export function OverviewTab({ taskId }: { taskId: string }) {
});
} catch {
// transient errors are ignored; the next poll will retry
} finally {
loading = false;
}
};
load();
void load();
void loadScope();
void loadTokens();
void loadGoals();
+107 -39
View File
@@ -1,22 +1,27 @@
"use client";
import * as React from "react";
import Link from "next/link";
import { useSearchParams } from "next/navigation";
import { toast } from "sonner";
import {
ArrowLeftIcon,
PauseIcon,
PlayIcon,
BrainIcon,
CheckIcon,
CircleAlertIcon,
} from "lucide-react";
import Link from "next/link";
import { useRouter, useSearchParams } from "next/navigation";
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 { SidebarTrigger } from "@/components/ui/sidebar";
import { Separator } from "@/components/ui/separator";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
AlertDialogTrigger,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
Popover,
@@ -26,30 +31,32 @@ import {
PopoverTitle,
PopoverTrigger,
} from "@/components/ui/popover";
import { Separator } from "@/components/ui/separator";
import { SidebarTrigger } from "@/components/ui/sidebar";
import { Spinner } from "@/components/ui/spinner";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { StatusBadge } from "@/components/status-badge";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { api } from "@/lib/api";
import type { LLMProfile, Task } from "@/lib/types";
import { SessionsTab } from "./_tabs/sessions-tab";
import { OverviewTab } from "./_tabs/overview-tab";
import { GraphTab } from "./_tabs/graph-tab";
import { FindingsTab } from "./_tabs/findings-tab";
import { ReportTab } from "./_tabs/report-tab";
import { InterceptTab } from "./_tabs/intercept-tab";
import { AssetsTab } from "./_tabs/assets-tab";
import { CoverageGraphTab } from "./_tabs/coverage-graph-tab";
import { FindingsTab } from "./_tabs/findings-tab";
import { GraphTab } from "./_tabs/graph-tab";
import { InterceptTab } from "./_tabs/intercept-tab";
import { OverviewTab } from "./_tabs/overview-tab";
import { ReportTab } from "./_tabs/report-tab";
import { SessionsTab } from "./_tabs/sessions-tab";
const TABS = [
{ value: "sessions", label: "会话" },
{ value: "overview", label: "总览" },
{ value: "graph", label: "探索链路" },
{ value: "findings", label: "发现" },
{ value: "assets", label: "测试资产" },
{ value: "coverage", label: "资产覆盖图" },
{ value: "sessions", label: "会话" },
{ value: "overview", label: "总览" },
{ value: "graph", label: "探索链路" },
{ value: "findings", label: "发现" },
{ value: "assets", label: "测试资产" },
{ value: "coverage", label: "资产覆盖图" },
{ value: "intercept", label: "拦截审批" },
{ value: "report", label: "报告" },
{ value: "report", label: "报告" },
];
function taskProfileIDs(task: Task): string[] {
@@ -81,7 +88,7 @@ function TaskLLMControl({ task, profiles, onUpdated }: { task: Task; profiles: L
const activeProfile = profiles.find((profile) => profile.id === activeID);
const currentLabel = exhausted
? "配置链已耗尽"
: activeProfile?.name ?? (activeID ? `配置 #${activeID}` : "跟随默认配置");
: (activeProfile?.name ?? (activeID ? `配置 #${activeID}` : "跟随默认配置"));
const activeIndex = chain.indexOf(activeID);
const backupCount = activeIndex >= 0 ? Math.max(0, chain.length - activeIndex - 1) : 0;
const currentTitle = [currentLabel, activeProfile?.model, backupCount > 0 ? `${backupCount} 个备用` : ""]
@@ -151,11 +158,7 @@ function TaskLLMControl({ task, profiles, onUpdated }: { task: Task; profiles: L
{backupCount > 0 && <span className="hidden text-muted-foreground xl:inline">+{backupCount}</span>}
</Button>
</PopoverTrigger>
<PopoverContent
ref={popoverContentRef}
align="start"
className="w-[min(28rem,calc(100vw-2rem))] gap-4 p-4"
>
<PopoverContent ref={popoverContentRef} align="start" className="w-[min(28rem,calc(100vw-2rem))] gap-4 p-4">
<PopoverHeader>
<PopoverTitle>任务 LLM 配置链</PopoverTitle>
<PopoverDescription>{editorDescription}</PopoverDescription>
@@ -199,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);
@@ -207,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
@@ -234,12 +239,16 @@ function TaskDetailInner() {
};
}, [id]);
const taskLoadInFlight = React.useRef<string | null>(null);
const load = React.useCallback(() => {
Promise.all([api.tasks(), api.stats(id).catch(() => null)])
.then(([r, s]) => {
const base = r.tasks.find((t) => t.id === id) ?? null;
if (taskLoadInFlight.current === id) return;
taskLoadInFlight.current = id;
Promise.all([api.task(id), api.stats(id).catch(() => null)])
.then(([task, s]) => {
if (taskLoadInFlight.current !== id) return;
const base = { ...task };
const at = s?.active_task;
if (base && at) {
if (at) {
base.in_flight = at.in_flight;
base.goals_total = at.goals_total;
base.goals_met = at.goals_met;
@@ -247,12 +256,17 @@ function TaskDetailInner() {
base.paused = at.paused;
}
setTask(base);
setPaused(at?.paused ?? base?.paused ?? false);
setPaused(at?.paused ?? base.paused ?? false);
})
.catch(() => {
// Keep the last rendered task state during a transient poll failure.
})
.finally(() => setLoaded(true));
.finally(() => {
if (taskLoadInFlight.current === id) {
taskLoadInFlight.current = null;
setLoaded(true);
}
});
}, [id]);
React.useEffect(() => {
load();
@@ -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}
File diff suppressed because it is too large. Load diff
+50 -17
View File
@@ -3,6 +3,8 @@
import * as React from "react";
import {
ArrowDownWideNarrowIcon,
ArrowUpNarrowWideIcon,
CheckIcon,
ChevronLeftIcon,
ChevronRightIcon,
@@ -39,7 +41,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { api } from "@/lib/api";
import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types";
import { cn } from "@/lib/utils";
import { cn, copyText } from "@/lib/utils";
function fmtTime(ts: string) {
return new Date(ts).toLocaleString("zh-CN", {
@@ -199,12 +201,12 @@ function HttpCodeBlock({ raw }: { raw: string }) {
}, [copied]);
const copyPacket = async () => {
try {
await navigator.clipboard.writeText(value);
const ok = await copyText(value);
if (ok) {
setCopied(true);
} catch {
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
return;
}
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
};
const renderLine = (line: string, index: number) => {
@@ -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}`}
+8 -6
View File
@@ -545,7 +545,8 @@ function ProfileSheet({
onChange={(e) => setProxy(e.target.value)}
/>
<p className="text-muted-foreground text-xs">
仅 LLM 出站请求走此代理,支持 http/https/socks5,可带账号密码(如 socks5://user:pass@host:port,密码含特殊字符需 URL 编码);留空表示不使用代理(直连)。
仅 LLM 出站请求走此代理,支持 http/https/socks5,可带账号密码(如
socks5://user:pass@host:port,密码含特殊字符需 URL 编码);留空表示不使用代理(直连)。
</p>
</div>
@@ -620,7 +621,8 @@ function ProfileSheet({
<Label className="text-sm">流式输出 · streaming</Label>
<p className="text-muted-foreground text-xs">
开启(默认)走流式 SSE,有运行中实时进度与实时 token 计数。 关闭则走真·非流式(stream:false,
一次性返回完整响应)——可绕开部分网关糟糕的 SSE 实现(空帧 / 思考字段丢帧), 代价是失去运行中的实时进度。
一次性返回完整响应)——可绕开部分网关糟糕的 SSE 实现(空帧 / 思考字段丢帧),
代价是失去运行中的实时进度。
</p>
</div>
<Switch checked={streaming} onCheckedChange={setStreaming} aria-label="流式输出" />
@@ -632,8 +634,8 @@ function ProfileSheet({
<div className="grid gap-0.5">
<Label className="text-sm">思考开关 · thinking.type</Label>
<p className="text-muted-foreground text-xs">
控制是否发送 thinking 字段。不发送=不带该字段(兼容 MiniMax 等不支持
的模型);关闭=发 disabled;开启=发 enabled。与下面的强度互相独立。
控制是否发送 thinking 字段。不发送=不带该字段(兼容 MiniMax 等不支持 的模型);关闭=发
disabled;开启=发 enabled。与下面的强度互相独立。
</p>
</div>
<Select value={thinkingType} onValueChange={setThinkingType}>
@@ -653,8 +655,8 @@ function ProfileSheet({
<div className="grid gap-0.5">
<Label className="text-sm">思考强度 · reasoning_effort</Label>
<p className="text-muted-foreground text-xs">
独立的强度档位(OpenAI reasoning_effort / Anthropic output_config.effort)。
有些接口没有 thinking 字段、只靠强度即可激活思考,故可单独设置、不发送思考开关。
独立的强度档位(OpenAI reasoning_effort / Anthropic output_config.effort)。 有些接口没有 thinking
字段、只靠强度即可激活思考,故可单独设置、不发送思考开关。
</p>
</div>
<Select value={effort} onValueChange={setEffort}>
+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>
);
}
+34 -2
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,
@@ -189,6 +192,7 @@ export const api = {
// ---- tasks ----
tasks: () =>
get<{ tasks: Task[]; active: string }>("/tasks").then((r) => ({ tasks: arr(r.tasks), active: r.active ?? "" })),
task: (id: string) => get<Task>(`/tasks/${encodeURIComponent(id)}`),
createTask: (input: {
name?: string;
categoryId?: number;
@@ -253,6 +257,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;
@@ -284,7 +309,8 @@ export const api = {
usageStats: (days = 365) => get<UsageStats>(`/tokens/usage?days=${days}`),
// ---- 目标管理(总览)----
// 本任务全部目标(text/vulnclass/state)。
taskGoals: (id: string) => get<{ goals: TaskGoal[] }>(`/tasks/${id}/goals`),
taskGoals: (id: string) =>
get<{ goals: TaskGoal[] | null }>(`/tasks/${id}/goals`).then((response) => ({ goals: arr(response.goals) })),
// 人工新增目标:写入图谱并通知 planner、复活任务。
addGoal: (id: string, text: string, vulnclass?: string) =>
post<TaskGoal>(`/tasks/${id}/goals`, { text, vulnclass: vulnclass ?? "" }),
@@ -296,7 +322,10 @@ export const api = {
// ---- 操作约束管理(总览)----
// 本任务全部操作约束(allow/deny)。
taskConstraints: (id: string) => get<{ constraints: TaskConstraint[] }>(`/tasks/${id}/constraints`),
taskConstraints: (id: string) =>
get<{ constraints: TaskConstraint[] | null }>(`/tasks/${id}/constraints`).then((response) => ({
constraints: arr(response.constraints),
})),
// 新增约束(不通知 planner,下一轮规划自然读到)。
addConstraint: (id: string, text: string, kind: TaskConstraint["kind"]) =>
post<TaskConstraint>(`/tasks/${id}/constraints`, { text, kind }),
@@ -425,6 +454,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()}`);
},
@@ -434,6 +464,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()}`);
},
@@ -455,6 +486,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();
+382 -17
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,11 +42,235 @@ 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">;
const mockTaskAssetSources = new Map<string, MockTaskAssetSource>();
let nextMockTaskAssetID = D.tasks.length + 1;
let mockActiveTask = D.ACTIVE_TASK;
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);
}
function mockAssetForTask(taskID: string, asset: Asset): Asset {
const source = mockTaskAssetSources.get(mockTaskAssetSourceKey(taskID, asset.id));
return source ? { ...asset, ...source } : asset;
}
function deleteMockTaskAssetSources(taskID: string, assetID?: number) {
if (assetID !== undefined) {
mockTaskAssetSources.delete(mockTaskAssetSourceKey(taskID, assetID));
return;
}
for (const key of mockTaskAssetSources.keys()) {
const [linkedTaskID] = JSON.parse(key) as [string, number];
if (linkedTaskID === taskID) mockTaskAssetSources.delete(key);
}
}
function mockTaskAssetID(taskID: string): number | undefined {
const numeric = Number(taskID);
if (Number.isInteger(numeric) && numeric > 0) return numeric;
@@ -72,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,
@@ -277,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;
@@ -289,7 +626,10 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
const now = new Date();
const profileIDs = [...((b.llm_profile_ids as number[] | undefined) ?? [])];
const sourceTaskIDs = [...((b.source_task_ids as string[] | undefined) ?? [])];
const companyIDs = [...((b.company_ids as number[] | undefined) ?? [])];
const companyIDs = [...new Set((b.company_ids as number[] | undefined) ?? [])];
if (companyIDs.some((companyID) => !mockCompanies.some((company) => company.id === companyID))) {
throw new Error("关联企业不存在或无效");
}
const categoryID = typeof b.category_id === "number" ? b.category_id : undefined;
const category = categoryID === undefined ? undefined : mockTaskCategories.find((item) => item.id === categoryID);
if (categoryID !== undefined && !category) throw new Error("任务分类不存在");
@@ -318,7 +658,19 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
source_task_ids: sourceTaskIDs,
company_ids: companyIDs,
};
mockTaskAssetIDs.set(id, nextMockTaskAssetID++);
const numericTaskID = nextMockTaskAssetID++;
mockTaskAssetIDs.set(id, numericTaskID);
const selectedCompanies = new Set(companyIDs);
for (const asset of mockAssets) {
if (asset.company_id === undefined || !selectedCompanies.has(asset.company_id)) continue;
if (!asset.task_ids.includes(numericTaskID)) asset.task_ids.push(numericTaskID);
const company = mockCompanies.find((candidate) => candidate.id === asset.company_id);
setMockTaskAssetSource(id, asset.id, {
task_source: "company",
task_source_summary: `任务创建时关联企业:${company?.name ?? `#${asset.company_id}`}`,
task_source_node_id: undefined,
});
}
for (const item of mockTasks) item.active = false;
mockTasks.unshift(created);
mockActiveTask = id;
@@ -449,6 +801,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (index >= 0) mockTaskTemplates.splice(index, 1);
return { deleted: id };
}
if (seg[0] === "tasks" && seg.length === 2 && m === "GET") {
const task = mockTasks.find((item) => item.id === seg[1]);
if (!task) throw new Error("任务不存在");
return publicMockTask(task);
}
if (seg[0] === "tasks" && seg.length === 2 && m === "PATCH") {
const task = mockTasks.find((item) => item.id === seg[1]);
if (!task) throw new Error("任务不存在");
@@ -462,9 +819,14 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
}
if (seg[0] === "tasks" && seg.length === 2 && m === "DELETE") {
const id = seg[1];
const numericTaskID = mockTaskAssetID(id);
const index = mockTasks.findIndex((item) => item.id === id);
if (index >= 0) mockTasks.splice(index, 1);
mockTaskAssetIDs.delete(id);
deleteMockTaskAssetSources(id);
if (numericTaskID !== undefined) {
for (const asset of mockAssets) asset.task_ids = asset.task_ids.filter((taskID) => taskID !== numericTaskID);
}
let findingsDeleted = 0;
if (b.delete_findings) {
@@ -701,7 +1063,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
});
const limit = Number(q.get("limit") ?? 50);
const offset = Number(q.get("offset") ?? 0);
const page = list.slice(offset, offset + limit);
const page = list.slice(offset, offset + limit).map((asset) => (taskID ? mockAssetForTask(taskID, asset) : asset));
return { count: page.length, total: list.length, assets: page };
}
if (path === "/assets" && m === "DELETE") {
@@ -792,9 +1154,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
asset.task_ids.push(numericTaskID);
mutation.assets_linked++;
}
asset.task_source = "manual";
asset.task_source_summary = "用户在测试资产页手工新增";
asset.task_source_node_id = undefined;
setMockTaskAssetSource(seg[1], asset.id, {
task_source: "manual",
task_source_summary: "用户在测试资产页手工新增",
task_source_node_id: undefined,
});
}
mockTaskScopes.set(seg[1], currentScopes);
return mutation;
@@ -812,9 +1176,11 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
asset.task_ids.push(numericTaskID);
mutation.attached++;
}
asset.task_source = "manual";
asset.task_source_summary = sourceSummary;
asset.task_source_node_id = undefined;
setMockTaskAssetSource(seg[1], asset.id, {
task_source: "manual",
task_source_summary: sourceSummary,
task_source_node_id: undefined,
});
}
return mutation;
}
@@ -824,11 +1190,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (numericTaskID === undefined || !asset) throw new Error("任务或资产不存在");
if (!asset.task_ids.includes(numericTaskID)) throw new Error("资产未关联当前任务");
asset.task_ids = asset.task_ids.filter((id) => id !== numericTaskID);
if (asset.task_ids.length === 0) {
asset.task_source = undefined;
asset.task_source_summary = undefined;
asset.task_source_node_id = undefined;
}
deleteMockTaskAssetSources(seg[1], asset.id);
return { detached: asset.id };
}
if (seg[0] === "tasks" && seg[2] === "intent-assets" && seg.length === 3 && m === "GET") {
@@ -843,8 +1205,9 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
}
const sourceTaskID = mockTaskAssetID(seg[1]) ?? 0;
const assets: IntentAsset[] = mappings.flatMap((mapping) => {
const asset = mockAssets.find((item) => item.id === mapping.assetID);
if (!asset) return [];
const storedAsset = mockAssets.find((item) => item.id === mapping.assetID);
if (!storedAsset) return [];
const asset = mockAssetForTask(seg[1], storedAsset);
return [
{
intent_id: mapping.intentID,
@@ -946,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);
@@ -1056,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);
+1
View File
@@ -13,6 +13,7 @@ const TASK_ASSET_SOURCE_LABELS: Record<string, string> = {
agent: "Agent 发现",
anchor: "黑板锚点",
api: "资产 API",
company: "企业关联",
legacy: "历史关联",
manual: "人工加入",
system: "系统关联",
+71 -1
View File
@@ -35,7 +35,8 @@ export interface Task {
llm_failover_state?: "default" | "ready" | "chain_exhausted" | string;
llm_failover_reason?: string;
source_task_ids?: string[]; // directly related tasks inherited as read-only context
company_ids?: number[]; // associated company asset scopes available as task context
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";
}