diff --git a/packages/cli/src/commands/session.ts b/packages/cli/src/commands/session.ts index 1d872a214..c93b7bef6 100644 --- a/packages/cli/src/commands/session.ts +++ b/packages/cli/src/commands/session.ts @@ -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 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); + } + }); } diff --git a/packages/core/src/__tests__/lifecycle-manager.test.ts b/packages/core/src/__tests__/lifecycle-manager.test.ts index daecf853b..eb6a9efd1 100644 --- a/packages/core/src/__tests__/lifecycle-manager.test.ts +++ b/packages/core/src/__tests__/lifecycle-manager.test.ts @@ -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), diff --git a/packages/core/src/__tests__/session-manager.test.ts b/packages/core/src/__tests__/session-manager.test.ts index 801b54af5..0e1376490 100644 --- a/packages/core/src/__tests__/session-manager.test.ts +++ b/packages/core/src/__tests__/session-manager.test.ts @@ -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).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).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"); diff --git a/packages/core/src/session-manager.ts b/packages/core/src/session-manager.ts index 599787d5c..47c9b80ab 100644 --- a/packages/core/src/session-manager.ts +++ b/packages/core/src/session-manager.ts @@ -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 { + async function cleanup( + projectId?: string, + options?: { dryRun?: boolean }, + ): Promise { const result: CleanupResult = { killed: [], skipped: [], errors: [] }; const sessions = await list(projectId); @@ -869,5 +880,152 @@ 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 { + // 1. Find session metadata across all projects + let raw: Record | 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 and enrich with live runtime state. + // metadataToSession sets activity: null, so without enrichment a crashed + // session (status "working", agent exited) would not be detected as terminal + // and isRestorable would reject it. + const session = metadataToSession(sessionId, raw); + const plugins = resolvePlugins(project); + await enrichSessionWithRuntimeState(session, plugins, true); + + // 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. Validate required plugins (plugins already resolved above for enrichment) + 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) { + throw new WorkspaceMissingError(workspacePath, "workspace plugin does not support restore"); + } + if (!session.branch) { + throw new WorkspaceMissingError(workspacePath, "branch metadata is missing"); + } + 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)}`, + ); + } + } + + // 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 }; } diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 6de8565a4..25bcbe221 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -83,6 +83,41 @@ export const SESSION_STATUS = { TERMINATED: "terminated" as const, } satisfies Record; +/** Statuses that indicate the session is in a terminal (dead) state. */ +export const TERMINAL_STATUSES: ReadonlySet = new Set([ + "killed", + "terminated", + "done", + "cleanup", + "errored", + "merged", +]); + +/** Activity states that indicate the session is no longer running. */ +export const TERMINAL_ACTIVITIES: ReadonlySet = new Set(["exited"]); + +/** Statuses that must never be restored (e.g. already merged). */ +export const NON_RESTORABLE_STATUSES: ReadonlySet = new Set(["merged"]); + +/** 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 +156,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; } @@ -243,6 +281,12 @@ export interface Agent { /** Extract information from agent's internal data (summary, cost, session ID) */ getSessionInfo(session: Session): Promise; + /** + * 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; + /** Optional: run setup after agent is launched (e.g. configure MCP servers) */ postLaunchSetup?(session: Session): Promise; @@ -339,6 +383,12 @@ export interface Workspace { /** Optional: run hooks after workspace creation (symlinks, installs, etc.) */ postCreate?(info: WorkspaceInfo, project: ProjectConfig): Promise; + + /** Optional: check if a workspace exists and is a valid git repo */ + exists?(workspacePath: string): Promise; + + /** Optional: restore a workspace (e.g. recreate a worktree for an existing branch) */ + restore?(config: WorkspaceCreateConfig, workspacePath: string): Promise; } export interface WorkspaceCreateConfig { @@ -904,6 +954,7 @@ export interface SessionMetadata { project?: string; createdAt?: string; runtimeHandle?: string; + restoredAt?: string; dashboardPort?: number; terminalWsPort?: number; directTerminalWsPort?: number; @@ -917,6 +968,7 @@ export interface SessionMetadata { export interface SessionManager { spawn(config: SessionSpawnConfig): Promise; spawnOrchestrator(config: OrchestratorSpawnConfig): Promise; + restore(sessionId: SessionId): Promise; list(projectId?: string): Promise; get(sessionId: SessionId): Promise; kill(sessionId: SessionId): Promise; @@ -996,3 +1048,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"; + } +} diff --git a/packages/plugins/agent-claude-code/src/index.ts b/packages/plugins/agent-claude-code/src/index.ts index 7be873dd7..21f1109c8 100644 --- a/packages/plugins/agent-claude-code/src/index.ts +++ b/packages/plugins/agent-claude-code/src/index.ts @@ -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 { + 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 { // Use absolute path for hook command (specific to this workspace) const hookScriptPath = join(workspacePath, ".claude", "metadata-updater.sh"); diff --git a/packages/plugins/workspace-clone/src/index.ts b/packages/plugins/workspace-clone/src/index.ts index 7d97289fd..d6bc7f924 100644 --- a/packages/plugins/workspace-clone/src/index.ts +++ b/packages/plugins/workspace-clone/src/index.ts @@ -166,6 +166,64 @@ export function create(config?: Record): Workspace { return infos; }, + async exists(workspacePath: string): Promise { + 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 { + 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 { // Run postCreate hooks // NOTE: commands run with full shell privileges — they come from trusted YAML config diff --git a/packages/plugins/workspace-worktree/src/index.ts b/packages/plugins/workspace-worktree/src/index.ts index 92f80fafb..11b9b49e1 100644 --- a/packages/plugins/workspace-worktree/src/index.ts +++ b/packages/plugins/workspace-worktree/src/index.ts @@ -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): Workspace { return infos; }, + async exists(workspacePath: string): Promise { + 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 { + 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 { const repoPath = expandPath(project.path); diff --git a/packages/web/src/__tests__/api-routes.test.ts b/packages/web/src/__tests__/api-routes.test.ts index d306cad99..baea1a098 100644 --- a/packages/web/src/__tests__/api-routes.test.ts +++ b/packages/web/src/__tests__/api-routes.test.ts @@ -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": { diff --git a/packages/web/src/app/api/sessions/[id]/restore/route.ts b/packages/web/src/app/api/sessions/[id]/restore/route.ts index 4e5ac683f..07a034f64 100644 --- a/packages/web/src/app/api/sessions/[id]/restore/route.ts +++ b/packages/web/src/app/api/sessions/[id]/restore/route.ts @@ -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 }); diff --git a/packages/web/src/components/SessionCard.tsx b/packages/web/src/components/SessionCard.tsx index 93d4d9b8d..e5a2c77bc 100644 --- a/packages/web/src/components/SessionCard.tsx +++ b/packages/web/src/components/SessionCard.tsx @@ -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 ( diff --git a/packages/web/src/lib/__tests__/types.test.ts b/packages/web/src/lib/__tests__/types.test.ts index f96862b73..42bdcfd3a 100644 --- a/packages/web/src/lib/__tests__/types.test.ts +++ b/packages/web/src/lib/__tests__/types.test.ts @@ -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 { @@ -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); + }); +}); diff --git a/packages/web/src/lib/types.ts b/packages/web/src/lib/types.ts index e0f5efe58..76f6e58cb 100644 --- a/packages/web/src/lib/types.ts +++ b/packages/web/src/lib/types.ts @@ -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: *