Merge pull request #79 from neouks/codex/sync-v0.3.6
feat: 增加任务冷归档与大任务性能优化
This commit is contained in:
52 files changed
+7366
-511
No files matched your search
@@ -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
@@ -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
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
{
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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).
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"`
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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"}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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);
|
||||
};
|
||||
}, []);
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
@@ -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}`}
|
||||
|
||||
@@ -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}>
|
||||
|
||||
@@ -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
@@ -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
@@ -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);
|
||||
|
||||
@@ -13,6 +13,7 @@ const TASK_ASSET_SOURCE_LABELS: Record<string, string> = {
|
||||
agent: "Agent 发现",
|
||||
anchor: "黑板锚点",
|
||||
api: "资产 API",
|
||||
company: "企业关联",
|
||||
legacy: "历史关联",
|
||||
manual: "人工加入",
|
||||
system: "系统关联",
|
||||
|
||||
+71
-1
@@ -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";
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user