Merge pull request #51 from getpaseo/implement/issue-43

fix: archive worktree when last agent is archived

packages/server/src/server/daemon-e2e/git-operations.e2e.test.ts

@@ -105,6 +105,21 @@ async function waitForPathExists(
   );
 }
 
+async function waitForPathRemoved(
+  options: { targetPath: string; timeoutMs: number; label: string }
+): Promise<void> {
+  const start = Date.now();
+  while (Date.now() - start < options.timeoutMs) {
+    if (!existsSync(options.targetPath)) {
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 50));
+  }
+  throw new Error(
+    `Timed out after ${options.timeoutMs}ms waiting for removal of ${options.label}: ${options.targetPath}`
+  );
+}
+
 async function withShell<T>(shell: string, run: () => Promise<T>): Promise<T> {
   const originalShell = process.env.SHELL;
   process.env.SHELL = shell;
@@ -830,6 +845,122 @@ describe("daemon E2E", () => {
       },
       60000
     );
+
+    test(
+      "archives the worktree when the last agent in it is archived",
+      async () => {
+        const repoRoot = tmpCwd();
+
+        const { execSync } = await import("child_process");
+        execSync("git init -b main", { cwd: repoRoot, stdio: "pipe" });
+        execSync("git config user.email 'test@test.com'", {
+          cwd: repoRoot,
+          stdio: "pipe",
+        });
+        execSync("git config user.name 'Test'", { cwd: repoRoot, stdio: "pipe" });
+
+        writeFileSync(path.join(repoRoot, "file.txt"), "hello\n");
+        execSync("git add .", { cwd: repoRoot, stdio: "pipe" });
+        execSync("git -c commit.gpgsign=false commit -m 'initial'", {
+          cwd: repoRoot,
+          stdio: "pipe",
+        });
+        execSync("git branch -M main", { cwd: repoRoot, stdio: "pipe" });
+
+        const agent = await ctx.client.createAgent({
+          provider: "codex",
+          model: CODEX_TEST_MODEL,
+          thinkingOptionId: CODEX_TEST_THINKING_OPTION_ID,
+          cwd: repoRoot,
+          title: "Archive Last Agent Cleanup Test",
+          git: {
+            createWorktree: true,
+            createNewBranch: true,
+            baseBranch: "main",
+            newBranchName: "archive-last-agent",
+            worktreeSlug: "archive-last-agent",
+          },
+        });
+
+        expect(existsSync(agent.cwd)).toBe(true);
+
+        const result = await ctx.client.archiveAgent(agent.id);
+        expect(result.archivedAt).toBeTruthy();
+
+        await waitForPathRemoved({
+          targetPath: agent.cwd,
+          timeoutMs: 10000,
+          label: "archived worktree",
+        });
+
+        rmSync(repoRoot, { recursive: true, force: true });
+      },
+      60000
+    );
+
+    test(
+      "does not archive the worktree until all agents in it are archived",
+      async () => {
+        const repoRoot = tmpCwd();
+
+        const { execSync } = await import("child_process");
+        execSync("git init -b main", { cwd: repoRoot, stdio: "pipe" });
+        execSync("git config user.email 'test@test.com'", {
+          cwd: repoRoot,
+          stdio: "pipe",
+        });
+        execSync("git config user.name 'Test'", { cwd: repoRoot, stdio: "pipe" });
+
+        writeFileSync(path.join(repoRoot, "file.txt"), "hello\n");
+        execSync("git add .", { cwd: repoRoot, stdio: "pipe" });
+        execSync("git -c commit.gpgsign=false commit -m 'initial'", {
+          cwd: repoRoot,
+          stdio: "pipe",
+        });
+        execSync("git branch -M main", { cwd: repoRoot, stdio: "pipe" });
+
+        const firstAgent = await ctx.client.createAgent({
+          provider: "codex",
+          model: CODEX_TEST_MODEL,
+          thinkingOptionId: CODEX_TEST_THINKING_OPTION_ID,
+          cwd: repoRoot,
+          title: "Archive Multi-Agent Test 1",
+          git: {
+            createWorktree: true,
+            createNewBranch: true,
+            baseBranch: "main",
+            newBranchName: "archive-multi-agent",
+            worktreeSlug: "archive-multi-agent",
+          },
+        });
+
+        const secondAgent = await ctx.client.createAgent({
+          provider: "codex",
+          model: CODEX_TEST_MODEL,
+          thinkingOptionId: CODEX_TEST_THINKING_OPTION_ID,
+          cwd: firstAgent.cwd,
+          title: "Archive Multi-Agent Test 2",
+        });
+
+        expect(existsSync(firstAgent.cwd)).toBe(true);
+
+        await ctx.client.archiveAgent(firstAgent.id);
+
+        await new Promise((resolve) => setTimeout(resolve, 300));
+        expect(existsSync(firstAgent.cwd)).toBe(true);
+
+        await ctx.client.archiveAgent(secondAgent.id);
+
+        await waitForPathRemoved({
+          targetPath: firstAgent.cwd,
+          timeoutMs: 10000,
+          label: "worktree after final agent archive",
+        });
+
+        rmSync(repoRoot, { recursive: true, force: true });
+      },
+      60000
+    );
   });
 
 

packages/server/src/server/session.ts

@@ -1767,14 +1767,16 @@ export class Session {
     );
 
     const archivedAt = new Date().toISOString();
+    let archivedRecord: StoredAgentRecord | null = null;
 
     try {
       const existing = await this.agentStorage.get(agentId);
       if (existing) {
-        await this.agentStorage.upsert({
+        archivedRecord = {
           ...existing,
           archivedAt,
-        });
+        };
+        await this.agentStorage.upsert(archivedRecord);
       }
       this.agentManager.notifyAgentState(agentId);
     } catch (error: any) {
@@ -1792,6 +1794,14 @@ export class Session {
         requestId,
       },
     });
+
+    if (archivedRecord) {
+      await this.maybeArchiveWorktreeAfterLastAgentArchived({
+        archivedAgentId: agentId,
+        archivedAgentCwd: archivedRecord.cwd,
+        requestId,
+      });
+    }
   }
 
   private async handleUpdateAgentRequest(
@@ -4459,12 +4469,144 @@ export class Session {
     }
   }
 
-	  private async handlePaseoWorktreeArchiveRequest(
+  private async maybeArchiveWorktreeAfterLastAgentArchived(options: {
-	    msg: Extract<SessionInboundMessage, { type: "paseo_worktree_archive_request" }>
+    archivedAgentId: string;
-	  ): Promise<void> {
+    archivedAgentCwd: string;
-	    const { requestId } = msg;
+    requestId: string;
-	    let targetPath = msg.worktreePath;
+  }): Promise<void> {
-	    let repoRoot = msg.repoRoot ?? null;
+    try {
+      const ownership = await isPaseoOwnedWorktreeCwd(options.archivedAgentCwd, {
+        paseoHome: this.paseoHome,
+      });
+      if (!ownership.allowed) {
+        return;
+      }
+
+      const resolvedWorktree = await resolvePaseoWorktreeRootForCwd(options.archivedAgentCwd, {
+        paseoHome: this.paseoHome,
+      });
+      if (!resolvedWorktree) {
+        return;
+      }
+
+      const records = await this.agentStorage.list();
+      const recordsById = new Map(records.map((record) => [record.id, record]));
+      const targetPath = resolvedWorktree.worktreePath;
+      const hasRemainingNonArchivedRecord = records.some((record) => {
+        if (record.id === options.archivedAgentId || record.archivedAt) {
+          return false;
+        }
+        return this.isPathWithinRoot(targetPath, record.cwd);
+      });
+      if (hasRemainingNonArchivedRecord) {
+        return;
+      }
+
+      const hasUnknownLiveAgent = this.agentManager.listAgents().some((agent) => {
+        if (agent.id === options.archivedAgentId) {
+          return false;
+        }
+        if (!this.isPathWithinRoot(targetPath, agent.cwd)) {
+          return false;
+        }
+        return !recordsById.has(agent.id);
+      });
+      if (hasUnknownLiveAgent) {
+        return;
+      }
+
+      const repoRoot = ownership.repoRoot;
+      if (!repoRoot) {
+        this.sessionLogger.warn(
+          { agentId: options.archivedAgentId, worktreePath: targetPath },
+          "Unable to resolve repo root for auto-archive after agent archive"
+        );
+        return;
+      }
+
+      await this.archivePaseoWorktree({
+        targetPath,
+        repoRoot,
+        requestId: options.requestId,
+      });
+    } catch (error: any) {
+      this.sessionLogger.warn(
+        { err: error, agentId: options.archivedAgentId, cwd: options.archivedAgentCwd },
+        "Failed to auto-archive worktree after agent archive"
+      );
+    }
+  }
+
+  private async archivePaseoWorktree(options: {
+    targetPath: string;
+    repoRoot: string;
+    requestId: string;
+  }): Promise<string[]> {
+    let targetPath = options.targetPath;
+    const resolvedWorktree = await resolvePaseoWorktreeRootForCwd(targetPath, {
+      paseoHome: this.paseoHome,
+    });
+    if (resolvedWorktree) {
+      targetPath = resolvedWorktree.worktreePath;
+    }
+
+    const removedAgents = new Set<string>();
+    const agents = this.agentManager.listAgents();
+    for (const agent of agents) {
+      if (this.isPathWithinRoot(targetPath, agent.cwd)) {
+        removedAgents.add(agent.id);
+        try {
+          await this.agentManager.closeAgent(agent.id);
+        } catch {
+          // ignore cleanup errors
+        }
+        try {
+          await this.agentStorage.remove(agent.id);
+        } catch {
+          // ignore cleanup errors
+        }
+      }
+    }
+
+    const registryRecords = await this.agentStorage.list();
+    for (const record of registryRecords) {
+      if (this.isPathWithinRoot(targetPath, record.cwd)) {
+        removedAgents.add(record.id);
+        try {
+          await this.agentStorage.remove(record.id);
+        } catch {
+          // ignore cleanup errors
+        }
+      }
+    }
+
+    await this.killTerminalsUnderPath(targetPath);
+
+    await deletePaseoWorktree({
+      cwd: options.repoRoot,
+      worktreePath: targetPath,
+      paseoHome: this.paseoHome,
+    });
+
+    for (const agentId of removedAgents) {
+      this.emit({
+        type: "agent_deleted",
+        payload: {
+          agentId,
+          requestId: options.requestId,
+        },
+      });
+    }
+
+    return Array.from(removedAgents);
+  }
+
+  private async handlePaseoWorktreeArchiveRequest(
+    msg: Extract<SessionInboundMessage, { type: "paseo_worktree_archive_request" }>
+  ): Promise<void> {
+    const { requestId } = msg;
+    let targetPath = msg.worktreePath;
+    let repoRoot = msg.repoRoot ?? null;
 
     try {
       if (!targetPath) {
@@ -4479,7 +4621,9 @@ export class Session {
         targetPath = match.path;
       }
 
-      const ownership = await isPaseoOwnedWorktreeCwd(targetPath, { paseoHome: this.paseoHome });
+      const ownership = await isPaseoOwnedWorktreeCwd(targetPath, {
+        paseoHome: this.paseoHome,
+      });
       if (!ownership.allowed) {
         this.emit({
           type: "paseo_worktree_archive_response",
@@ -4496,71 +4640,22 @@ export class Session {
         return;
       }
 
-	      repoRoot = ownership.repoRoot ?? repoRoot ?? null;
+      repoRoot = ownership.repoRoot ?? repoRoot ?? null;
-	      if (!repoRoot) {
+      if (!repoRoot) {
-	        throw new Error("Unable to resolve repo root for worktree");
+        throw new Error("Unable to resolve repo root for worktree");
-	      }
-
-	      const resolvedWorktree = await resolvePaseoWorktreeRootForCwd(targetPath, {
-	        paseoHome: this.paseoHome,
-	      });
-	      if (resolvedWorktree) {
-	        targetPath = resolvedWorktree.worktreePath;
-	      }
-
-	      const removedAgents = new Set<string>();
-	      const agents = this.agentManager.listAgents();
-	      for (const agent of agents) {
-	        if (this.isPathWithinRoot(targetPath, agent.cwd)) {
-	          removedAgents.add(agent.id);
-          try {
-            await this.agentManager.closeAgent(agent.id);
-          } catch {
-            // ignore cleanup errors
-          }
-          try {
-            await this.agentStorage.remove(agent.id);
-          } catch {
-            // ignore cleanup errors
-          }
-        }
       }
 
-      const registryRecords = await this.agentStorage.list();
+      const removedAgents = await this.archivePaseoWorktree({
-      for (const record of registryRecords) {
+        targetPath,
-        if (this.isPathWithinRoot(targetPath, record.cwd)) {
+        repoRoot,
-          removedAgents.add(record.id);
+        requestId,
-          try {
-            await this.agentStorage.remove(record.id);
-          } catch {
-            // ignore cleanup errors
-          }
-        }
-      }
-
-      await this.killTerminalsUnderPath(targetPath);
-
-      await deletePaseoWorktree({
-        cwd: repoRoot,
-        worktreePath: targetPath,
-        paseoHome: this.paseoHome,
       });
 
-      for (const agentId of removedAgents) {
-        this.emit({
-          type: "agent_deleted",
-          payload: {
-            agentId,
-            requestId,
-          },
-        });
-      }
-
       this.emit({
         type: "paseo_worktree_archive_response",
         payload: {
           success: true,
-          removedAgents: Array.from(removedAgents),
+          removedAgents,
           error: null,
           requestId,
         },