This commit is contained in:
pphua committed 2026-08-28 12:56:29 +08:00
1 parent 594436fddd
commit 9da521ee4e
31 files changed
+5768 -267

No files matched your search

+14
View File
@@ -56,10 +56,24 @@ type TaskDTO struct {
LLMFailoverState string `json:"llm_failover_state"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []string `json:"source_task_ids"`
ArchiveBlockedBy string `json:"archive_blocked_by_task_id,omitempty"`
CompanyIDs []int64 `json:"company_ids"`
CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定)
}
func applyTaskArchiveBlocker(dto *TaskDTO, blockers map[int64]int64) {
if dto == nil || len(blockers) == 0 {
return
}
taskID, err := strconv.ParseInt(dto.ID, 10, 64)
if err != nil {
return
}
if dependentID := blockers[taskID]; dependentID > 0 {
dto.ArchiveBlockedBy = strconv.FormatInt(dependentID, 10)
}
}
// TokenTotalDTO is a whole-task (all agents) token aggregate.
type TokenTotalDTO struct {
InputTokens int `json:"input_tokens"`
+1
View File
@@ -37,6 +37,7 @@ func (s *Server) findingGroups(w http.ResponseWriter, r *http.Request) {
Status: normFilter(q.Get("status")),
VulnClass: normFilter(q.Get("vulnclass")),
TaskID: normFilter(q.Get("task_id")),
Query: q.Get("q"),
Sort: q.Get("sort"),
}
groups, total, findingTotal, err := s.m.pg.ListFindingGroups(filter, page, limit)
+22
View File
@@ -102,6 +102,28 @@ func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) {
t.Fatalf("unexpected first group: %+v", page.Items[0])
}
request = httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&q=finding+2&page=1&limit=10", nil)
recorder = httptest.NewRecorder()
s.findingGroups(recorder, request)
if recorder.Code != http.StatusOK {
t.Fatalf("searched groups status=%d body=%s", recorder.Code, recorder.Body.String())
}
page = struct {
Items []db.FindingGroup `json:"items"`
Total int `json:"total"`
FindingTotal int `json:"finding_total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
}{}
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
t.Fatal(err)
}
if len(page.Items) != 1 || page.Total != 1 || page.FindingTotal != 1 || page.Items[0].TaskID == nil ||
*page.Items[0].TaskID != mustTaskID(t, taskB.ID) {
t.Fatalf("unexpected searched group page: %+v", page)
}
request = httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&page=-4&limit=0", nil)
recorder = httptest.NewRecorder()
+19
View File
@@ -358,6 +358,25 @@ func NewManager(dir, proxyAddr string) (*Manager, error) {
if err != nil {
log.Printf("[traffic] disabled: %v", err)
} else {
err = tr.RecoverHostDeleteStages(func(_ int64, taskID int64) (bool, error) {
if taskID <= 0 {
return false, errors.New("归档流量暂存日志缺少任务 ID")
}
task, taskErr := pg.GetTask(taskID)
if taskErr != nil {
return false, taskErr
}
// Archived/permanently deleted tasks are hidden from GetTask. A
// restored task is visible and needs the staged traffic put back.
return task == nil, nil
})
if err != nil {
_ = tr.Close()
_ = pg.Close()
return nil, fmt.Errorf("recover traffic delete staging: %w", err)
}
}
if tr != nil {
m.traffic = tr
go func() {
log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr())
+24 -3
View File
@@ -116,6 +116,11 @@ type Server struct {
// the persisted current profile at every new LLM call.
taskAgentMu sync.Mutex
taskAgents map[string]*taskAgentBundle
// Cold task archives run through one persistent FIFO worker. The buffered wake
// channel coalesces enqueue bursts; the database remains the source of truth.
archiveWake chan struct{}
archiveWG sync.WaitGroup
}
// profBundle is a planner/worker pair built from one LLM profile.
@@ -151,7 +156,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{},
profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{},
provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg),
taskAgents: map[string]*taskAgentBundle{}}
taskAgents: map[string]*taskAgentBundle{}, archiveWake: make(chan struct{}, 1)}
// Every task uses a stable task router. An empty explicit chain is resolved by
// that router through Agent bindings and then the global provider, so adding a
// first chain to a running task takes effect on its very next LLM call.
@@ -264,6 +269,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
}
}
go s.reconcileConcurrency()
s.startTaskArchiveWorker()
s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定
return s
}
@@ -690,6 +696,14 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("PATCH /api/tasks/{id}", s.updateTaskMetadata)
mux.HandleFunc("PATCH /api/tasks/{id}/category", s.updateTaskCategory)
mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch)
mux.HandleFunc("GET /api/task-archives", s.listTaskArchives)
mux.HandleFunc("GET /api/task-archives/{id}", s.getTaskArchive)
mux.HandleFunc("POST /api/tasks/{id}/archive", s.queueTaskArchive)
mux.HandleFunc("POST /api/tasks/archive/batch", s.queueTaskArchivesBatch)
mux.HandleFunc("POST /api/task-archives/{id}/restore", s.queueTaskArchiveRestore)
mux.HandleFunc("POST /api/task-archives/restore/batch", s.restoreTaskArchivesBatch)
mux.HandleFunc("DELETE /api/task-archives/{id}", s.queueTaskArchiveDelete)
mux.HandleFunc("POST /api/task-archives/delete/batch", s.deleteTaskArchivesBatch)
mux.HandleFunc("GET /api/tasks/{id}/coverage", s.taskCoverage)
mux.HandleFunc("GET /api/tasks/{id}/coverage-graph", s.taskCoverageGraph)
mux.HandleFunc("GET /api/tasks/{id}/asset-refs", s.taskAssetRefs)
@@ -996,9 +1010,11 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
}
list := s.m.List()
metrics, _ := s.m.PG().TaskListMetricsAll()
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
dtos := make([]TaskDTO, 0, len(list))
for _, t := range list {
dto := taskDTO(t, s.resolvedTaskStatus(t))
applyTaskArchiveBlocker(&dto, archiveBlockers)
metric := metrics[t.ExpID]
dto.Tokens = tokenTotalDTO(metric.Tokens)
// prefer the live in-memory heartbeat (fresher) and fall back to the
@@ -1640,7 +1656,10 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
writeErr(w, 404, "task not found")
return
}
writeJSON(w, 200, taskDTO(t, s.resolvedTaskStatus(t)))
dto := taskDTO(t, s.resolvedTaskStatus(t))
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
applyTaskArchiveBlocker(&dto, archiveBlockers)
writeJSON(w, 200, dto)
}
// taskCoverage returns a task's rough asset test coverage (denominator/tested/backlog).
@@ -1842,6 +1861,7 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) {
VulnClass: normFilter(q.Get("vulnclass")),
// task_id(独立于会切到「按任务节点」分支的 task 参数):全局表按任务筛选。
TaskID: normFilter(q.Get("task_id")),
Query: q.Get("q"),
Sort: q.Get("sort"),
}
fs, total, err := s.m.pg.ListFindingsPage(filter, page, limit)
@@ -2011,7 +2031,7 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) {
// format = md-single(整合一份 .md)| md-zip(一漏洞一 .md,打包 zip)
// | csv | json
// ids = 逗号分隔的 finding id(scope=selected 时必填)
// 筛选参数 severity/status/vulnclass/task_id/sort 与列表接口一致(scope=filtered 用)。
// 筛选参数 severity/status/vulnclass/task_id/q/sort 与列表接口一致(scope=filtered 用)。
func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
q := r.URL.Query()
scope := q.Get("scope")
@@ -2045,6 +2065,7 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
Status: normFilter(q.Get("status")),
VulnClass: normFilter(q.Get("vulnclass")),
TaskID: normFilter(q.Get("task_id")),
Query: q.Get("q"),
Sort: q.Get("sort"),
}
default:
+655
View File
@@ -0,0 +1,655 @@
package server
import (
"archive/tar"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
"os"
"path/filepath"
"strconv"
"strings"
pgdb "github.com/Autumn-27/artex/db"
"github.com/klauspost/compress/zstd"
)
const (
archiveDirMode = 0o700
archiveFileMode = 0o600
maxArchiveFiles = 1_000_000
maxArchiveBytes = int64(1 << 47) // 128 TiB safety ceiling for corrupt headers.
)
type archiveFileMove struct {
Source string `json:"source"`
Relative string `json:"relative"`
}
type archiveStageJournal struct {
ArchiveID int64 `json:"archive_id"`
TaskID string `json:"task_id"`
Moves []archiveFileMove `json:"moves"`
}
type archiveRestoreJournal struct {
ArchiveID int64 `json:"archive_id"`
TaskID int64 `json:"task_id"`
Moves []restoredArchivePath `json:"moves"`
}
type taskArchiveFileStage struct {
root string
payload string
journal archiveStageJournal
done bool
}
type restoredArchivePath struct {
Source string
Destination string
}
type taskArchiveRestoreFiles struct {
extracted string
moves []restoredArchivePath
done bool
}
func writeArchiveJournal(path string, value any) error {
raw, err := json.Marshal(value)
if err != nil {
return err
}
file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, archiveFileMode)
if err != nil {
return err
}
if _, err := file.Write(raw); err != nil {
_ = file.Close()
return err
}
if err := file.Sync(); err != nil {
_ = file.Close()
return err
}
return file.Close()
}
func taskArchiveRoot(dataDir string) string {
return filepath.Join(dataDir, "archives", "tasks")
}
func taskArchivePath(dataDir string, archiveID int64, taskID string) string {
return filepath.Join(taskArchiveRoot(dataDir), fmt.Sprintf("task-%s-%d.tar.zst", taskID, archiveID))
}
func stageTaskArchiveFiles(dataDir string, archiveID int64, taskID string, explorationID int64) (*taskArchiveFileStage, error) {
root := filepath.Join(taskArchiveRoot(dataDir), ".staging", strconv.FormatInt(archiveID, 10))
stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload")}
journalPath := filepath.Join(root, "journal.json")
if raw, err := os.ReadFile(journalPath); err == nil {
if err := json.Unmarshal(raw, &stage.journal); err != nil {
return nil, fmt.Errorf("read task archive staging journal: %w", err)
}
return stage, nil
} else if !os.IsNotExist(err) {
return nil, err
}
if err := os.MkdirAll(filepath.Join(stage.payload, "files"), archiveDirMode); err != nil {
return nil, err
}
stage.journal = archiveStageJournal{ArchiveID: archiveID, TaskID: taskID}
targets := []archiveFileMove{}
taskDir := filepath.Join(dataDir, "tasks", taskID)
if _, err := os.Lstat(taskDir); err == nil {
targets = append(targets, archiveFileMove{Source: taskDir, Relative: filepath.Join("files", "tasks", taskID)})
} else if !os.IsNotExist(err) {
return nil, err
}
transcriptDir := filepath.Join(dataDir, "transcripts")
entries, err := os.ReadDir(transcriptDir)
if err != nil && !os.IsNotExist(err) {
return nil, err
}
prefix := fmt.Sprintf("exp%d-", explorationID)
for _, entry := range entries {
name := entry.Name()
if !strings.HasPrefix(name, prefix) || (!entry.IsDir() && !strings.HasSuffix(name, ".jsonl")) {
continue
}
targets = append(targets, archiveFileMove{
Source: filepath.Join(transcriptDir, name), Relative: filepath.Join("files", "transcripts", name),
})
}
// Persist the complete plan before the first rename. Recovery can therefore
// roll back any prefix of the moves after an abrupt process termination.
stage.journal.Moves = targets
if err := writeArchiveJournal(journalPath, stage.journal); err != nil {
_ = stage.rollback()
return nil, err
}
for _, move := range targets {
destination := filepath.Join(stage.payload, move.Relative)
if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil {
_ = stage.rollback()
return nil, err
}
if err := os.Rename(move.Source, destination); err != nil {
_ = stage.rollback()
return nil, fmt.Errorf("stage task archive path %s: %w", move.Source, err)
}
}
return stage, nil
}
func (s *taskArchiveFileStage) rollback() error {
if s == nil || s.done {
return nil
}
var errs []error
for i := len(s.journal.Moves) - 1; i >= 0; i-- {
move := s.journal.Moves[i]
staged := filepath.Join(s.payload, move.Relative)
if _, err := os.Lstat(staged); os.IsNotExist(err) {
continue
} else if err != nil {
errs = append(errs, err)
continue
}
if _, err := os.Lstat(move.Source); err == nil {
errs = append(errs, fmt.Errorf("archive rollback destination exists: %s", move.Source))
continue
} else if !os.IsNotExist(err) {
errs = append(errs, err)
continue
}
if err := os.MkdirAll(filepath.Dir(move.Source), 0o755); err != nil {
errs = append(errs, err)
continue
}
if err := os.Rename(staged, move.Source); err != nil {
errs = append(errs, err)
}
}
if len(errs) == 0 {
if err := os.RemoveAll(s.root); err != nil {
errs = append(errs, err)
}
}
s.done = true
return errors.Join(errs...)
}
func (s *taskArchiveFileStage) commit() error {
if s == nil || s.done {
return nil
}
s.done = true
return os.RemoveAll(s.root)
}
func installTaskArchiveFiles(dataDir, extractedDir, taskID string, archiveID int64) (*taskArchiveRestoreFiles, error) {
stage := &taskArchiveRestoreFiles{extracted: extractedDir}
numericTaskID, err := strconv.ParseInt(taskID, 10, 64)
if err != nil || numericTaskID <= 0 {
return nil, fmt.Errorf("invalid restore task id %q", taskID)
}
sources := []restoredArchivePath{}
workspace := filepath.Join(extractedDir, "files", "tasks", taskID)
if _, err := os.Lstat(workspace); err == nil {
sources = append(sources, restoredArchivePath{Source: workspace, Destination: filepath.Join(dataDir, "tasks", taskID)})
} else if !os.IsNotExist(err) {
return nil, err
}
transcripts := filepath.Join(extractedDir, "files", "transcripts")
entries, err := os.ReadDir(transcripts)
if err != nil && !os.IsNotExist(err) {
return nil, err
}
for _, entry := range entries {
sources = append(sources, restoredArchivePath{
Source: filepath.Join(transcripts, entry.Name()), Destination: filepath.Join(dataDir, "transcripts", entry.Name()),
})
}
for _, move := range sources {
if _, err := os.Lstat(move.Destination); err == nil {
return nil, fmt.Errorf("restore destination already exists: %s", move.Destination)
} else if !os.IsNotExist(err) {
return nil, err
}
}
stage.moves = sources
journal := archiveRestoreJournal{ArchiveID: archiveID, TaskID: numericTaskID, Moves: sources}
if err := writeArchiveJournal(filepath.Join(extractedDir, "restore-journal.json"), journal); err != nil {
return nil, err
}
for _, move := range sources {
if err := os.MkdirAll(filepath.Dir(move.Destination), 0o755); err != nil {
_ = stage.rollback()
return nil, err
}
if err := os.Rename(move.Source, move.Destination); err != nil {
_ = stage.rollback()
return nil, err
}
}
return stage, nil
}
func (s *taskArchiveRestoreFiles) rollback() error {
if s == nil || s.done {
return nil
}
var errs []error
for i := len(s.moves) - 1; i >= 0; i-- {
move := s.moves[i]
if _, err := os.Lstat(move.Destination); os.IsNotExist(err) {
continue
} else if err != nil {
errs = append(errs, err)
continue
}
if _, err := os.Lstat(move.Source); err == nil {
errs = append(errs, fmt.Errorf("restore rollback source exists: %s", move.Source))
continue
} else if !os.IsNotExist(err) {
errs = append(errs, err)
continue
}
if err := os.MkdirAll(filepath.Dir(move.Source), archiveDirMode); err != nil {
errs = append(errs, err)
continue
}
if err := os.Rename(move.Destination, move.Source); err != nil {
errs = append(errs, err)
}
}
s.done = true
return errors.Join(errs...)
}
func (s *taskArchiveRestoreFiles) commit() error {
if s == nil || s.done {
return nil
}
s.done = true
return os.RemoveAll(s.extracted)
}
// recoverTaskArchiveRestoreStages resolves file installs left by an interrupted
// restore. If PostgreSQL committed, installed files are authoritative and only
// the extraction directory is stale. Otherwise all completed renames are moved
// back so the persistent restore job can retry from a clean destination.
func recoverTaskArchiveRestoreStages(dataDir string, pg *pgdb.DB) error {
parent := filepath.Join(taskArchiveRoot(dataDir), ".restore")
entries, err := os.ReadDir(parent)
if os.IsNotExist(err) {
return nil
}
if err != nil {
return err
}
var errs []error
for _, entry := range entries {
if !entry.IsDir() {
continue
}
root := filepath.Join(parent, entry.Name())
raw, err := os.ReadFile(filepath.Join(root, "restore-journal.json"))
if os.IsNotExist(err) {
// Extraction was interrupted before any destination rename.
if err := os.RemoveAll(root); err != nil {
errs = append(errs, err)
}
continue
}
if err != nil {
errs = append(errs, err)
continue
}
var journal archiveRestoreJournal
if err := json.Unmarshal(raw, &journal); err != nil {
errs = append(errs, err)
continue
}
restored, err := pg.IsTaskArchiveRestored(journal.ArchiveID)
if errors.Is(err, pgdb.ErrTaskArchiveNotFound) && journal.TaskID > 0 {
task, taskErr := pg.GetTask(journal.TaskID)
if taskErr != nil {
errs = append(errs, taskErr)
continue
}
restored = task != nil
err = nil
}
if err != nil {
errs = append(errs, err)
continue
}
if restored {
if err := os.RemoveAll(root); err != nil {
errs = append(errs, err)
}
continue
}
stage := &taskArchiveRestoreFiles{extracted: root, moves: journal.Moves}
if err := stage.rollback(); err != nil {
errs = append(errs, err)
continue
}
if err := os.RemoveAll(root); err != nil {
errs = append(errs, err)
}
}
return errors.Join(errs...)
}
// recoverTaskArchiveStages resolves file renames left by an interrupted archive.
// A committed cold task already has a verified package, so stale staging can be
// discarded; otherwise files are moved back before the persistent job retries.
func recoverTaskArchiveStages(dataDir string, pg *pgdb.DB) error {
parent := filepath.Join(taskArchiveRoot(dataDir), ".staging")
entries, err := os.ReadDir(parent)
if os.IsNotExist(err) {
return nil
}
if err != nil {
return err
}
var errs []error
for _, entry := range entries {
if !entry.IsDir() {
continue
}
root := filepath.Join(parent, entry.Name())
raw, err := os.ReadFile(filepath.Join(root, "journal.json"))
if err != nil {
errs = append(errs, err)
continue
}
var journal archiveStageJournal
if err := json.Unmarshal(raw, &journal); err != nil {
errs = append(errs, err)
continue
}
archive, getErr := pg.GetTaskArchive(journal.ArchiveID)
if getErr != nil {
errs = append(errs, getErr)
continue
}
if archive != nil && archive.State == pgdb.ArchiveReady {
if err := os.RemoveAll(root); err != nil {
errs = append(errs, err)
}
continue
}
stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload"), journal: journal}
if err := stage.rollback(); err != nil {
errs = append(errs, err)
}
if archive != nil && archive.ArchivePath != "" {
_ = os.Remove(archive.ArchivePath)
}
}
return errors.Join(errs...)
}
func recoverTaskArchiveDeletePackages(dataDir string, pg *pgdb.DB) error {
root := taskArchiveRoot(dataDir)
entries, err := os.ReadDir(root)
if os.IsNotExist(err) {
return nil
}
if err != nil {
return err
}
var errs []error
for _, entry := range entries {
if entry.IsDir() {
continue
}
name := entry.Name()
marker := strings.LastIndex(name, ".deleting-")
if marker < 0 {
continue
}
archiveID, err := strconv.ParseInt(name[marker+len(".deleting-"):], 10, 64)
if err != nil || archiveID <= 0 {
continue
}
staged := filepath.Join(root, name)
original := filepath.Join(root, name[:marker])
archive, err := pg.GetTaskArchive(archiveID)
if err != nil {
errs = append(errs, err)
continue
}
if archive == nil {
if err := os.Remove(staged); err != nil && !os.IsNotExist(err) {
errs = append(errs, err)
}
continue
}
if archive.State == pgdb.DeleteQueued || archive.State == pgdb.Deleting || archive.State == pgdb.DeleteFailed {
// The idempotent delete worker consumes the staged path directly.
continue
}
if _, err := os.Lstat(original); err == nil {
errs = append(errs, fmt.Errorf("archive delete recovery destination exists: %s", original))
continue
} else if !os.IsNotExist(err) {
errs = append(errs, err)
continue
}
if err := os.Rename(staged, original); err != nil {
errs = append(errs, err)
}
}
return errors.Join(errs...)
}
func stageTaskArchivePackageDelete(archivePath string, archiveID int64) (string, bool, error) {
staged := archivePath + fmt.Sprintf(".deleting-%d", archiveID)
if _, err := os.Lstat(staged); err == nil {
if _, originalErr := os.Lstat(archivePath); originalErr == nil {
return staged, false, errors.New("归档包原文件和删除暂存文件同时存在")
} else if !os.IsNotExist(originalErr) {
return staged, false, originalErr
}
return staged, true, nil
} else if !os.IsNotExist(err) {
return staged, false, err
}
if err := os.Rename(archivePath, staged); err == nil {
return staged, true, nil
} else if !os.IsNotExist(err) {
return staged, false, err
}
return staged, false, nil
}
func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchiveSnapshot) (originalSize, compressedSize int64, checksum string, err error) {
if snapshot == nil {
return 0, 0, "", errors.New("nil task archive snapshot")
}
if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil {
return 0, 0, "", err
}
manifest, err := json.Marshal(snapshot)
if err != nil {
return 0, 0, "", err
}
if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil {
return 0, 0, "", err
}
temporary := path + ".partial"
_ = os.Remove(temporary)
file, err := os.OpenFile(temporary, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode)
if err != nil {
return 0, 0, "", err
}
committed := false
defer func() {
_ = file.Close()
if !committed {
_ = os.Remove(temporary)
}
}()
hasher := sha256.New()
multi := io.MultiWriter(file, hasher)
encoder, err := zstd.NewWriter(multi,
zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(10)),
zstd.WithEncoderCRC(true),
)
if err != nil {
return 0, 0, "", err
}
tarWriter := tar.NewWriter(encoder)
walkErr := filepath.WalkDir(payloadDir, func(current string, entry fs.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
relative, err := filepath.Rel(payloadDir, current)
if err != nil || relative == "." {
return err
}
info, err := entry.Info()
if err != nil {
return err
}
if info.Mode()&os.ModeSymlink != 0 {
return fmt.Errorf("task archive refuses symlink: %s", current)
}
header, err := tar.FileInfoHeader(info, "")
if err != nil {
return err
}
header.Name = filepath.ToSlash(relative)
if err := tarWriter.WriteHeader(header); err != nil {
return err
}
if entry.IsDir() {
return nil
}
originalSize += info.Size()
input, err := os.Open(current)
if err != nil {
return err
}
_, copyErr := io.Copy(tarWriter, input)
closeErr := input.Close()
return errors.Join(copyErr, closeErr)
})
if walkErr != nil {
_ = tarWriter.Close()
_ = encoder.Close()
return 0, 0, "", walkErr
}
if err := tarWriter.Close(); err != nil {
_ = encoder.Close()
return 0, 0, "", err
}
if err := encoder.Close(); err != nil {
return 0, 0, "", err
}
if err := file.Sync(); err != nil {
return 0, 0, "", err
}
if err := file.Close(); err != nil {
return 0, 0, "", err
}
stat, err := os.Stat(temporary)
if err != nil {
return 0, 0, "", err
}
compressedSize = stat.Size()
checksum = hex.EncodeToString(hasher.Sum(nil))
if err := os.Rename(temporary, path); err != nil {
return 0, 0, "", err
}
if err := os.Chmod(path, archiveFileMode); err != nil {
return 0, 0, "", err
}
committed = true
return originalSize, compressedSize, checksum, nil
}
func extractTaskArchivePackage(path, expectedSHA, destination string) error {
file, err := os.Open(path)
if err != nil {
return err
}
defer file.Close()
if expectedSHA != "" {
hasher := sha256.New()
if _, err := io.Copy(hasher, file); err != nil {
return err
}
if hex.EncodeToString(hasher.Sum(nil)) != strings.ToLower(expectedSHA) {
return errors.New("task archive checksum mismatch")
}
if _, err := file.Seek(0, io.SeekStart); err != nil {
return err
}
}
decoder, err := zstd.NewReader(file)
if err != nil {
return err
}
defer decoder.Close()
tarReader := tar.NewReader(decoder)
if err := os.MkdirAll(destination, archiveDirMode); err != nil {
return err
}
var entries int
var total int64
for {
header, err := tarReader.Next()
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return err
}
entries++
if entries > maxArchiveFiles || header.Size < 0 || total+header.Size > maxArchiveBytes {
return errors.New("task archive exceeds extraction safety limits")
}
total += header.Size
clean := filepath.Clean(filepath.FromSlash(header.Name))
if clean == "." || filepath.IsAbs(clean) || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
return fmt.Errorf("unsafe task archive path %q", header.Name)
}
target := filepath.Join(destination, clean)
if relative, err := filepath.Rel(destination, target); err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return fmt.Errorf("unsafe task archive target %q", header.Name)
}
switch header.Typeflag {
case tar.TypeDir:
if err := os.MkdirAll(target, archiveDirMode); err != nil {
return err
}
case tar.TypeReg, tar.TypeRegA:
if err := os.MkdirAll(filepath.Dir(target), archiveDirMode); err != nil {
return err
}
output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, fs.FileMode(header.Mode)&0o777)
if err != nil {
return err
}
_, copyErr := io.CopyN(output, tarReader, header.Size)
closeErr := output.Close()
if err := errors.Join(copyErr, closeErr); err != nil {
return err
}
default:
return fmt.Errorf("unsupported task archive entry type %d", header.Typeflag)
}
}
return nil
}
+209
View File
@@ -0,0 +1,209 @@
package server
import (
"archive/tar"
"bytes"
"encoding/json"
"os"
"path/filepath"
"testing"
pgdb "github.com/Autumn-27/artex/db"
"github.com/klauspost/compress/zstd"
)
func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
dataDir := t.TempDir()
taskID := "42"
explorationID := int64(73)
taskFile := filepath.Join(dataDir, "tasks", taskID, "uploads", "evidence.txt")
transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker-1.jsonl")
unrelatedTranscript := filepath.Join(dataDir, "transcripts", "exp74-worker-1.jsonl")
for path, body := range map[string]string{
taskFile: "task evidence",
transcriptFile: "transcript",
unrelatedTranscript: "leave me hot",
} {
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
t.Fatal(err)
}
}
stage, err := stageTaskArchiveFiles(dataDir, 1, taskID, explorationID)
if err != nil {
t.Fatal(err)
}
if _, err := os.Stat(taskFile); !os.IsNotExist(err) {
t.Fatalf("task file still hot after staging: %v", err)
}
if _, err := os.Stat(transcriptFile); !os.IsNotExist(err) {
t.Fatalf("task transcript still hot after staging: %v", err)
}
if _, err := os.Stat(unrelatedTranscript); err != nil {
t.Fatalf("unrelated transcript was staged: %v", err)
}
archivePath := taskArchivePath(dataDir, 1, taskID)
snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID}
original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot)
if err != nil {
t.Fatal(err)
}
if original == 0 || compressed == 0 || checksum == "" {
t.Fatalf("invalid package metrics original=%d compressed=%d checksum=%q", original, compressed, checksum)
}
if err := stage.commit(); err != nil {
t.Fatal(err)
}
extracted := filepath.Join(dataDir, "restore")
if err := extractTaskArchivePackage(archivePath, checksum, extracted); err != nil {
t.Fatal(err)
}
installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1)
if err != nil {
t.Fatal(err)
}
if err := installed.commit(); err != nil {
t.Fatal(err)
}
for path, want := range map[string]string{taskFile: "task evidence", transcriptFile: "transcript"} {
got, err := os.ReadFile(path)
if err != nil || string(got) != want {
t.Fatalf("restored %s = %q, %v; want %q", path, got, err, want)
}
}
if err := extractTaskArchivePackage(archivePath, "deadbeef", filepath.Join(dataDir, "bad-checksum")); err == nil {
t.Fatal("checksum mismatch was accepted")
}
}
func TestTaskArchiveStageJournalRollsBackInterruptedMoves(t *testing.T) {
dataDir := t.TempDir()
taskFile := filepath.Join(dataDir, "tasks", "42", "evidence.txt")
transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker.jsonl")
for _, path := range []string{taskFile, transcriptFile} {
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte(path), 0o600); err != nil {
t.Fatal(err)
}
}
stage, err := stageTaskArchiveFiles(dataDir, 9, "42", 73)
if err != nil {
t.Fatal(err)
}
raw, err := os.ReadFile(filepath.Join(stage.root, "journal.json"))
if err != nil {
t.Fatal(err)
}
var journal archiveStageJournal
if err := json.Unmarshal(raw, &journal); err != nil {
t.Fatal(err)
}
if len(journal.Moves) != 2 {
t.Fatalf("journal moves=%d, want 2", len(journal.Moves))
}
recovered := &taskArchiveFileStage{root: stage.root, payload: stage.payload, journal: journal}
if err := recovered.rollback(); err != nil {
t.Fatal(err)
}
for _, path := range []string{taskFile, transcriptFile} {
if _, err := os.Stat(path); err != nil {
t.Fatalf("interrupted archive did not restore %s: %v", path, err)
}
}
}
func TestTaskArchiveRestoreJournalRollsBackInterruptedInstall(t *testing.T) {
dataDir := t.TempDir()
extracted := filepath.Join(dataDir, "archives", "tasks", ".restore", "11-test")
source := filepath.Join(extracted, "files", "tasks", "42", "evidence.txt")
if err := os.MkdirAll(filepath.Dir(source), 0o700); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(source, []byte("evidence"), 0o600); err != nil {
t.Fatal(err)
}
_, err := installTaskArchiveFiles(dataDir, extracted, "42", 11)
if err != nil {
t.Fatal(err)
}
destination := filepath.Join(dataDir, "tasks", "42", "evidence.txt")
if _, err := os.Stat(destination); err != nil {
t.Fatal(err)
}
raw, err := os.ReadFile(filepath.Join(extracted, "restore-journal.json"))
if err != nil {
t.Fatal(err)
}
var journal archiveRestoreJournal
if err := json.Unmarshal(raw, &journal); err != nil {
t.Fatal(err)
}
if journal.ArchiveID != 11 || len(journal.Moves) != 1 {
t.Fatalf("unexpected restore journal: %+v", journal)
}
recovered := &taskArchiveRestoreFiles{extracted: extracted, moves: journal.Moves}
if err := recovered.rollback(); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(destination); !os.IsNotExist(err) {
t.Fatalf("hot destination remains after restore rollback: %v", err)
}
if got, err := os.ReadFile(source); err != nil || string(got) != "evidence" {
t.Fatalf("restore source=%q err=%v", got, err)
}
}
func TestTaskArchiveDeletePackageCanResumeFromStagedPath(t *testing.T) {
dataDir := t.TempDir()
archivePath := taskArchivePath(dataDir, 21, "42")
if err := os.MkdirAll(filepath.Dir(archivePath), archiveDirMode); err != nil {
t.Fatal(err)
}
staged := archivePath + ".deleting-21"
if err := os.WriteFile(staged, []byte("archive"), archiveFileMode); err != nil {
t.Fatal(err)
}
got, moved, err := stageTaskArchivePackageDelete(archivePath, 21)
if err != nil || !moved || got != staged {
t.Fatalf("resume staged package path=%q moved=%v err=%v", got, moved, err)
}
}
func TestTaskArchivePackageRejectsTraversal(t *testing.T) {
var compressed bytes.Buffer
encoder, err := zstd.NewWriter(&compressed)
if err != nil {
t.Fatal(err)
}
tw := tar.NewWriter(encoder)
content := []byte("escape")
if err := tw.WriteHeader(&tar.Header{Name: "../escape.txt", Mode: 0o600, Size: int64(len(content))}); err != nil {
t.Fatal(err)
}
if _, err := tw.Write(content); err != nil {
t.Fatal(err)
}
if err := tw.Close(); err != nil {
t.Fatal(err)
}
if err := encoder.Close(); err != nil {
t.Fatal(err)
}
path := filepath.Join(t.TempDir(), "malicious.tar.zst")
if err := os.WriteFile(path, compressed.Bytes(), 0o600); err != nil {
t.Fatal(err)
}
root := t.TempDir()
if err := extractTaskArchivePackage(path, "", root); err == nil {
t.Fatal("path traversal archive was accepted")
}
if _, err := os.Stat(filepath.Join(filepath.Dir(root), "escape.txt")); !os.IsNotExist(err) {
t.Fatalf("traversal wrote outside destination: %v", err)
}
}
+656
View File
@@ -0,0 +1,656 @@
package server
import (
"context"
"encoding/json"
"errors"
"fmt"
"log"
"net/http"
"os"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/Autumn-27/artex/agent"
pgdb "github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/traffic"
)
const taskArchivePollInterval = 2 * time.Second
type archiveBatchRequest struct {
TaskIDs []string `json:"task_ids"`
ArchiveIDs []int64 `json:"archive_ids"`
}
type archiveBatchItem struct {
ID string `json:"id"`
OK bool `json:"ok"`
Queued bool `json:"queued"`
Error string `json:"error,omitempty"`
Archive int64 `json:"archive_id,omitempty"`
}
func (s *Server) startTaskArchiveWorker() {
if s.m == nil || s.m.pg == nil {
return
}
if err := s.m.pg.RecoverTaskArchiveJobs(); err != nil {
log.Printf("[task-archive] recover jobs: %v", err)
}
recoveryFailed := false
if err := recoverTaskArchiveStages(s.m.dir, s.m.pg); err != nil {
log.Printf("[task-archive] recover file staging: %v", err)
recoveryFailed = true
}
if err := recoverTaskArchiveRestoreStages(s.m.dir, s.m.pg); err != nil {
log.Printf("[task-archive] recover restore staging: %v", err)
recoveryFailed = true
}
if err := recoverTaskArchiveDeletePackages(s.m.dir, s.m.pg); err != nil {
log.Printf("[task-archive] recover delete packages: %v", err)
recoveryFailed = true
}
if recoveryFailed {
log.Printf("[task-archive] worker disabled because startup recovery is incomplete")
return
}
s.archiveWG.Add(1)
go func() {
defer s.archiveWG.Done()
ticker := time.NewTicker(taskArchivePollInterval)
defer ticker.Stop()
for {
if err := s.runOneTaskArchiveJob(); err != nil {
log.Printf("[task-archive] worker: %v", err)
}
select {
case <-s.ctx.Done():
return
case <-s.archiveWake:
case <-ticker.C:
}
}
}()
s.notifyTaskArchiveWorker()
}
func (s *Server) notifyTaskArchiveWorker() {
if s == nil || s.archiveWake == nil {
return
}
select {
case s.archiveWake <- struct{}{}:
default:
}
}
func (s *Server) runOneTaskArchiveJob() error {
job, err := s.m.pg.ClaimTaskArchiveJob(s.ctx)
if err != nil || job == nil {
return err
}
var runErr error
switch job.State {
case pgdb.Archiving:
runErr = s.archiveTask(job)
case pgdb.Restoring:
runErr = s.restoreTaskArchive(job)
case pgdb.Deleting:
runErr = s.deleteTaskArchive(job)
default:
runErr = fmt.Errorf("unknown claimed archive state %q", job.State)
}
if runErr != nil {
if failErr := s.m.pg.FailTaskArchiveJob(job.ID, job.State, runErr); failErr != nil {
return errors.Join(runErr, failErr)
}
return fmt.Errorf("archive job %d (%s): %w", job.ID, job.State, runErr)
}
// Drain another queued item without waiting for the periodic poll.
s.notifyTaskArchiveWorker()
return nil
}
func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
taskID := strconv.FormatInt(job.TaskID, 10)
if !s.beginTaskDelete(taskID) {
return errors.New("任务正在进行其他归档或删除操作")
}
committed := false
defer func() {
if !committed {
s.abortTaskDelete(taskID)
}
}()
s.cancelTaskChat(taskID, agent.AbortTaskDeleted)
drainCtx, cancel := context.WithTimeout(s.ctx, taskDeleteDrainTimeout)
defer cancel()
if err := s.waitTaskQuiescent(drainCtx, taskID); err != nil {
return errors.New("任务仍有运行中的 Agent,请先暂停后重试归档")
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 10)
snapshot, err := s.m.pg.SnapshotTaskArchive(job.TaskID)
if err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 25)
fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, snapshot.ExplorationID)
if err != nil {
return err
}
defer func() {
if !committed {
runErr = errors.Join(runErr, fileStage.rollback())
}
}()
if s.m.traffic != nil && len(snapshot.Hosts) > 0 {
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_traffic", 38)
trafficCount, err := s.m.traffic.ExportHosts(snapshot.Hosts, filepath.Join(fileStage.payload, "traffic"))
if err != nil {
return err
}
snapshot.DataCounts["traffic"] = trafficCount
}
archivePath := taskArchivePath(s.m.dir, job.ID, taskID)
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50)
originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot)
if err != nil {
return err
}
removePackage := true
defer func() {
if removePackage {
_ = os.Remove(archivePath)
}
}()
var trafficStage *traffic.HostDeleteStage
if s.m.traffic != nil && len(snapshot.ExclusiveHosts) > 0 {
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_traffic", 68)
trafficStage, err = s.m.traffic.StageDeleteHostsExactForArchive(snapshot.ExclusiveHosts, job.ID, job.TaskID)
if err != nil {
return err
}
defer func() {
if !committed {
runErr = errors.Join(runErr, trafficStage.Rollback())
}
}()
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_database", 78)
if err := s.m.pg.CompleteTaskArchive(job.ID, snapshot, archivePath, checksum, originalSize, compressedSize); err != nil {
return err
}
committed = true
removePackage = false
if trafficStage != nil {
if err := trafficStage.Commit(); err != nil {
warning := "归档已完成,但独占流量热存储清理失败(归档包仍可恢复):" + err.Error()
log.Printf("[task-archive] task %s: %s", taskID, warning)
_ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning)
}
}
if err := fileStage.commit(); err != nil {
warning := "归档已完成,但暂存目录清理失败:" + err.Error()
log.Printf("[task-archive] task %s: %s", taskID, warning)
_ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning)
}
s.engine.StopTask(taskID)
s.taskAgentMu.Lock()
delete(s.taskAgents, taskID)
s.taskAgentMu.Unlock()
s.m.forgetTask(taskID, job.TaskID)
return nil
}
func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil {
return err
}
// A prior attempt may have committed PostgreSQL and files, then failed only
// while consuming the package. Finish that cleanup without replaying rows.
if restored, err := s.m.pg.IsTaskArchiveRestored(job.ID); err != nil {
return err
} else if restored {
if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) {
return err
}
if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil {
return err
}
s.loadRestoredTask(job.TaskID)
return nil
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "verify_package", 10)
restoreParent := filepath.Join(taskArchiveRoot(s.m.dir), ".restore")
if err := os.MkdirAll(restoreParent, archiveDirMode); err != nil {
return err
}
extracted, err := os.MkdirTemp(restoreParent, fmt.Sprintf("%d-", job.ID))
if err != nil {
return err
}
defer func() { _ = os.RemoveAll(extracted) }()
if err := extractTaskArchivePackage(job.ArchivePath, job.SHA256, extracted); err != nil {
return err
}
manifestRaw, err := os.ReadFile(filepath.Join(extracted, "manifest.json"))
if err != nil {
return err
}
var snapshot pgdb.TaskArchiveSnapshot
if err := json.Unmarshal(manifestRaw, &snapshot); err != nil {
return err
}
if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion {
return pgdb.ErrTaskArchiveFormatMismatch
}
if s.m.traffic != nil {
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_traffic", 35)
if _, err := s.m.traffic.ImportArchive(filepath.Join(extracted, "traffic")); err != nil {
return err
}
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_files", 55)
files, err := installTaskArchiveFiles(s.m.dir, extracted, strconv.FormatInt(job.TaskID, 10), job.ID)
if err != nil {
return err
}
databaseRestored := false
defer func() {
if !databaseRestored {
runErr = errors.Join(runErr, files.rollback())
}
}()
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70)
warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
if err != nil {
return err
}
for _, warning := range warnings {
log.Printf("[task-archive] restored task %d with warning: %s", job.TaskID, warning)
}
if len(warnings) > 0 {
if _, err := s.m.pg.Exploration(snapshot.ExplorationID).AppendActivity(pgdb.Activity{
Worker: "system", Kind: "system", Summary: "任务还原完成,但部分关联已降级", Detail: strings.Join(warnings, "\n"),
}); err != nil {
log.Printf("[task-archive] persist restore warnings for task %d: %v", job.TaskID, err)
}
}
databaseRestored = true
if err := files.commit(); err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "consume_package", 95)
if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) {
return err
}
if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil {
return err
}
s.loadRestoredTask(job.TaskID)
return nil
}
func (s *Server) loadRestoredTask(taskID int64) {
id := strconv.FormatInt(taskID, 10)
var restored *Task
for _, task := range s.m.LoadExisting() {
if task.ID == id {
restored = task
break
}
}
if restored == nil {
if task, ok := s.m.Task(id); ok {
restored = task
}
}
if restored == nil {
return
}
lifecycle := restored.lifecycleSnapshot()
if lifecycle.Paused {
s.engine.Pause(id, agent.AbortPausedOnReload)
}
if !isTerminalStatus(lifecycle.Status) {
s.engine.startDeadlineCoordinator(s.ctx, restored)
}
}
func (s *Server) deleteTaskArchive(job *pgdb.TaskArchive) error {
if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_package_delete", 35)
staged, moved, err := stageTaskArchivePackageDelete(job.ArchivePath, job.ID)
if err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "delete_metadata", 70)
if err := s.m.pg.DeleteTaskArchiveStub(job.ID); err != nil {
if moved {
_ = os.Rename(staged, job.ArchivePath)
}
return err
}
if moved {
if err := os.Remove(staged); err != nil && !os.IsNotExist(err) {
log.Printf("[task-archive] archived task %d metadata deleted; stale package cleanup failed: %v", job.TaskID, err)
}
}
return nil
}
func validateArchivePath(dataDir, candidate string) error {
if strings.TrimSpace(candidate) == "" {
return errors.New("归档包路径为空")
}
root, err := filepath.Abs(taskArchiveRoot(dataDir))
if err != nil {
return err
}
path, err := filepath.Abs(candidate)
if err != nil {
return err
}
relative, err := filepath.Rel(root, path)
if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return errors.New("归档包路径不在受管目录内")
}
return nil
}
func (s *Server) listTaskArchives(w http.ResponseWriter, r *http.Request) {
page, _ := strconv.Atoi(r.URL.Query().Get("page"))
size, _ := strconv.Atoi(r.URL.Query().Get("size"))
result, err := s.m.pg.ListTaskArchives(r.URL.Query().Get("q"), r.URL.Query().Get("state"), page, size)
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, result)
}
func (s *Server) getTaskArchive(w http.ResponseWriter, r *http.Request) {
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, 400, "归档 id 无效")
return
}
item, err := s.m.pg.GetTaskArchive(id)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if item == nil {
writeErr(w, 404, "归档不存在")
return
}
writeJSON(w, 200, item)
}
func (s *Server) queueTaskArchive(w http.ResponseWriter, r *http.Request) {
id, ok := canonicalTaskID(r.PathValue("id"))
if !ok {
writeErr(w, 400, "任务 id 无效")
return
}
numeric, _ := strconv.ParseInt(id, 10, 64)
item, err := s.m.pg.QueueTaskArchive(numeric)
if err != nil {
writeArchiveError(w, err)
return
}
s.notifyTaskArchiveWorker()
writeJSON(w, http.StatusAccepted, item)
}
func (s *Server) queueTaskArchiveRestore(w http.ResponseWriter, r *http.Request) {
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, 400, "归档 id 无效")
return
}
item, err := s.m.pg.QueueTaskArchiveRestore(id)
if err != nil {
writeArchiveError(w, err)
return
}
s.notifyTaskArchiveWorker()
writeJSON(w, http.StatusAccepted, item)
}
func (s *Server) queueTaskArchiveDelete(w http.ResponseWriter, r *http.Request) {
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, 400, "归档 id 无效")
return
}
item, err := s.m.pg.QueueTaskArchiveDelete(id)
if err != nil {
writeArchiveError(w, err)
return
}
s.notifyTaskArchiveWorker()
writeJSON(w, http.StatusAccepted, item)
}
func (s *Server) queueTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
var request archiveBatchRequest
if err := decode(r, &request); err != nil {
writeErr(w, 400, err.Error())
return
}
parsed := normalizeBatchTaskIDs(request.TaskIDs)
if len(parsed) == 0 {
writeErr(w, 400, "task_ids 不能为空")
return
}
if len(parsed) > 100 {
writeErr(w, 400, "一次最多处理 100 个任务")
return
}
ids := make([]string, 0, len(parsed))
items := make([]archiveBatchItem, 0, len(parsed))
for _, item := range parsed {
if !item.valid {
items = append(items, archiveBatchItem{ID: item.id, Error: "任务 id 无效"})
continue
}
ids = append(ids, item.id)
}
ids = s.orderLiveTaskArchiveBatch(ids)
for _, id := range ids {
numeric, _ := strconv.ParseInt(id, 10, 64)
archive, err := s.m.pg.QueueTaskArchive(numeric)
item := archiveBatchItem{ID: id, OK: err == nil, Queued: err == nil}
if err != nil {
item.Error = err.Error()
} else {
item.Archive = archive.ID
}
items = append(items, item)
}
s.notifyTaskArchiveWorker()
writeJSON(w, http.StatusAccepted, map[string]any{"items": items})
}
func (s *Server) queueTaskArchiveActionBatch(w http.ResponseWriter, r *http.Request, action string) {
var request archiveBatchRequest
if err := decode(r, &request); err != nil {
writeErr(w, 400, err.Error())
return
}
ids, err := normalizeArchiveIDs(request.ArchiveIDs)
if err != nil {
writeErr(w, 400, err.Error())
return
}
ids = s.orderColdTaskArchiveBatch(ids, action == "restore")
items := make([]archiveBatchItem, 0, len(ids))
for _, id := range ids {
var archive *pgdb.TaskArchive
if action == "restore" {
archive, err = s.m.pg.QueueTaskArchiveRestore(id)
} else {
archive, err = s.m.pg.QueueTaskArchiveDelete(id)
}
item := archiveBatchItem{ID: strconv.FormatInt(id, 10), Archive: id, OK: err == nil, Queued: err == nil}
if err != nil {
item.Error = err.Error()
} else if archive != nil {
item.Archive = archive.ID
}
items = append(items, item)
}
s.notifyTaskArchiveWorker()
writeJSON(w, http.StatusAccepted, map[string]any{"items": items})
}
func (s *Server) restoreTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
s.queueTaskArchiveActionBatch(w, r, "restore")
}
func (s *Server) deleteTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
s.queueTaskArchiveActionBatch(w, r, "delete")
}
func normalizeArchiveIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, errors.New("archive_ids 不能为空")
}
if len(ids) > 100 {
return nil, errors.New("一次最多处理 100 个归档")
}
seen := map[int64]bool{}
out := make([]int64, 0, len(ids))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("归档 id %d 无效", id)
}
if !seen[id] {
seen[id] = true
out = append(out, id)
}
}
return out, nil
}
func (s *Server) orderLiveTaskArchiveBatch(ids []string) []string {
sources := map[string][]string{}
for _, id := range ids {
if task, ok := s.m.Task(id); ok {
for _, sourceID := range task.lifecycleSnapshot().SourceTaskIDs {
sources[id] = append(sources[id], strconv.FormatInt(sourceID, 10))
}
}
}
return topoTaskIDs(ids, sources, false)
}
func (s *Server) orderColdTaskArchiveBatch(ids []int64, restore bool) []int64 {
taskToArchive := map[string]int64{}
sources := map[string][]string{}
var taskIDs []string
for _, id := range ids {
archive, err := s.m.pg.GetTaskArchive(id)
if err != nil || archive == nil {
continue
}
taskID := strconv.FormatInt(archive.TaskID, 10)
taskToArchive[taskID] = id
taskIDs = append(taskIDs, taskID)
for _, sourceID := range archive.SourceTaskIDs {
sources[taskID] = append(sources[taskID], strconv.FormatInt(sourceID, 10))
}
}
orderedTasks := topoTaskIDs(taskIDs, sources, restore)
ordered := make([]int64, 0, len(ids))
seen := map[int64]bool{}
for _, taskID := range orderedTasks {
id := taskToArchive[taskID]
ordered = append(ordered, id)
seen[id] = true
}
for _, id := range ids {
if !seen[id] {
ordered = append(ordered, id)
}
}
return ordered
}
// topoTaskIDs orders source->dependent for restore and dependent->source for
// archive/delete. Cycles are impossible at creation, but stable fallback keeps a
// damaged database operable.
func topoTaskIDs(ids []string, sources map[string][]string, restore bool) []string {
set := map[string]bool{}
for _, id := range ids {
set[id] = true
}
adjacency := map[string][]string{}
indegree := map[string]int{}
for id := range set {
indegree[id] = 0
}
for dependent, sourceIDs := range sources {
for _, source := range sourceIDs {
if !set[dependent] || !set[source] {
continue
}
from, to := dependent, source
if restore {
from, to = source, dependent
}
adjacency[from] = append(adjacency[from], to)
indegree[to]++
}
}
queue := []string{}
for _, id := range ids {
if indegree[id] == 0 {
queue = append(queue, id)
}
}
ordered := make([]string, 0, len(ids))
for len(queue) > 0 {
id := queue[0]
queue = queue[1:]
ordered = append(ordered, id)
for _, next := range adjacency[id] {
indegree[next]--
if indegree[next] == 0 {
queue = append(queue, next)
}
}
}
if len(ordered) != len(ids) {
seen := map[string]bool{}
for _, id := range ordered {
seen[id] = true
}
for _, id := range ids {
if !seen[id] {
ordered = append(ordered, id)
}
}
}
return ordered
}
func writeArchiveError(w http.ResponseWriter, err error) {
status := http.StatusConflict
if errors.Is(err, pgdb.ErrTaskArchiveNotFound) {
status = http.StatusNotFound
} else if !errors.Is(err, pgdb.ErrTaskArchiveIneligible) &&
!errors.Is(err, pgdb.ErrTaskArchiveQueued) &&
!errors.Is(err, pgdb.ErrTaskArchiveDependent) &&
!errors.Is(err, pgdb.ErrTaskArchiveState) &&
!errors.Is(err, pgdb.ErrTaskArchiveDeleteBlocked) {
status = http.StatusInternalServerError
}
writeErr(w, status, err.Error())
}
+88
View File
@@ -0,0 +1,88 @@
package server
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
)
func TestTaskArchiveAPIQueueListAndLimits(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("archive api", "verify 202", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
taskID, err := strconv.ParseInt(task.ID, 10, 64)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, taskID)
_ = m.pg.DeleteTask(taskID)
}()
// Start with an already-cancelled context so the background archive worker
// exits before these route-contract assertions enqueue work.
ctx, cancel := context.WithCancel(context.Background())
cancel()
dataDir := t.TempDir()
s := New(ctx, m, dataDir, dataDir, dataDir)
s.archiveWG.Wait()
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
request := func(method, path, body string) *httptest.ResponseRecorder {
req := httptest.NewRequest(method, path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
s.Handler().ServeHTTP(recorder, req)
return recorder
}
recorder := request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}")
if recorder.Code != http.StatusConflict {
t.Fatalf("running archive status=%d body=%s", recorder.Code, recorder.Body.String())
}
if err := m.pg.SetPaused(taskID, true); err != nil {
t.Fatal(err)
}
recorder = request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}")
if recorder.Code != http.StatusAccepted {
t.Fatalf("archive queue status=%d body=%s", recorder.Code, recorder.Body.String())
}
var queued struct {
ID int64 `json:"id"`
State string `json:"state"`
}
if err := json.Unmarshal(recorder.Body.Bytes(), &queued); err != nil {
t.Fatal(err)
}
if queued.ID <= 0 || queued.State != "archive_queued" {
t.Fatalf("unexpected queue response: %+v", queued)
}
recorder = request(http.MethodGet, "/api/task-archives?page=1&size=10&q="+task.ID, "")
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), fmt.Sprintf(`"task_id":%d`, taskID)) {
t.Fatalf("archive list status=%d body=%s", recorder.Code, recorder.Body.String())
}
ids := make([]string, 101)
for index := range ids {
ids[index] = strconv.Itoa(index + 1)
}
body, _ := json.Marshal(map[string]any{"task_ids": ids})
recorder = request(http.MethodPost, "/api/tasks/archive/batch", string(body))
if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), "100") {
t.Fatalf("oversized batch status=%d body=%s", recorder.Code, recorder.Body.String())
}
}
+17
View File
@@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"iter"
"strings"
"sync"
"testing"
"time"
@@ -148,6 +149,22 @@ func TestTaskDTOUsesEmptyArrays(t *testing.T) {
}
}
func TestApplyTaskArchiveBlocker(t *testing.T) {
t.Parallel()
dto := taskDTO(&Task{ID: "7"}, "paused")
applyTaskArchiveBlocker(&dto, map[int64]int64{7: 11})
if dto.ArchiveBlockedBy != "11" {
t.Fatalf("archive blocker=%q, want 11", dto.ArchiveBlockedBy)
}
raw, err := json.Marshal(dto)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), `"archive_blocked_by_task_id":"11"`) {
t.Fatalf("archive blocker missing from payload: %s", raw)
}
}
func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
t.Parallel()
task := &Task{ID: "7"}