From 86b34c89236cc34fbd467ad56214a11f2778fddf Mon Sep 17 00:00:00 2001 From: Harsh Date: Fri, 6 Mar 2026 14:19:16 +0530 Subject: [PATCH] fix: restrict cleanup to AO-managed worktrees --- .../src/__tests__/session-manager.test.ts | 36 ++++-- packages/core/src/session-manager.ts | 119 +++++++++++------- 2 files changed, 104 insertions(+), 51 deletions(-) diff --git a/packages/core/src/__tests__/session-manager.test.ts b/packages/core/src/__tests__/session-manager.test.ts index a88ad5cfc..f5ff0e86e 100644 --- a/packages/core/src/__tests__/session-manager.test.ts +++ b/packages/core/src/__tests__/session-manager.test.ts @@ -5,7 +5,7 @@ import { tmpdir } from "node:os"; import { randomUUID } from "node:crypto"; import { createSessionManager } from "../session-manager.js"; import { writeMetadata, readMetadata, readMetadataRaw, deleteMetadata } from "../metadata.js"; -import { getSessionsDir, getProjectBaseDir } from "../paths.js"; +import { getSessionsDir, getProjectBaseDir, getWorktreesDir } from "../paths.js"; import { SessionNotRestorableError, WorkspaceMissingError, @@ -193,7 +193,8 @@ describe("spawn", () => { const session = await sm.spawn({ projectId: "my-app", - issueId: "this is a very long issue description that should be truncated to sixty characters maximum", + issueId: + "this is a very long issue description that should be truncated to sixty characters maximum", }); expect(session.branch!.replace("feat/", "").length).toBeLessThanOrEqual(60); @@ -341,9 +342,9 @@ describe("spawn", () => { it("throws when agent override plugin is not found", async () => { const sm = createSessionManager({ config, registry: registryWithMultipleAgents }); - await expect( - sm.spawn({ projectId: "my-app", agent: "nonexistent" }), - ).rejects.toThrow("Agent plugin 'nonexistent' not found"); + await expect(sm.spawn({ projectId: "my-app", agent: "nonexistent" })).rejects.toThrow( + "Agent plugin 'nonexistent' not found", + ); }); it("uses default agent when no override specified", async () => { @@ -957,6 +958,27 @@ describe("get", () => { describe("kill", () => { it("destroys runtime, workspace, and archives metadata", async () => { + const managedWorktree = join( + getWorktreesDir(config.configPath, config.projects["my-app"]!.path), + "ws-1", + ); + writeMetadata(sessionsDir, "app-1", { + worktree: managedWorktree, + branch: "main", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-1")), + }); + + const sm = createSessionManager({ config, registry: mockRegistry }); + await sm.kill("app-1"); + + expect(mockRuntime.destroy).toHaveBeenCalledWith(makeHandle("rt-1")); + expect(mockWorkspace.destroy).toHaveBeenCalledWith(managedWorktree); + expect(readMetadata(sessionsDir, "app-1")).toBeNull(); // archived + deleted + }); + + it("does not destroy workspace outside AO-managed worktree root", async () => { writeMetadata(sessionsDir, "app-1", { worktree: "/tmp/ws", branch: "main", @@ -968,9 +990,7 @@ describe("kill", () => { const sm = createSessionManager({ config, registry: mockRegistry }); await sm.kill("app-1"); - expect(mockRuntime.destroy).toHaveBeenCalledWith(makeHandle("rt-1")); - expect(mockWorkspace.destroy).toHaveBeenCalledWith("/tmp/ws"); - expect(readMetadata(sessionsDir, "app-1")).toBeNull(); // archived + deleted + expect(mockWorkspace.destroy).not.toHaveBeenCalled(); }); it("throws for nonexistent session", async () => { diff --git a/packages/core/src/session-manager.ts b/packages/core/src/session-manager.ts index 126d4934e..1d60a91f4 100644 --- a/packages/core/src/session-manager.ts +++ b/packages/core/src/session-manager.ts @@ -11,8 +11,8 @@ * Reference: scripts/claude-ao-session, scripts/send-to-session */ -import { statSync, existsSync, readdirSync, writeFileSync, mkdirSync } from "node:fs"; -import { join } from "node:path"; +import { statSync, existsSync, readdirSync, writeFileSync, mkdirSync, realpathSync } from "node:fs"; +import { join, resolve } from "node:path"; import { isIssueNotFoundError, isRestorable, @@ -50,6 +50,7 @@ import { import { buildPrompt } from "./prompt-builder.js"; import { getSessionsDir, + getWorktreesDir, getProjectBaseDir, generateTmuxName, generateConfigHash, @@ -173,6 +174,38 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { return getSessionsDir(config.configPath, project.path); } + function getProjectWorktreesDir(project: ProjectConfig): string { + return getWorktreesDir(config.configPath, project.path); + } + + function canonicalPath(path: string): string { + try { + return realpathSync(path); + } catch { + return resolve(path); + } + } + + function normalizePath(path: string): string { + return canonicalPath(path).replace(/\/$/, ""); + } + + function isPathInside(path: string, parentPath: string): boolean { + const normalizedPath = normalizePath(path); + const normalizedParent = normalizePath(parentPath); + return normalizedPath === normalizedParent || normalizedPath.startsWith(`${normalizedParent}/`); + } + + function shouldDestroyWorkspacePath( + project: ProjectConfig | undefined, + workspacePath: string, + ): boolean { + if (!project) return true; + const isProjectPath = normalizePath(workspacePath) === normalizePath(project.path); + if (isProjectPath) return false; + return isPathInside(workspacePath, getProjectWorktreesDir(project)); + } + /** * List all session files across all projects (or filtered by projectId). * Scans project-specific directories under ~/.agent-orchestrator/{hash}-{projectId}/sessions/ @@ -213,7 +246,10 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { /** Resolve which plugins to use for a project. */ function resolvePlugins(project: ProjectConfig, agentOverride?: string) { const runtime = registry.get("runtime", project.runtime ?? config.defaults.runtime); - const agent = registry.get("agent", agentOverride ?? project.agent ?? config.defaults.agent); + const agent = registry.get( + "agent", + agentOverride ?? project.agent ?? config.defaults.agent, + ); const workspace = registry.get( "workspace", project.workspace ?? config.defaults.workspace, @@ -251,9 +287,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { * Enrich session with live runtime state (alive/exited) and activity detection. * Mutates the session object in place. */ - const TERMINAL_SESSION_STATUSES = new Set([ - "killed", "done", "merged", "terminated", "cleanup", - ]); + const TERMINAL_SESSION_STATUSES = new Set(["killed", "done", "merged", "terminated", "cleanup"]); async function enrichSessionWithRuntimeState( session: Session, @@ -398,8 +432,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { // If the issueId is already branch-safe (e.g. "INT-9999"), use as-is. // Otherwise sanitize free-text (e.g. "fix login bug") into a valid slug. const id = spawnConfig.issueId; - const isBranchSafe = - /^[A-Za-z0-9][A-Za-z0-9._-]*$/.test(id) && !id.includes(".."); + const isBranchSafe = /^[A-Za-z0-9][A-Za-z0-9._-]*$/.test(id) && !id.includes(".."); const slug = isBranchSafe ? id : id @@ -429,7 +462,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { try { await plugins.workspace.postCreate(wsInfo, project); } catch (err) { - if (workspacePath !== project.path) { + if (shouldDestroyWorkspacePath(project, workspacePath)) { try { await plugins.workspace.destroy(workspacePath); } catch { @@ -498,7 +531,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { }); } catch (err) { // Clean up workspace and reserved ID if agent config or runtime creation failed - if (plugins.workspace && workspacePath !== project.path) { + if (plugins.workspace && shouldDestroyWorkspacePath(project, workspacePath)) { try { await plugins.workspace.destroy(workspacePath); } catch { @@ -553,7 +586,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { } catch { /* best effort */ } - if (plugins.workspace && workspacePath !== project.path) { + if (plugins.workspace && shouldDestroyWorkspacePath(project, workspacePath)) { try { await plugins.workspace.destroy(workspacePath); } catch { @@ -712,36 +745,41 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { async function list(projectId?: string): Promise { const allSessions = listAllSessions(projectId); - const sessionPromises = allSessions.map(async ({ sessionName, projectId: sessionProjectId }) => { - const project = config.projects[sessionProjectId]; - if (!project) return null; + const sessionPromises = allSessions.map( + async ({ sessionName, projectId: sessionProjectId }) => { + const project = config.projects[sessionProjectId]; + if (!project) return null; - const sessionsDir = getProjectSessionsDir(project); - const raw = readMetadataRaw(sessionsDir, sessionName); - if (!raw) return null; + const sessionsDir = getProjectSessionsDir(project); + const raw = readMetadataRaw(sessionsDir, sessionName); + if (!raw) return null; - // Get file timestamps for createdAt/lastActivityAt - let createdAt: Date | undefined; - let modifiedAt: Date | undefined; - try { - const metaPath = join(sessionsDir, sessionName); - const stats = statSync(metaPath); - createdAt = stats.birthtime; - modifiedAt = stats.mtime; - } catch { - // If stat fails, timestamps will fall back to current time - } + // Get file timestamps for createdAt/lastActivityAt + let createdAt: Date | undefined; + let modifiedAt: Date | undefined; + try { + const metaPath = join(sessionsDir, sessionName); + const stats = statSync(metaPath); + createdAt = stats.birthtime; + modifiedAt = stats.mtime; + } catch { + // If stat fails, timestamps will fall back to current time + } - const session = metadataToSession(sessionName, raw, createdAt, modifiedAt); + const session = metadataToSession(sessionName, raw, createdAt, modifiedAt); - const plugins = resolvePlugins(project, raw["agent"]); - // Cap per-session enrichment at 2s — subprocess calls (tmux/ps) can be - // slow under load. If we time out, session keeps its metadata values. - const enrichTimeout = new Promise((resolve) => setTimeout(resolve, 2_000)); - await Promise.race([ensureHandleAndEnrich(session, sessionName, project, plugins), enrichTimeout]); + const plugins = resolvePlugins(project, raw["agent"]); + // Cap per-session enrichment at 2s — subprocess calls (tmux/ps) can be + // slow under load. If we time out, session keeps its metadata values. + const enrichTimeout = new Promise((resolve) => setTimeout(resolve, 2_000)); + await Promise.race([ + ensureHandleAndEnrich(session, sessionName, project, plugins), + enrichTimeout, + ]); - return session; - }); + return session; + }, + ); const results = await Promise.all(sessionPromises); return results.filter((s): s is Session => s !== null); @@ -817,10 +855,8 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { } } - // Destroy workspace — skip if worktree is the project path (no isolation was used) const worktree = raw["worktree"]; - const isProjectPath = project && worktree === project.path; - if (worktree && !isProjectPath) { + if (worktree && shouldDestroyWorkspacePath(project, worktree)) { const workspacePlugin = project ? resolvePlugins(project).workspace : registry.get("workspace", config.defaults.workspace); @@ -849,10 +885,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { // Never clean up orchestrator sessions — they manage the lifecycle. // Check explicit role metadata first, fall back to naming convention // for pre-existing sessions spawned before the role field was added. - if ( - session.metadata["role"] === "orchestrator" || - session.id.endsWith("-orchestrator") - ) { + if (session.metadata["role"] === "orchestrator" || session.id.endsWith("-orchestrator")) { result.skipped.push(session.id); continue; }