fix: prevent OOM when archiving large tasks

server/task_archive_package.go

@@ -480,11 +480,20 @@ func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchive
 	if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil {
 		return 0, 0, "", err
 	}
-	manifest, err := json.Marshal(snapshot)
+	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
 	}
-	if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil {
+	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"

server/task_archive_package_test.go

@@ -47,7 +47,18 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
 	}
 
 	archivePath := taskArchivePath(dataDir, 1, taskID)
-	snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID}
+	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)
@@ -62,6 +73,9 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
 	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)

server/task_archives.go

@@ -131,13 +131,15 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
 	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)
+	task, err := s.m.pg.GetTask(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 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
 	}
@@ -146,6 +148,27 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
 			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"))
@@ -244,9 +267,27 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
 	if err := json.Unmarshal(manifestRaw, &snapshot); err != nil {
 		return err
 	}
-	if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion {
+	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 {
@@ -265,7 +306,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
 		}
 	}()
 	_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70)
-	warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
+	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
 	}