feat: implement session restore for crashed/exited agents

Add true in-place session restore: same session ID, same worktree, same
metadata — optionally resuming the Claude Code conversation via --resume.

Core changes:
- Add TERMINAL_STATUSES, TERMINAL_ACTIVITIES, NON_RESTORABLE_STATUSES sets
  and isTerminalSession/isRestorable helpers to types.ts
- Add SessionNotRestorableError and WorkspaceMissingError error classes
- Add restore() to SessionManager with 9-step flow: find metadata →
  validate restorability → check/recreate workspace → get restore or
  launch command → create runtime → update metadata
- Add restoredAt field to Session and SessionMetadata

Plugin extensions:
- workspace-worktree: exists() + restore() (git worktree prune + re-add)
- workspace-clone: exists() + restore() (git clone + checkout)
- scm-github: branchExists() via git rev-parse
- agent-claude-code: getRestoreCommand() finds latest JSONL session file
  and builds claude --resume command

CLI + Web:
- Add `ao session restore <id>` subcommand
- Web restore API route uses sessionManager.restore() instead of spawn()
- SessionCard uses centralized TERMINAL_STATUSES/TERMINAL_ACTIVITIES
- Web types re-export core constants with sync tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Prateek 2026-02-18 22:04:53 +05:30
parent de6653e258
commit 6cb18fc6b9
14 changed files with 760 additions and 68 deletions

View File

@ -1,6 +1,6 @@
import chalk from "chalk";
import type { Command } from "commander";
import { loadConfig } from "@composio/ao-core";
import { loadConfig, SessionNotRestorableError, WorkspaceMissingError } from "@composio/ao-core";
import { git, getTmuxActivity } from "../lib/shell.js";
import { formatAge } from "../lib/format.js";
import { getSessionManager } from "../lib/create-session-manager.js";
@ -146,10 +146,39 @@ export function registerSession(program: Command): void {
console.error(chalk.red(` Error cleaning ${sessionId}: ${error}`));
}
}
console.log(
chalk.green(`\nCleanup complete. ${result.killed.length} sessions cleaned.`),
);
console.log(chalk.green(`\nCleanup complete. ${result.killed.length} sessions cleaned.`));
}
}
});
session
.command("restore")
.description("Restore a terminated/crashed session in-place")
.argument("<session>", "Session name to restore")
.action(async (sessionName: string) => {
const config = loadConfig();
const sm = await getSessionManager(config);
try {
const restored = await sm.restore(sessionName);
console.log(chalk.green(`\nSession ${sessionName} restored.`));
if (restored.workspacePath) {
console.log(chalk.dim(` Worktree: ${restored.workspacePath}`));
}
if (restored.branch) {
console.log(chalk.dim(` Branch: ${restored.branch}`));
}
const tmuxTarget = restored.runtimeHandle?.id ?? sessionName;
console.log(chalk.dim(` Attach: tmux attach -t ${tmuxTarget}`));
} catch (err) {
if (err instanceof SessionNotRestorableError) {
console.error(chalk.red(`Cannot restore: ${err.reason}`));
} else if (err instanceof WorkspaceMissingError) {
console.error(chalk.red(`Workspace missing: ${err.message}`));
} else {
console.error(chalk.red(`Failed to restore session ${sessionName}: ${err}`));
}
process.exit(1);
}
});
}

View File

@ -104,6 +104,7 @@ beforeEach(() => {
mockSessionManager = {
spawn: vi.fn(),
spawnOrchestrator: vi.fn(),
restore: vi.fn(),
list: vi.fn().mockResolvedValue([]),
get: vi.fn().mockResolvedValue(null),
kill: vi.fn().mockResolvedValue(undefined),

View File

@ -4,17 +4,19 @@ import { join } from "node:path";
import { tmpdir } from "node:os";
import { randomUUID } from "node:crypto";
import { createSessionManager } from "../session-manager.js";
import { writeMetadata, readMetadata } from "../metadata.js";
import { writeMetadata, readMetadata, readMetadataRaw } from "../metadata.js";
import { getSessionsDir, getProjectBaseDir } from "../paths.js";
import type {
OrchestratorConfig,
PluginRegistry,
Runtime,
Agent,
Workspace,
Tracker,
SCM,
RuntimeHandle,
import {
SessionNotRestorableError,
WorkspaceMissingError,
type OrchestratorConfig,
type PluginRegistry,
type Runtime,
type Agent,
type Workspace,
type Tracker,
type SCM,
type RuntimeHandle,
} from "../types.js";
let tmpDir: string;
@ -875,6 +877,248 @@ describe("spawnOrchestrator", () => {
});
});
describe("restore", () => {
it("restores a killed session with existing workspace", async () => {
// Create a workspace directory that exists
const wsPath = join(tmpDir, "ws-app-1");
mkdirSync(wsPath, { recursive: true });
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
issue: "TEST-1",
pr: "https://github.com/org/my-app/pull/10",
createdAt: "2025-01-01T00:00:00.000Z",
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: mockRegistry });
const restored = await sm.restore("app-1");
expect(restored.id).toBe("app-1");
expect(restored.status).toBe("spawning");
expect(restored.activity).toBe("active");
expect(restored.workspacePath).toBe(wsPath);
expect(restored.branch).toBe("feat/TEST-1");
expect(restored.runtimeHandle).toEqual(makeHandle("rt-1"));
expect(restored.restoredAt).toBeInstanceOf(Date);
// Verify runtime was created
expect(mockRuntime.create).toHaveBeenCalled();
// Verify metadata was updated (not rewritten)
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["status"]).toBe("spawning");
expect(meta!["restoredAt"]).toBeDefined();
// Verify original fields are preserved
expect(meta!["issue"]).toBe("TEST-1");
expect(meta!["pr"]).toBe("https://github.com/org/my-app/pull/10");
expect(meta!["createdAt"]).toBe("2025-01-01T00:00:00.000Z");
});
it("recreates workspace when missing and plugin supports restore", async () => {
const wsPath = join(tmpDir, "ws-app-1");
// DO NOT create the directory — it's missing
const mockWorkspaceWithRestore: Workspace = {
...mockWorkspace,
exists: vi.fn().mockResolvedValue(false),
restore: vi.fn().mockResolvedValue({
path: wsPath,
branch: "feat/TEST-1",
sessionId: "app-1",
projectId: "my-app",
}),
};
const registryWithRestore: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgent;
if (slot === "workspace") return mockWorkspaceWithRestore;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "terminated",
project: "my-app",
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: registryWithRestore });
const restored = await sm.restore("app-1");
expect(restored.id).toBe("app-1");
expect(mockWorkspaceWithRestore.restore).toHaveBeenCalled();
expect(mockRuntime.create).toHaveBeenCalled();
});
it("throws SessionNotRestorableError for merged sessions", async () => {
writeMetadata(sessionsDir, "app-1", {
worktree: "/tmp",
branch: "main",
status: "merged",
project: "my-app",
});
const sm = createSessionManager({ config, registry: mockRegistry });
await expect(sm.restore("app-1")).rejects.toThrow(SessionNotRestorableError);
});
it("throws SessionNotRestorableError for working sessions", async () => {
writeMetadata(sessionsDir, "app-1", {
worktree: "/tmp",
branch: "main",
status: "working",
project: "my-app",
});
const sm = createSessionManager({ config, registry: mockRegistry });
await expect(sm.restore("app-1")).rejects.toThrow(SessionNotRestorableError);
});
it("throws WorkspaceMissingError when workspace gone and no restore method", async () => {
const wsPath = join(tmpDir, "nonexistent-ws");
const mockWorkspaceNoRestore: Workspace = {
...mockWorkspace,
exists: vi.fn().mockResolvedValue(false),
// No restore method
};
const registryNoRestore: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgent;
if (slot === "workspace") return mockWorkspaceNoRestore;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: registryNoRestore });
await expect(sm.restore("app-1")).rejects.toThrow(WorkspaceMissingError);
});
it("throws for nonexistent session", async () => {
const sm = createSessionManager({ config, registry: mockRegistry });
await expect(sm.restore("nonexistent")).rejects.toThrow("not found");
});
it("uses getRestoreCommand when available", async () => {
const wsPath = join(tmpDir, "ws-app-1");
mkdirSync(wsPath, { recursive: true });
const mockAgentWithRestore: Agent = {
...mockAgent,
getRestoreCommand: vi.fn().mockResolvedValue("claude --resume abc123"),
};
const registryWithAgentRestore: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgentWithRestore;
if (slot === "workspace") return mockWorkspace;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "errored",
project: "my-app",
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: registryWithAgentRestore });
await sm.restore("app-1");
expect(mockAgentWithRestore.getRestoreCommand).toHaveBeenCalled();
// Verify runtime.create was called with the restore command
const createCall = (mockRuntime.create as ReturnType<typeof vi.fn>).mock.calls[0][0];
expect(createCall.launchCommand).toBe("claude --resume abc123");
});
it("falls back to getLaunchCommand when getRestoreCommand returns null", async () => {
const wsPath = join(tmpDir, "ws-app-1");
mkdirSync(wsPath, { recursive: true });
const mockAgentWithNullRestore: Agent = {
...mockAgent,
getRestoreCommand: vi.fn().mockResolvedValue(null),
};
const registryWithNullRestore: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgentWithNullRestore;
if (slot === "workspace") return mockWorkspace;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: registryWithNullRestore });
await sm.restore("app-1");
expect(mockAgentWithNullRestore.getRestoreCommand).toHaveBeenCalled();
expect(mockAgent.getLaunchCommand).toHaveBeenCalled();
const createCall = (mockRuntime.create as ReturnType<typeof vi.fn>).mock.calls[0][0];
expect(createCall.launchCommand).toBe("mock-agent --start");
});
it("preserves original createdAt/issue/PR metadata", async () => {
const wsPath = join(tmpDir, "ws-app-1");
mkdirSync(wsPath, { recursive: true });
const originalCreatedAt = "2024-06-15T10:00:00.000Z";
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-42",
status: "killed",
project: "my-app",
issue: "TEST-42",
pr: "https://github.com/org/my-app/pull/99",
summary: "Implementing feature X",
createdAt: originalCreatedAt,
runtimeHandle: JSON.stringify(makeHandle("rt-old")),
});
const sm = createSessionManager({ config, registry: mockRegistry });
await sm.restore("app-1");
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["createdAt"]).toBe(originalCreatedAt);
expect(meta!["issue"]).toBe("TEST-42");
expect(meta!["pr"]).toBe("https://github.com/org/my-app/pull/99");
expect(meta!["summary"]).toBe("Implementing feature X");
expect(meta!["branch"]).toBe("feat/TEST-42");
});
});
describe("PluginRegistry.loadBuiltins importFn", () => {
it("should use provided importFn instead of built-in import", async () => {
const { createPluginRegistry: createReg } = await import("../plugin-registry.js");

View File

@ -15,6 +15,10 @@ import { statSync, existsSync, readdirSync, writeFileSync, mkdirSync } from "nod
import { join } from "node:path";
import {
isIssueNotFoundError,
isRestorable,
NON_RESTORABLE_STATUSES,
SessionNotRestorableError,
WorkspaceMissingError,
type SessionManager,
type Session,
type SessionId,
@ -37,12 +41,19 @@ import {
import {
readMetadataRaw,
writeMetadata,
updateMetadata,
deleteMetadata,
listMetadata,
reserveSessionId,
} from "./metadata.js";
import { buildPrompt } from "./prompt-builder.js";
import { getSessionsDir, getProjectBaseDir, generateTmuxName, generateConfigHash, validateAndStoreOrigin } from "./paths.js";
import {
getSessionsDir,
getProjectBaseDir,
generateTmuxName,
generateConfigHash,
validateAndStoreOrigin,
} from "./paths.js";
/** Escape regex metacharacters in a string. */
function escapeRegex(str: string): string {
@ -266,10 +277,7 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
// external scripts (which don't store runtimeHandle) to always show "unknown".
if (plugins.agent) {
try {
const detected = await plugins.agent.getActivityState(
session,
config.readyThresholdMs,
);
const detected = await plugins.agent.getActivityState(session, config.readyThresholdMs);
if (detected !== null) {
session.activity = detected;
}
@ -762,7 +770,10 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
deleteMetadata(sessionsDir, sessionId, true);
}
async function cleanup(projectId?: string, options?: { dryRun?: boolean }): Promise<CleanupResult> {
async function cleanup(
projectId?: string,
options?: { dryRun?: boolean },
): Promise<CleanupResult> {
const result: CleanupResult = { killed: [], skipped: [], errors: [] };
const sessions = await list(projectId);
@ -869,5 +880,146 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager {
await runtimePlugin.sendMessage(handle, message);
}
return { spawn, spawnOrchestrator, list, get, kill, cleanup, send };
async function restore(sessionId: SessionId): Promise<Session> {
// 1. Find session metadata across all projects
let raw: Record<string, string> | null = null;
let sessionsDir: string | null = null;
let project: ProjectConfig | undefined;
let projectId: string | undefined;
for (const [key, proj] of Object.entries(config.projects)) {
const dir = getProjectSessionsDir(proj);
const metadata = readMetadataRaw(dir, sessionId);
if (metadata) {
raw = metadata;
sessionsDir = dir;
project = proj;
projectId = key;
break;
}
}
if (!raw || !sessionsDir || !project || !projectId) {
throw new Error(`Session ${sessionId} not found`);
}
// 2. Reconstruct Session from metadata
const session = metadataToSession(sessionId, raw);
// 3. Validate restorability
if (!isRestorable(session)) {
if (NON_RESTORABLE_STATUSES.has(session.status)) {
throw new SessionNotRestorableError(sessionId, `status is "${session.status}"`);
}
throw new SessionNotRestorableError(sessionId, "session is not in a terminal state");
}
// 4. Resolve plugins
const plugins = resolvePlugins(project);
if (!plugins.runtime) {
throw new Error(`Runtime plugin '${project.runtime ?? config.defaults.runtime}' not found`);
}
if (!plugins.agent) {
throw new Error(`Agent plugin '${project.agent ?? config.defaults.agent}' not found`);
}
// 5. Check workspace
const workspacePath = raw["worktree"] || project.path;
const workspaceExists = plugins.workspace?.exists
? await plugins.workspace.exists(workspacePath)
: existsSync(workspacePath);
if (!workspaceExists) {
// Try to restore workspace if plugin supports it
if (plugins.workspace?.restore && session.branch) {
try {
const wsInfo = await plugins.workspace.restore(
{
projectId,
project,
sessionId,
branch: session.branch,
},
workspacePath,
);
// Run post-create hooks on restored workspace
if (plugins.workspace.postCreate) {
await plugins.workspace.postCreate(wsInfo, project);
}
} catch (err) {
throw new WorkspaceMissingError(
workspacePath,
`restore failed: ${err instanceof Error ? err.message : String(err)}`,
);
}
} else {
throw new WorkspaceMissingError(workspacePath, "workspace plugin does not support restore");
}
}
// 6. Get launch command — try restore command first, fall back to fresh launch
let launchCommand: string;
const agentLaunchConfig = {
sessionId,
projectConfig: project,
issueId: session.issueId ?? undefined,
permissions: project.agentConfig?.permissions,
model: project.agentConfig?.model,
};
if (plugins.agent.getRestoreCommand) {
const restoreCmd = await plugins.agent.getRestoreCommand(session, project);
launchCommand = restoreCmd ?? plugins.agent.getLaunchCommand(agentLaunchConfig);
} else {
launchCommand = plugins.agent.getLaunchCommand(agentLaunchConfig);
}
const environment = plugins.agent.getEnvironment(agentLaunchConfig);
// 7. Create runtime (reuse tmuxName from metadata)
const tmuxName = raw["tmuxName"];
const handle = await plugins.runtime.create({
sessionId: tmuxName ?? sessionId,
workspacePath,
launchCommand,
environment: {
...environment,
AO_SESSION: sessionId,
AO_DATA_DIR: sessionsDir,
AO_SESSION_NAME: sessionId,
...(tmuxName && { AO_TMUX_NAME: tmuxName }),
},
});
// 8. Update metadata — merge updates, preserving existing fields
const now = new Date().toISOString();
updateMetadata(sessionsDir, sessionId, {
status: "spawning",
runtimeHandle: JSON.stringify(handle),
restoredAt: now,
});
// 9. Run postLaunchSetup (non-fatal)
const restoredSession: Session = {
...session,
status: "spawning",
activity: "active",
workspacePath,
runtimeHandle: handle,
restoredAt: new Date(now),
};
if (plugins.agent.postLaunchSetup) {
try {
await plugins.agent.postLaunchSetup(restoredSession);
} catch {
// Non-fatal — session is already running
}
}
return restoredSession;
}
return { spawn, spawnOrchestrator, restore, list, get, kill, cleanup, send };
}

View File

@ -83,6 +83,40 @@ export const SESSION_STATUS = {
TERMINATED: "terminated" as const,
} satisfies Record<string, SessionStatus>;
/** Statuses that indicate the session is in a terminal (dead) state. */
export const TERMINAL_STATUSES: ReadonlySet<SessionStatus> = new Set([
"killed",
"terminated",
"done",
"cleanup",
"errored",
]);
/** Activity states that indicate the session is no longer running. */
export const TERMINAL_ACTIVITIES: ReadonlySet<ActivityState> = new Set(["exited"]);
/** Statuses that must never be restored (e.g. already merged, actively working). */
export const NON_RESTORABLE_STATUSES: ReadonlySet<SessionStatus> = new Set(["merged", "working"]);
/** Check if a session is in a terminal (dead) state. */
export function isTerminalSession(session: {
status: SessionStatus;
activity: ActivityState | null;
}): boolean {
return (
TERMINAL_STATUSES.has(session.status) ||
(session.activity !== null && TERMINAL_ACTIVITIES.has(session.activity))
);
}
/** Check if a session can be restored. */
export function isRestorable(session: {
status: SessionStatus;
activity: ActivityState | null;
}): boolean {
return isTerminalSession(session) && !NON_RESTORABLE_STATUSES.has(session.status);
}
/** A running agent session */
export interface Session {
/** Unique session ID, e.g. "my-app-3" */
@ -121,6 +155,9 @@ export interface Session {
/** Last activity timestamp */
lastActivityAt: Date;
/** When this session was last restored (undefined if never restored) */
restoredAt?: Date;
/** Metadata key-value pairs */
metadata: Record<string, string>;
}
@ -243,6 +280,12 @@ export interface Agent {
/** Extract information from agent's internal data (summary, cost, session ID) */
getSessionInfo(session: Session): Promise<AgentSessionInfo | null>;
/**
* Optional: get a launch command that resumes a previous session.
* Returns null if no previous session is found (caller falls back to getLaunchCommand).
*/
getRestoreCommand?(session: Session, project: ProjectConfig): Promise<string | null>;
/** Optional: run setup after agent is launched (e.g. configure MCP servers) */
postLaunchSetup?(session: Session): Promise<void>;
@ -339,6 +382,12 @@ export interface Workspace {
/** Optional: run hooks after workspace creation (symlinks, installs, etc.) */
postCreate?(info: WorkspaceInfo, project: ProjectConfig): Promise<void>;
/** Optional: check if a workspace exists and is a valid git repo */
exists?(workspacePath: string): Promise<boolean>;
/** Optional: restore a workspace (e.g. recreate a worktree for an existing branch) */
restore?(config: WorkspaceCreateConfig, workspacePath: string): Promise<WorkspaceInfo>;
}
export interface WorkspaceCreateConfig {
@ -485,6 +534,9 @@ export interface SCM {
/** Check if PR is ready to merge */
getMergeability(pr: PRInfo): Promise<MergeReadiness>;
/** Optional: check if a branch exists in the repo */
branchExists?(repoPath: string, branch: string): Promise<boolean>;
}
// --- PR Types ---
@ -904,6 +956,7 @@ export interface SessionMetadata {
project?: string;
createdAt?: string;
runtimeHandle?: string;
restoredAt?: string;
dashboardPort?: number;
terminalWsPort?: number;
directTerminalWsPort?: number;
@ -917,6 +970,7 @@ export interface SessionMetadata {
export interface SessionManager {
spawn(config: SessionSpawnConfig): Promise<Session>;
spawnOrchestrator(config: OrchestratorSpawnConfig): Promise<Session>;
restore(sessionId: SessionId): Promise<Session>;
list(projectId?: string): Promise<Session[]>;
get(sessionId: SessionId): Promise<Session | null>;
kill(sessionId: SessionId): Promise<void>;
@ -996,3 +1050,25 @@ export function isIssueNotFoundError(err: unknown): boolean {
message.includes("no issue with identifier")
);
}
/** Thrown when a session cannot be restored (e.g. merged, still working). */
export class SessionNotRestorableError extends Error {
constructor(
public readonly sessionId: string,
public readonly reason: string,
) {
super(`Session ${sessionId} cannot be restored: ${reason}`);
this.name = "SessionNotRestorableError";
}
}
/** Thrown when a workspace is missing and cannot be recreated. */
export class WorkspaceMissingError extends Error {
constructor(
public readonly path: string,
public readonly detail?: string,
) {
super(`Workspace missing at ${path}${detail ? `: ${detail}` : ""}`);
this.name = "WorkspaceMissingError";
}
}

View File

@ -8,6 +8,7 @@ import {
type ActivityState,
type CostEstimate,
type PluginModule,
type ProjectConfig,
type RuntimeHandle,
type Session,
type WorkspaceHooksConfig,
@ -722,6 +723,35 @@ function createClaudeCodeAgent(): Agent {
};
},
async getRestoreCommand(session: Session, project: ProjectConfig): Promise<string | null> {
if (!session.workspacePath) return null;
// Find Claude's project directory for this workspace
const projectPath = toClaudeProjectPath(session.workspacePath);
const projectDir = join(homedir(), ".claude", "projects", projectPath);
// Find the latest session JSONL file
const sessionFile = await findLatestSessionFile(projectDir);
if (!sessionFile) return null;
// Extract session UUID from filename (e.g. "abc123-def456.jsonl" → "abc123-def456")
const sessionUuid = basename(sessionFile, ".jsonl");
if (!sessionUuid) return null;
// Build resume command
const parts: string[] = ["claude", "--resume", shellEscape(sessionUuid)];
if (project.agentConfig?.permissions === "skip") {
parts.push("--dangerously-skip-permissions");
}
if (project.agentConfig?.model) {
parts.push("--model", shellEscape(project.agentConfig.model as string));
}
return parts.join(" ");
},
async setupWorkspaceHooks(workspacePath: string, _config: WorkspaceHooksConfig): Promise<void> {
// Use absolute path for hook command (specific to this workspace)
const hookScriptPath = join(workspacePath, ".claude", "metadata-updater.sh");

View File

@ -560,6 +560,26 @@ function createGitHubSCM(): SCM {
blockers,
};
},
async branchExists(repoPath: string, branch: string): Promise<boolean> {
try {
await execFileAsync("git", ["rev-parse", "--verify", `refs/heads/${branch}`], {
cwd: repoPath,
timeout: 30_000,
});
return true;
} catch {
// Try remote ref
try {
await execFileAsync("git", ["rev-parse", "--verify", `refs/remotes/origin/${branch}`], {
cwd: repoPath,
timeout: 30_000,
});
return true;
} catch {
return false;
}
}
},
};
}

View File

@ -166,6 +166,64 @@ export function create(config?: Record<string, unknown>): Workspace {
return infos;
},
async exists(workspacePath: string): Promise<boolean> {
if (!existsSync(workspacePath)) return false;
try {
await execFileAsync("git", ["rev-parse", "--is-inside-work-tree"], {
cwd: workspacePath,
timeout: 30_000,
});
return true;
} catch {
return false;
}
},
async restore(cfg: WorkspaceCreateConfig, workspacePath: string): Promise<WorkspaceInfo> {
const repoPath = expandPath(cfg.project.path);
// Get remote URL
let remoteUrl: string;
try {
remoteUrl = await git(repoPath, "remote", "get-url", "origin");
} catch {
remoteUrl = repoPath;
}
// Clone fresh
await execFileAsync("git", [
"clone",
"--reference",
repoPath,
"--branch",
cfg.project.defaultBranch,
remoteUrl,
workspacePath,
]);
// Try to checkout the branch
try {
await git(workspacePath, "checkout", cfg.branch);
} catch {
try {
await git(workspacePath, "checkout", "-b", cfg.branch);
} catch (checkoutErr: unknown) {
rmSync(workspacePath, { recursive: true, force: true });
const msg = checkoutErr instanceof Error ? checkoutErr.message : String(checkoutErr);
throw new Error(`Failed to checkout branch "${cfg.branch}" during restore: ${msg}`, {
cause: checkoutErr,
});
}
}
return {
path: workspacePath,
branch: cfg.branch,
sessionId: cfg.sessionId,
projectId: cfg.projectId,
};
},
async postCreate(info: WorkspaceInfo, project: ProjectConfig): Promise<void> {
// Run postCreate hooks
// NOTE: commands run with full shell privileges — they come from trusted YAML config

View File

@ -11,6 +11,9 @@ import type {
ProjectConfig,
} from "@composio/ao-core";
/** Timeout for git commands (30 seconds) */
const GIT_TIMEOUT = 30_000;
const execFileAsync = promisify(execFile);
export const manifest = {
@ -190,6 +193,59 @@ export function create(config?: Record<string, unknown>): Workspace {
return infos;
},
async exists(workspacePath: string): Promise<boolean> {
if (!existsSync(workspacePath)) return false;
try {
await execFileAsync("git", ["rev-parse", "--is-inside-work-tree"], {
cwd: workspacePath,
timeout: GIT_TIMEOUT,
});
return true;
} catch {
return false;
}
},
async restore(cfg: WorkspaceCreateConfig, workspacePath: string): Promise<WorkspaceInfo> {
const repoPath = expandPath(cfg.project.path);
// Prune stale worktree entries
try {
await git(repoPath, "worktree", "prune");
} catch {
// Best effort
}
// Fetch latest
try {
await git(repoPath, "fetch", "origin", "--quiet");
} catch {
// May fail if offline
}
// Try to create worktree on the existing branch
try {
await git(repoPath, "worktree", "add", workspacePath, cfg.branch);
} catch {
// Branch might not exist locally — try from origin
const remoteBranch = `origin/${cfg.branch}`;
try {
await git(repoPath, "worktree", "add", "-b", cfg.branch, workspacePath, remoteBranch);
} catch {
// Last resort: create from default branch
const baseRef = `origin/${cfg.project.defaultBranch}`;
await git(repoPath, "worktree", "add", "-b", cfg.branch, workspacePath, baseRef);
}
}
return {
path: workspacePath,
branch: cfg.branch,
sessionId: cfg.sessionId,
projectId: cfg.projectId,
};
},
async postCreate(info: WorkspaceInfo, project: ProjectConfig): Promise<void> {
const repoPath = expandPath(project.path);

View File

@ -1,11 +1,12 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { NextRequest } from "next/server";
import type {
Session,
SessionManager,
OrchestratorConfig,
PluginRegistry,
SCM,
import {
SessionNotRestorableError,
type Session,
type SessionManager,
type OrchestratorConfig,
type PluginRegistry,
type SCM,
} from "@composio/ao-core";
// ── Mock Data ─────────────────────────────────────────────────────────
@ -82,6 +83,17 @@ const mockSessionManager: SessionManager = {
}),
cleanup: vi.fn(async () => ({ killed: [], skipped: [], errors: [] })),
spawnOrchestrator: vi.fn(),
restore: vi.fn(async (id: string) => {
const session = testSessions.find((s) => s.id === id);
if (!session) {
throw new Error(`Session ${id} not found`);
}
// Simulate SessionNotRestorableError for non-terminal sessions
if (session.status === "working" && session.activity !== "exited") {
throw new SessionNotRestorableError(id, "session is not in a terminal state");
}
return { ...session, status: "spawning" as const, activity: "active" as const };
}),
};
const mockSCM: SCM = {
@ -114,9 +126,9 @@ const mockRegistry: PluginRegistry = {
};
const mockConfig: OrchestratorConfig = {
dataDir: "/tmp/ao-test",
worktreeDir: "/tmp/ao-worktrees",
configPath: "/tmp/ao-test/agent-orchestrator.yaml",
port: 3000,
readyThresholdMs: 300_000,
defaults: { runtime: "tmux", agent: "claude-code", workspace: "worktree", notifiers: [] },
projects: {
"my-app": {

View File

@ -2,13 +2,7 @@ import { type NextRequest, NextResponse } from "next/server";
import { validateIdentifier } from "@/lib/validation";
import { getServices } from "@/lib/services";
import { sessionToDashboard } from "@/lib/serialize";
/** Terminal states that can be restored */
const RESTORABLE_STATUSES = new Set(["killed", "cleanup", "terminated", "done"]);
const RESTORABLE_ACTIVITIES = new Set(["exited"]);
/** Statuses that must never be restored (e.g. already merged) */
const NON_RESTORABLE_STATUSES = new Set(["merged"]);
import { SessionNotRestorableError, WorkspaceMissingError } from "@composio/ao-core";
/** POST /api/sessions/:id/restore — Restore a terminated session */
export async function POST(_request: NextRequest, { params }: { params: Promise<{ id: string }> }) {
@ -20,36 +14,20 @@ export async function POST(_request: NextRequest, { params }: { params: Promise<
try {
const { sessionManager } = await getServices();
const session = await sessionManager.get(id);
if (!session) {
return NextResponse.json({ error: "Session not found" }, { status: 404 });
}
if (NON_RESTORABLE_STATUSES.has(session.status)) {
return NextResponse.json({ error: "Cannot restore a merged session" }, { status: 409 });
}
const isTerminal =
RESTORABLE_STATUSES.has(session.status) ||
(session.activity !== null && RESTORABLE_ACTIVITIES.has(session.activity));
if (!isTerminal) {
return NextResponse.json({ error: "Session is not in a terminal state" }, { status: 409 });
}
// Re-spawn with the same project and issue to create a fresh session
const newSession = await sessionManager.spawn({
projectId: session.projectId,
issueId: session.issueId ?? undefined,
branch: session.branch ?? undefined,
});
const restored = await sessionManager.restore(id);
return NextResponse.json({
ok: true,
sessionId: id,
newSession: sessionToDashboard(newSession),
session: sessionToDashboard(restored),
});
} catch (err) {
if (err instanceof SessionNotRestorableError) {
return NextResponse.json({ error: err.message }, { status: 409 });
}
if (err instanceof WorkspaceMissingError) {
return NextResponse.json({ error: err.message }, { status: 422 });
}
const msg = err instanceof Error ? err.message : "Failed to restore session";
const status = msg.includes("not found") ? 404 : 500;
return NextResponse.json({ error: msg }, { status });

View File

@ -1,7 +1,13 @@
"use client";
import { useState, useEffect, useRef } from "react";
import { type DashboardSession, type AttentionLevel, getAttentionLevel } from "@/lib/types";
import {
type DashboardSession,
type AttentionLevel,
getAttentionLevel,
TERMINAL_STATUSES,
TERMINAL_ACTIVITIES,
} from "@/lib/types";
import { CI_STATUS } from "@composio/ao-core/types";
import { cn } from "@/lib/cn";
import { PRStatus } from "./PRStatus";
@ -56,12 +62,11 @@ export function SessionCard({ session, onSend, onKill, onMerge, onRestore }: Ses
const alerts = getAlerts(session);
const isReadyToMerge = pr?.mergeability.mergeable && pr.state === "open";
const isTerminal =
session.status === "killed" ||
session.status === "cleanup" ||
session.status === "terminated" ||
session.status === "done" ||
session.status === "merged" ||
session.activity === "exited";
TERMINAL_STATUSES.has(session.status) ||
(session.activity !== null && TERMINAL_ACTIVITIES.has(session.activity));
// UI restore button: more permissive than core isRestorable() — shows restore
// when agent has exited even if status is "working" (agent crashed mid-work).
// Only block "merged" (server-side restore will reject anyway).
const isRestorable = isTerminal && session.status !== "merged";
return (

View File

@ -3,7 +3,18 @@
*/
import { describe, it, expect } from "vitest";
import { getAttentionLevel, type DashboardSession } from "../types";
import {
getAttentionLevel,
TERMINAL_STATUSES,
TERMINAL_ACTIVITIES,
NON_RESTORABLE_STATUSES,
type DashboardSession,
} from "../types";
import {
TERMINAL_STATUSES as CORE_TERMINAL_STATUSES,
TERMINAL_ACTIVITIES as CORE_TERMINAL_ACTIVITIES,
NON_RESTORABLE_STATUSES as CORE_NON_RESTORABLE_STATUSES,
} from "@composio/ao-core/types";
// Helper to create a minimal DashboardSession for testing
function createSession(overrides?: Partial<DashboardSession>): DashboardSession {
@ -459,3 +470,17 @@ describe("getAttentionLevel", () => {
});
});
});
describe("constants sync with core", () => {
it("TERMINAL_STATUSES matches core", () => {
expect(TERMINAL_STATUSES).toBe(CORE_TERMINAL_STATUSES);
});
it("TERMINAL_ACTIVITIES matches core", () => {
expect(TERMINAL_ACTIVITIES).toBe(CORE_TERMINAL_ACTIVITIES);
});
it("NON_RESTORABLE_STATUSES matches core", () => {
expect(NON_RESTORABLE_STATUSES).toBe(CORE_NON_RESTORABLE_STATUSES);
});
});

View File

@ -21,6 +21,9 @@ import {
ACTIVITY_STATE,
SESSION_STATUS,
CI_STATUS,
TERMINAL_STATUSES,
TERMINAL_ACTIVITIES,
NON_RESTORABLE_STATUSES,
type CICheck as CoreCICheck,
type MergeReadiness,
type CIStatus,
@ -29,6 +32,9 @@ import {
type ReviewDecision,
} from "@composio/ao-core/types";
// Re-export for use in client components
export { TERMINAL_STATUSES, TERMINAL_ACTIVITIES, NON_RESTORABLE_STATUSES };
/**
* Attention zone priority level, ordered by human action urgency:
*