@posthog/agent 2.3.388 → 2.3.401

package/src/sagas/apply-snapshot-saga.ts

@@ -1,4 +1,4 @@
-import { mkdir, rm, writeFile } from "node:fs/promises";
+import { mkdir, readdir, rm, rmdir, writeFile } from "node:fs/promises";
 import { join } from "node:path";
 import { ApplyTreeSaga as GitApplyTreeSaga } from "@posthog/git/sagas/tree";
 import { Saga } from "@posthog/shared";
@@ -37,64 +37,78 @@ export class ApplySnapshotSaga extends Saga<
 
     const archiveUrl = snapshot.archiveUrl;
 
-    await this.step({
-      name: "create_tmp_dir",
-      execute: () => mkdir(tmpDir, { recursive: true }),
-      rollback: async () => {},
-    });
+    try {
+      await this.step({
+        name: "create_tmp_dir",
+        execute: () => mkdir(tmpDir, { recursive: true }),
+        rollback: async () => {},
+      });
 
-    const archivePath = join(tmpDir, `${snapshot.treeHash}.tar.gz`);
-    this.archivePath = archivePath;
-    await this.step({
-      name: "download_archive",
-      execute: async () => {
-        const arrayBuffer = await apiClient.downloadArtifact(
-          taskId,
-          runId,
-          archiveUrl,
-        );
-        if (!arrayBuffer) {
-          throw new Error("Failed to download archive");
-        }
-        const base64Content = Buffer.from(arrayBuffer).toString("utf-8");
-        const binaryContent = Buffer.from(base64Content, "base64");
-        await writeFile(archivePath, binaryContent);
-        this.log.info("Tree archive downloaded", {
-          treeHash: snapshot.treeHash,
-          snapshotBytes: binaryContent.byteLength,
-          snapshotWireBytes: arrayBuffer.byteLength,
-          totalBytes: binaryContent.byteLength,
-          totalWireBytes: arrayBuffer.byteLength,
-        });
-      },
-      rollback: async () => {
-        if (this.archivePath) {
-          await rm(this.archivePath, { force: true }).catch(() => {});
-        }
-      },
-    });
+      const archivePath = join(tmpDir, `${snapshot.treeHash}.tar.gz`);
+      this.archivePath = archivePath;
+      await this.step({
+        name: "download_archive",
+        execute: async () => {
+          const arrayBuffer = await apiClient.downloadArtifact(
+            taskId,
+            runId,
+            archiveUrl,
+          );
+          if (!arrayBuffer) {
+            throw new Error("Failed to download archive");
+          }
+          const base64Content = Buffer.from(arrayBuffer).toString("utf-8");
+          const binaryContent = Buffer.from(base64Content, "base64");
+          await writeFile(archivePath, binaryContent);
+          this.log.info("Tree archive downloaded", {
+            treeHash: snapshot.treeHash,
+            snapshotBytes: binaryContent.byteLength,
+            snapshotWireBytes: arrayBuffer.byteLength,
+            totalBytes: binaryContent.byteLength,
+            totalWireBytes: arrayBuffer.byteLength,
+          });
+        },
+        rollback: async () => {
+          if (this.archivePath) {
+            await rm(this.archivePath, { force: true }).catch(() => {});
+          }
+        },
+      });
 
-    const gitApplySaga = new GitApplyTreeSaga(this.log);
-    const applyResult = await gitApplySaga.run({
-      baseDir: repositoryPath,
-      treeHash: snapshot.treeHash,
-      baseCommit: snapshot.baseCommit,
-      changes: snapshot.changes,
-      archivePath: this.archivePath,
-    });
+      const gitApplySaga = new GitApplyTreeSaga(this.log);
+      const applyResult = await gitApplySaga.run({
+        baseDir: repositoryPath,
+        treeHash: snapshot.treeHash,
+        baseCommit: snapshot.baseCommit,
+        changes: snapshot.changes,
+        archivePath: this.archivePath,
+      });
 
-    if (!applyResult.success) {
-      throw new Error(`Failed to apply tree: ${applyResult.error}`);
-    }
+      if (!applyResult.success) {
+        throw new Error(`Failed to apply tree: ${applyResult.error}`);
+      }
 
-    await rm(this.archivePath, { force: true }).catch(() => {});
+      this.log.info("Tree snapshot applied", {
+        treeHash: snapshot.treeHash,
+        totalChanges: snapshot.changes.length,
+        deletedFiles: snapshot.changes.filter((c) => c.status === "D").length,
+      });
 
-    this.log.info("Tree snapshot applied", {
-      treeHash: snapshot.treeHash,
-      totalChanges: snapshot.changes.length,
-      deletedFiles: snapshot.changes.filter((c) => c.status === "D").length,
-    });
+      return { treeHash: snapshot.treeHash };
+    } finally {
+      if (this.archivePath) {
+        await rm(this.archivePath, { force: true }).catch(() => {});
+      }
+      await this.removeTmpDirIfEmpty(tmpDir);
+      this.archivePath = null;
+    }
+  }
 
-    return { treeHash: snapshot.treeHash };
+  private async removeTmpDirIfEmpty(tmpDir: string): Promise<void> {
+    const entries = await readdir(tmpDir).catch(() => null);
+    if (!entries || entries.length > 0) {
+      return;
+    }
+    await rmdir(tmpDir).catch(() => {});
   }
 }

package/src/sagas/capture-tree-saga.test.ts

@@ -366,6 +366,24 @@ describe("CaptureTreeSaga", () => {
       const indexFiles = files.filter((f: string) => f.startsWith("index-"));
       expect(indexFiles).toHaveLength(0);
     });
+
+    it("cleans up uploaded tree archive and tmp dir on success", async () => {
+      const mockApiClient = createMockApiClient();
+
+      await repo.writeFile("new.ts", "content");
+
+      const saga = new CaptureTreeSaga(mockLogger);
+      const result = await saga.run({
+        repositoryPath: repo.path,
+        taskId: "task-1",
+        runId: "run-1",
+        lastTreeHash: null,
+        apiClient: mockApiClient,
+      });
+
+      expect(result.success).toBe(true);
+      expect(repo.exists(".posthog/tmp")).toBe(false);
+    });
   });
 
   describe("git state isolation", () => {

package/src/sagas/capture-tree-saga.ts

@@ -1,5 +1,5 @@
 import { existsSync } from "node:fs";
-import { readFile, rm } from "node:fs/promises";
+import { readdir, readFile, rm, rmdir } from "node:fs/promises";
 import { join } from "node:path";
 import { CaptureTreeSaga as GitCaptureTreeSaga } from "@posthog/git/sagas/tree";
 import { Saga } from "@posthog/shared";
@@ -45,60 +45,67 @@ export class CaptureTreeSaga extends Saga<CaptureTreeInput, CaptureTreeOutput> {
       ? join(tmpDir, `tree-${Date.now()}.tar.gz`)
       : undefined;
 
-    const gitCaptureSaga = new GitCaptureTreeSaga(this.log);
-    const captureResult = await gitCaptureSaga.run({
-      baseDir: repositoryPath,
-      lastTreeHash,
-      archivePath,
-    });
+    try {
+      const gitCaptureSaga = new GitCaptureTreeSaga(this.log);
+      const captureResult = await gitCaptureSaga.run({
+        baseDir: repositoryPath,
+        lastTreeHash,
+        archivePath,
+      });
 
-    if (!captureResult.success) {
-      throw new Error(`Failed to capture tree: ${captureResult.error}`);
-    }
+      if (!captureResult.success) {
+        throw new Error(`Failed to capture tree: ${captureResult.error}`);
+      }
 
-    const {
-      snapshot: gitSnapshot,
-      archivePath: createdArchivePath,
-      changed,
-    } = captureResult.data;
-
-    if (!changed || !gitSnapshot) {
-      this.log.debug("No changes since last capture", { lastTreeHash });
-      return { snapshot: null, newTreeHash: lastTreeHash };
-    }
+      const {
+        snapshot: gitSnapshot,
+        archivePath: createdArchivePath,
+        changed,
+      } = captureResult.data;
 
-    let archiveUrl: string | undefined;
-    if (apiClient && createdArchivePath) {
-      try {
-        archiveUrl = await this.uploadArchive(
-          createdArchivePath,
-          gitSnapshot.treeHash,
-          apiClient,
-          taskId,
-          runId,
-        );
-      } finally {
-        await rm(createdArchivePath, { force: true }).catch(() => {});
+      if (!changed || !gitSnapshot) {
+        this.log.debug("No changes since last capture", { lastTreeHash });
+        return { snapshot: null, newTreeHash: lastTreeHash };
       }
-    }
 
-    const snapshot: TreeSnapshot = {
-      treeHash: gitSnapshot.treeHash,
-      baseCommit: gitSnapshot.baseCommit,
-      changes: gitSnapshot.changes,
-      timestamp: gitSnapshot.timestamp,
-      interrupted,
-      archiveUrl,
-    };
-
-    this.log.info("Tree captured", {
-      treeHash: snapshot.treeHash,
-      changes: snapshot.changes.length,
-      interrupted,
-      archiveUrl,
-    });
+      let archiveUrl: string | undefined;
+      if (apiClient && createdArchivePath) {
+        try {
+          archiveUrl = await this.uploadArchive(
+            createdArchivePath,
+            gitSnapshot.treeHash,
+            apiClient,
+            taskId,
+            runId,
+          );
+        } finally {
+          await rm(createdArchivePath, { force: true }).catch(() => {});
+        }
+      }
 
-    return { snapshot, newTreeHash: snapshot.treeHash };
+      const snapshot: TreeSnapshot = {
+        treeHash: gitSnapshot.treeHash,
+        baseCommit: gitSnapshot.baseCommit,
+        changes: gitSnapshot.changes,
+        timestamp: gitSnapshot.timestamp,
+        interrupted,
+        archiveUrl,
+      };
+
+      this.log.info("Tree captured", {
+        treeHash: snapshot.treeHash,
+        changes: snapshot.changes.length,
+        interrupted,
+        archiveUrl,
+      });
+
+      return { snapshot, newTreeHash: snapshot.treeHash };
+    } finally {
+      if (archivePath) {
+        await rm(archivePath, { force: true }).catch(() => {});
+      }
+      await this.removeTmpDirIfEmpty(tmpDir);
+    }
   }
 
   private async uploadArchive(
@@ -147,4 +154,12 @@ export class CaptureTreeSaga extends Saga<CaptureTreeInput, CaptureTreeOutput> {
 
     return archiveUrl;
   }
+
+  private async removeTmpDirIfEmpty(tmpDir: string): Promise<void> {
+    const entries = await readdir(tmpDir).catch(() => null);
+    if (!entries || entries.length > 0) {
+      return;
+    }
+    await rmdir(tmpDir).catch(() => {});
+  }
 }