Merge pull request #322 from ComposioHQ/fix/worktree-cleanup-safety-v2

fix: restrict cleanup to AO-managed worktree roots
This commit is contained in:
Harsh Batheja 2026-03-06 15:12:02 +05:30 committed by GitHub
commit 8a655e7625
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 150 additions and 52 deletions

View File

@ -1,11 +1,11 @@
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { mkdirSync, rmSync, writeFileSync, existsSync } from "node:fs";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { homedir, 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),
"app-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 paths outside managed roots", async () => {
writeMetadata(sessionsDir, "app-1", {
worktree: "/tmp/ws",
branch: "main",
@ -968,9 +990,38 @@ 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("destroys workspace under legacy ~/.worktrees root", async () => {
const legacyWorktree = join(homedir(), ".worktrees", "my-app", "app-1");
writeMetadata(sessionsDir, "app-1", {
worktree: legacyWorktree,
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(mockWorkspace.destroy).toHaveBeenCalledWith(legacyWorktree);
});
it("never destroys workspace equal to project path", async () => {
writeMetadata(sessionsDir, "app-1", {
worktree: config.projects["my-app"]!.path,
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(mockWorkspace.destroy).not.toHaveBeenCalled();
});
it("throws for nonexistent session", async () => {

View File

@ -12,7 +12,8 @@
*/
import { statSync, existsSync, readdirSync, writeFileSync, mkdirSync } from "node:fs";
import { join } from "node:path";
import { basename, join, resolve } from "node:path";
import { homedir } from "node:os";
import {
isIssueNotFoundError,
isRestorable,
@ -50,6 +51,7 @@ import {
import { buildPrompt } from "./prompt-builder.js";
import {
getSessionsDir,
getWorktreesDir,
getProjectBaseDir,
generateTmuxName,
generateConfigHash,
@ -173,6 +175,43 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
return getSessionsDir(config.configPath, project.path);
}
function normalizePath(path: string): string {
return resolve(path).replace(/\/$/, "");
}
function isPathInside(path: string, parentPath: string): boolean {
const normalizedPath = normalizePath(path);
const normalizedParent = normalizePath(parentPath);
return normalizedPath === normalizedParent || normalizedPath.startsWith(`${normalizedParent}/`);
}
function getManagedWorkspaceRoots(project: ProjectConfig, projectId?: string): string[] {
const roots = [getWorktreesDir(config.configPath, project.path)];
const legacyIds = new Set<string>();
if (projectId) {
legacyIds.add(projectId);
}
legacyIds.add(basename(project.path));
for (const id of legacyIds) {
roots.push(join(homedir(), ".worktrees", id));
}
return roots;
}
function shouldDestroyWorkspacePath(
project: ProjectConfig | undefined,
projectId: string | undefined,
workspacePath: string,
): boolean {
if (!project) return false;
if (normalizePath(workspacePath) === normalizePath(project.path)) return false;
const roots = getManagedWorkspaceRoots(project, projectId);
return roots.some((root) => isPathInside(workspacePath, root));
}
/**
* List all session files across all projects (or filtered by projectId).
* Scans project-specific directories under ~/.agent-orchestrator/{hash}-{projectId}/sessions/
@ -213,7 +252,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>("runtime", project.runtime ?? config.defaults.runtime);
const agent = registry.get<Agent>("agent", agentOverride ?? project.agent ?? config.defaults.agent);
const agent = registry.get<Agent>(
"agent",
agentOverride ?? project.agent ?? config.defaults.agent,
);
const workspace = registry.get<Workspace>(
"workspace",
project.workspace ?? config.defaults.workspace,
@ -251,9 +293,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 +438,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 +468,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
try {
await plugins.workspace.postCreate(wsInfo, project);
} catch (err) {
if (workspacePath !== project.path) {
if (shouldDestroyWorkspacePath(project, spawnConfig.projectId, workspacePath)) {
try {
await plugins.workspace.destroy(workspacePath);
} catch {
@ -498,7 +537,10 @@ 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, spawnConfig.projectId, workspacePath)
) {
try {
await plugins.workspace.destroy(workspacePath);
} catch {
@ -553,7 +595,10 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
} catch {
/* best effort */
}
if (plugins.workspace && workspacePath !== project.path) {
if (
plugins.workspace &&
shouldDestroyWorkspacePath(project, spawnConfig.projectId, workspacePath)
) {
try {
await plugins.workspace.destroy(workspacePath);
} catch {
@ -712,36 +757,41 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
async function list(projectId?: string): Promise<Session[]> {
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<void>((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<void>((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);
@ -782,14 +832,16 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
let raw: Record<string, string> | null = null;
let sessionsDir: string | null = null;
let project: ProjectConfig | undefined;
let projectId: string | undefined;
for (const proj of Object.values(config.projects)) {
for (const [projId, proj] of Object.entries(config.projects)) {
const dir = getProjectSessionsDir(proj);
const metadata = readMetadataRaw(dir, sessionId);
if (metadata) {
raw = metadata;
sessionsDir = dir;
project = proj;
projectId = projId;
break;
}
}
@ -817,10 +869,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, projectId, worktree)) {
const workspacePlugin = project
? resolvePlugins(project).workspace
: registry.get<Workspace>("workspace", config.defaults.workspace);
@ -849,10 +899,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;
}