diff --git a/.claude/skills/env-reference/SKILL.md b/.claude/skills/env-reference/SKILL.md index df3197bfa..06c75a369 100644 --- a/.claude/skills/env-reference/SKILL.md +++ b/.claude/skills/env-reference/SKILL.md @@ -66,7 +66,13 @@ See `apps/api/.env.example` for the full list. Key variables: - `CF_CONTAINER_SLEEP_AFTER` — Container idle sleep duration for instant-session runtime (default: `10m`) - `CF_CONTAINER_VM_AGENT_PORT` — vm-agent standalone HTTP port inside the raw container (default: `8080`) - `CF_CONTAINER_PORT_READY_TIMEOUT_MS` — Max wait for vm-agent port readiness (default: `30000`) -- `CF_CONTAINER_WORKSPACE_BASE_DIR` — Base checkout directory inside raw containers (default: `/workspaces`) +- `$1 +- `SESSION_SNAPSHOT_TTL_DAYS` — Retention for hibernated session snapshots; deployment also provisions matching R2 prefix expiration (default: `7`) +- `SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES` — Maximum bytes accepted for one snapshot artifact (default: `104857600`) +- `SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES` — Per-file threshold before snapshot content is visibly skipped (default: `52428800`) +- `SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS` — Progress-idle timeout for snapshot upload/download (default: `30000`) +- `SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES` — Maximum snapshot coordination JSON body (default: `262144`) +- `SESSION_SNAPSHOT_R2_PREFIX` — Private R2 object prefix for session snapshots (default: `session-snapshots`) ### Devcontainer Cache diff --git a/.github/workflows/deploy-reusable.yml b/.github/workflows/deploy-reusable.yml index ff4240435..090a3f4a7 100644 --- a/.github/workflows/deploy-reusable.yml +++ b/.github/workflows/deploy-reusable.yml @@ -358,7 +358,9 @@ jobs: HETZNER_BASE_IMAGE: ${{ vars.HETZNER_BASE_IMAGE }} ARTIFACTS_BINDING_ENABLED: ${{ vars.ARTIFACTS_BINDING_ENABLED }} CF_CONTAINER_ENABLED: ${{ vars.CF_CONTAINER_ENABLED }} + CF_CONTAINER_SLEEP_AFTER: ${{ vars.CF_CONTAINER_SLEEP_AFTER }} CF_CONTAINER_PORT_READY_TIMEOUT_MS: ${{ vars.CF_CONTAINER_PORT_READY_TIMEOUT_MS }} + CF_CONTAINER_WAKE_TIMEOUT_MS: ${{ vars.CF_CONTAINER_WAKE_TIMEOUT_MS }} CF_CONTAINER_VM_AGENT_PORT: ${{ vars.CF_CONTAINER_VM_AGENT_PORT }} SANDBOX_ENABLED: ${{ vars.SANDBOX_ENABLED }} SANDBOX_EXEC_TIMEOUT_MS: ${{ vars.SANDBOX_EXEC_TIMEOUT_MS }} diff --git a/apps/api/.env.example b/apps/api/.env.example index 9a56fb5a1..7fb68a852 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -52,6 +52,12 @@ BASE_DOMAIN=workspaces.example.com # CF_CONTAINER_VM_AGENT_PORT=8080 # CF_CONTAINER_PORT_READY_TIMEOUT_MS=30000 # CF_CONTAINER_WORKSPACE_BASE_DIR=/workspaces +# SESSION_SNAPSHOT_TTL_DAYS=7 +# SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES=104857600 +# SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES=52428800 +# SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS=30000 +# SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES=262144 +# SESSION_SNAPSHOT_R2_PREFIX=session-snapshots # User approval / invite-only mode # When "true", new users require admin approval before accessing the platform. diff --git a/apps/api/src/db/migrations/0091_session_snapshots.sql b/apps/api/src/db/migrations/0091_session_snapshots.sql new file mode 100644 index 000000000..46e25f730 --- /dev/null +++ b/apps/api/src/db/migrations/0091_session_snapshots.sql @@ -0,0 +1,33 @@ +-- Runtime-neutral hibernate/restore snapshots for agent sessions. +CREATE TABLE session_snapshots ( + id TEXT PRIMARY KEY, + project_id TEXT REFERENCES projects(id) ON DELETE SET NULL, + workspace_id TEXT REFERENCES workspaces(id) ON DELETE SET NULL, + node_id TEXT REFERENCES nodes(id) ON DELETE SET NULL, + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + chat_session_id TEXT NOT NULL, + agent_session_id TEXT, + runtime TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + degradation TEXT NOT NULL DEFAULT 'none', + home_r2_key TEXT, + wip_r2_key TEXT, + manifest_r2_key TEXT NOT NULL, + base_commit TEXT, + expires_at TEXT NOT NULL, + manifest_json TEXT, + restore_status TEXT, + restore_message TEXT, + restored_at TEXT, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE UNIQUE INDEX idx_session_snapshots_chat_session_id + ON session_snapshots(chat_session_id); + +CREATE INDEX idx_session_snapshots_workspace_id + ON session_snapshots(workspace_id); + +CREATE INDEX idx_session_snapshots_expires_at + ON session_snapshots(expires_at); diff --git a/apps/api/src/db/schema.ts b/apps/api/src/db/schema.ts index 963564a5e..4d62949e2 100644 --- a/apps/api/src/db/schema.ts +++ b/apps/api/src/db/schema.ts @@ -1044,6 +1044,47 @@ export const agentSessions = sqliteTable( }) ); +// ============================================================================= +// Runtime-neutral Session Snapshots +// ============================================================================= +export const sessionSnapshots = sqliteTable( + 'session_snapshots', + { + id: text('id').primaryKey(), + projectId: text('project_id').references(() => projects.id, { onDelete: 'set null' }), + workspaceId: text('workspace_id').references(() => workspaces.id, { onDelete: 'set null' }), + nodeId: text('node_id').references(() => nodes.id, { onDelete: 'set null' }), + userId: text('user_id') + .notNull() + .references(() => users.id, { onDelete: 'cascade' }), + chatSessionId: text('chat_session_id').notNull(), + agentSessionId: text('agent_session_id'), + runtime: text('runtime').notNull(), + status: text('status').notNull().default('pending'), + degradation: text('degradation').notNull().default('none'), + homeR2Key: text('home_r2_key'), + wipR2Key: text('wip_r2_key'), + manifestR2Key: text('manifest_r2_key').notNull(), + baseCommit: text('base_commit'), + expiresAt: text('expires_at').notNull(), + manifestJson: text('manifest_json'), + restoreStatus: text('restore_status'), + restoreMessage: text('restore_message'), + restoredAt: text('restored_at'), + createdAt: text('created_at') + .notNull() + .default(sql`CURRENT_TIMESTAMP`), + updatedAt: text('updated_at') + .notNull() + .default(sql`CURRENT_TIMESTAMP`), + }, + (table) => ({ + chatSessionIdUnique: uniqueIndex('idx_session_snapshots_chat_session_id').on(table.chatSessionId), + workspaceIdIdx: index('idx_session_snapshots_workspace_id').on(table.workspaceId), + expiresAtIdx: index('idx_session_snapshots_expires_at').on(table.expiresAt), + }) +); + // ============================================================================= // Agent Settings (per-user, per-agent configuration) // ============================================================================= @@ -1545,6 +1586,8 @@ export type Workspace = typeof workspaces.$inferSelect; export type NewWorkspace = typeof workspaces.$inferInsert; export type AgentSession = typeof agentSessions.$inferSelect; export type NewAgentSession = typeof agentSessions.$inferInsert; +export type SessionSnapshot = typeof sessionSnapshots.$inferSelect; +export type NewSessionSnapshot = typeof sessionSnapshots.$inferInsert; export type UIStandard = typeof uiStandards.$inferSelect; export type NewUIStandard = typeof uiStandards.$inferInsert; export type ThemeToken = typeof themeTokens.$inferSelect; diff --git a/apps/api/src/durable-objects/vm-agent-container.ts b/apps/api/src/durable-objects/vm-agent-container.ts index f9290484c..f80c6ca32 100644 --- a/apps/api/src/durable-objects/vm-agent-container.ts +++ b/apps/api/src/durable-objects/vm-agent-container.ts @@ -1,10 +1,11 @@ import { Container, switchPort } from '@cloudflare/containers'; -import { eq } from 'drizzle-orm'; +import { desc, eq } from 'drizzle-orm'; import { drizzle } from 'drizzle-orm/d1'; import * as schema from '../db/schema'; import type { Env } from '../env'; import { log } from '../lib/logger'; +import { signCallbackToken, signNodeCallbackToken, signNodeManagementToken } from '../services/jwt'; export const DEFAULT_CF_CONTAINER_SLEEP_AFTER = '1h'; export const DEFAULT_CF_CONTAINER_ACTIVE_WORK_MAX_MS = 2 * 60 * 60 * 1000; @@ -45,7 +46,7 @@ interface ActiveWorkState { const ACTIVE_WORK_KEY = 'activeWork'; const KEEPALIVE_CALLBACK = 'renewActiveWorkKeepalive'; -const SLEEPING_RESPONSE = 'Container is asleep; wake/rehydrate is not implemented yet.'; +const WAKE_DEGRADED_RESPONSE = 'Workspace woke with degraded snapshot restore; retry the prompt or fork from transcript history.'; export class VmAgentContainer extends Container { defaultPort = 8080; @@ -53,6 +54,12 @@ export class VmAgentContainer extends Container { sleepAfter = DEFAULT_CF_CONTAINER_SLEEP_AFTER; enableInternet = true; + // Serializes wake-from-snapshot so two concurrent requests to a sleeping + // container do not both launch a fresh container + restore. DO request + // handlers interleave across `await`, so the sleeping-check + wake must run + // as one critical section (see .claude/rules/45). + private wakeChain: Promise = Promise.resolve(); + constructor(ctx: DurableObjectState>, env: Env) { super(ctx, env); const configuredPort = Number.parseInt(env.CF_CONTAINER_VM_AGENT_PORT || env.SANDBOX_VM_AGENT_PORT || '', 10); @@ -108,12 +115,18 @@ export class VmAgentContainer extends Container { } async proxyHttp(request: Request, port?: number): Promise { - const state = await this.getState(); + let state = await this.getState(); const lifecycleStatus = await this.ctx.storage.get('lifecycleStatus'); if (lifecycleStatus === 'sleeping') { - // Phase 3 of idea 01KX4KSXEXQMP41KS34TW9EN01 will replace this temporary - // response with wake/rehydrate before proxying the request. - return new Response(SLEEPING_RESPONSE, { status: 503 }); + const wake = await this.ensureAwake(); + if (!wake.ok) { + return new Response(wake.message || WAKE_DEGRADED_RESPONSE, { status: 503 }); + } + // wakeFromSnapshot launched a fresh container and restored the session. + // Re-read the container state so the stopped-check below reflects the + // now-running container instead of the pre-wake stopped snapshot, + // otherwise a successfully-woken session is wrongly rejected with 410. + state = await this.getState(); } if (state.status === 'stopped' || state.status === 'stopped_with_code') { return new Response('Container is stopped; create a new instant session.', { status: 410 }); @@ -232,12 +245,16 @@ export class VmAgentContainer extends Container { } override async fetch(request: Request): Promise { - const state = await this.getState(); + let state = await this.getState(); const lifecycleStatus = await this.ctx.storage.get('lifecycleStatus'); if (lifecycleStatus === 'sleeping') { - // Phase 3 of idea 01KX4KSXEXQMP41KS34TW9EN01 will replace this temporary - // response with wake/rehydrate before proxying the request. - return new Response(SLEEPING_RESPONSE, { status: 503 }); + const wake = await this.ensureAwake(); + if (!wake.ok) { + return new Response(wake.message || WAKE_DEGRADED_RESPONSE, { status: 503 }); + } + // Re-read the container state after wake so the stopped-check reflects + // the freshly-launched container, not the pre-wake stopped snapshot. + state = await this.getState(); } if (state.status === 'stopped' || state.status === 'stopped_with_code') { return new Response('Container is stopped; create a new instant session.', { status: 410 }); @@ -299,6 +316,165 @@ export class VmAgentContainer extends Container { return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS; } + /** + * Wake a sleeping container exactly once under concurrency. Serializes on + * `wakeChain` and re-reads `lifecycleStatus` inside the critical section, so a + * second request that arrives while the first is waking sees `running` and + * skips a duplicate launch/restore instead of racing it (see rule 45). + */ + private async ensureAwake(): Promise<{ ok: boolean; message?: string }> { + const run = this.wakeChain.then(async () => { + const status = await this.ctx.storage.get('lifecycleStatus'); + if (status !== 'sleeping') { + return { ok: true }; + } + return this.wakeFromSnapshot(); + }); + // Keep the chain alive even if this wake rejects, so a failure does not + // permanently wedge all future wakes. + this.wakeChain = run.then( + () => undefined, + () => undefined + ); + return run; + } + + private async wakeFromSnapshot(): Promise<{ ok: boolean; message?: string }> { + const config = await this.ctx.storage.get('launchConfig'); + if (!config) { + return { ok: false, message: 'Container launch configuration is unavailable.' }; + } + const db = drizzle(this.env.DATABASE, { schema }); + const workspace = await db + .select({ + userId: schema.workspaces.userId, + chatSessionId: schema.workspaces.chatSessionId, + }) + .from(schema.workspaces) + .where(eq(schema.workspaces.id, config.workspaceId)) + .get(); + if (!workspace?.chatSessionId) { + return { ok: false, message: 'Workspace session metadata is unavailable.' }; + } + const agentSession = await db + .select({ id: schema.agentSessions.id, agentType: schema.agentSessions.agentType }) + .from(schema.agentSessions) + .where(eq(schema.agentSessions.workspaceId, config.workspaceId)) + .orderBy(desc(schema.agentSessions.updatedAt)) + .get(); + if (!agentSession) { + return { ok: false, message: 'Agent session metadata is unavailable.' }; + } + + log.info('vm_agent_container_wake_started', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + chatSessionId: workspace.chatSessionId, + agentSessionId: agentSession.id, + }); + + await this.ctx.storage.put('lifecycleStatus', 'launching' satisfies LifecycleStatus); + // The container's CALLBACK_TOKEN must be node-scoped to match the initial + // launch (see launchInstantSession): the vm-agent uses it for node callbacks + // (error/activity/message reporting) which reject workspace-scoped tokens. + // Using a workspace-scoped token here caused restored sessions to accept a + // prompt (200) but silently fail to report the agent's reply back (403 + // "Insufficient token scope"), so no answer appeared after wake. + const callbackToken = await signNodeCallbackToken(config.nodeId, this.env); + await this.launch(config, { nodeCallbackToken: callbackToken }); + + // The fresh container never ran create-workspace, so its workspace-scoped + // runtime.CallbackToken is unset. The message reporter and snapshot + // callbacks require it (they do NOT fall back to the node-scoped token), so + // pass it on the restore request; without it, restored sessions accept a + // prompt but silently discard the agent's reply ("no auth token"). + const workspaceCallbackToken = await signCallbackToken(config.workspaceId, this.env); + const { token } = await signNodeManagementToken(workspace.userId, config.nodeId, config.workspaceId, this.env); + const restoreUrl = new URL(`http://localhost:${config.vmAgentPort}/workspaces/${config.workspaceId}/agent-sessions/${agentSession.id}/restore`); + const restoreResponse = await this.containerFetch( + new Request(restoreUrl.toString(), { + method: 'POST', + headers: { + Authorization: `Bearer ${token}`, + 'Content-Type': 'application/json', + 'X-SAM-Node-Id': config.nodeId, + 'X-SAM-Workspace-Id': config.workspaceId, + }, + body: JSON.stringify({ + chatSessionId: workspace.chatSessionId, + runtime: 'cf-container', + agentType: agentSession.agentType, + workspaceCallbackToken, + }), + }), + config.vmAgentPort + ); + const restoreBody = await restoreResponse.text().catch(() => ''); + if (!restoreResponse.ok) { + await this.markWakeDegraded(config, restoreBody || `restore failed with HTTP ${restoreResponse.status}`); + return { ok: false, message: restoreBody || 'Session restore failed.' }; + } + let restoreStatus = ''; + try { + const parsed = JSON.parse(restoreBody) as { status?: unknown }; + restoreStatus = typeof parsed.status === 'string' ? parsed.status : ''; + } catch { + // A successful restore must provide an explicit machine-readable status. + } + if (restoreStatus !== 'restored') { + const message = restoreBody || 'Session restore did not report restored status.'; + await this.markWakeDegraded(config, message); + return { ok: false, message }; + } + + const now = new Date().toISOString(); + await db.update(schema.nodes).set({ + status: 'running', + healthStatus: 'healthy', + errorMessage: null, + updatedAt: now, + }).where(eq(schema.nodes.id, config.nodeId)); + await db.update(schema.workspaces).set({ + status: 'running', + errorMessage: null, + updatedAt: now, + }).where(eq(schema.workspaces.id, config.workspaceId)); + await db.update(schema.agentSessions).set({ + status: 'running', + errorMessage: null, + updatedAt: now, + }).where(eq(schema.agentSessions.id, agentSession.id)); + await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus); + + log.info('vm_agent_container_wake_completed', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + chatSessionId: workspace.chatSessionId, + agentSessionId: agentSession.id, + }); + return { ok: true }; + } + + private async markWakeDegraded(config: VmAgentContainerLaunchConfig, message: string): Promise { + const now = new Date().toISOString(); + const db = drizzle(this.env.DATABASE, { schema }); + await db.update(schema.workspaces).set({ + status: 'recovery', + errorMessage: message, + updatedAt: now, + }).where(eq(schema.workspaces.id, config.workspaceId)); + await db.update(schema.agentSessions).set({ + status: 'error', + errorMessage: message, + updatedAt: now, + }).where(eq(schema.agentSessions.workspaceId, config.workspaceId)); + log.warn('vm_agent_container_wake_degraded', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + message, + }); + } + private async replaceKeepaliveSchedule(delayMs: number): Promise { await this.clearKeepaliveSchedule(); await this.schedule(Math.max(1, Math.ceil(delayMs / 1000)), KEEPALIVE_CALLBACK); diff --git a/apps/api/src/env.ts b/apps/api/src/env.ts index 3e6d4c4aa..76e03a9a7 100644 --- a/apps/api/src/env.ts +++ b/apps/api/src/env.ts @@ -140,6 +140,12 @@ export interface Env { COMPOSE_IMAGE_ARTIFACT_CLEANUP_BATCH_SIZE?: string; // Max abandoned compose artifacts to delete per cleanup run (default: 50) COMPOSE_IMAGE_ARTIFACT_CLEANUP_INTERVAL_HOURS?: string; // Minimum hours between R2 cleanup scans from cron (default: 24) COMPOSE_IMAGE_ARTIFACT_CLEANUP_LAST_RUN_KV_KEY?: string; // KV key for cleanup interval gating + SESSION_SNAPSHOT_TTL_DAYS?: string; // Runtime hibernate snapshot retention (default: 7) + SESSION_SNAPSHOT_R2_PREFIX?: string; // R2 key prefix for runtime hibernate snapshots + SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES?: string; // Max combined home/WIP snapshot size (default: 104857600) + SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES?: string; // Max individual file/dir included by vm-agent scanner (default: 52428800) + SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS?: string; // No-progress upload/download watchdog window (default: 30000) + SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES?: string; // Max snapshot control-plane JSON request size (default: 262144) DEPLOY_ACME_EMAIL?: string; // Contact email for deployment-node ACME certificates DEPLOY_ACME_CA?: string; // ACME CA directory override for deployment nodes DEPLOY_COMPOSE_CMD?: string; // Docker Compose command override on deployment nodes @@ -772,6 +778,7 @@ export interface Env { CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS?: string; // Active-work renewActivityTimeout interval (default: 300000) CF_CONTAINER_VM_AGENT_PORT?: string; // vm-agent standalone HTTP port inside the raw container (default: 8080) CF_CONTAINER_PORT_READY_TIMEOUT_MS?: string; // Max time to wait for vm-agent port readiness (default: 30000) + CF_CONTAINER_WAKE_TIMEOUT_MS?: string; // Max time for launch + restore before forwarding a wake request (default: 120000) CF_CONTAINER_WORKSPACE_BASE_DIR?: string; // Base checkout dir inside raw container (default: /workspaces) // Legacy Sandbox SDK prototype (admin-only) SANDBOX_ENABLED?: string; // Legacy/fallback kill switch for sandbox routes and older staging config (default: false) diff --git a/apps/api/src/routes/chat-workspace-resolver.ts b/apps/api/src/routes/chat-workspace-resolver.ts index d4804d464..928723f52 100644 --- a/apps/api/src/routes/chat-workspace-resolver.ts +++ b/apps/api/src/routes/chat-workspace-resolver.ts @@ -76,7 +76,7 @@ export async function resolveLiveWorkspaceForSession( } /** - * Resolve the live workspace AND its running agent session for a chat action + * Resolve the live workspace AND its resumable agent session for a chat action * (e.g. /prompt, /cancel), enforcing tenant scoping at every layer. * * Consolidates the shared resolution path used by the /prompt and /cancel @@ -91,7 +91,7 @@ export async function resolveLiveWorkspaceForSession( export async function resolveLiveAgentSessionForChat( db: ChatDb, { projectId, sessionId, userId }: { projectId: string; sessionId: string; userId: string } -): Promise<{ workspace: { id: string; nodeId: string }; agentSession: { id: string } }> { +): Promise<{ workspace: { id: string; nodeId: string; nodeStatus: string }; agentSession: { id: string } }> { // Find the workspace linked to this chat session, scoped to the owning // project + user (see resolveLiveWorkspaceForSession). The node join also // verifies the node is still active: when a node is destroyed (e.g., after @@ -107,13 +107,11 @@ export async function resolveLiveAgentSessionForChat( // Verify the node is still reachable — prevents requests to destroyed VMs // whose DNS records no longer exist (would loop back via wildcard DNS). // D1 nodes.status uses 'running' for healthy nodes (not 'active'/'warm', which are DO states). - if (workspace.nodeStatus === 'sleeping') { - // Phase 3 of idea 01KX4KSXEXQMP41KS34TW9EN01 will wake and rehydrate the - // cf-container before forwarding this prompt. - throw errors.conflict('The workspace container is asleep. Send a new message after wake/rehydrate support lands.'); - } - - if (workspace.nodeStatus !== 'running') { + // A request routed to a sleeping cf-container is the wake signal: the + // container proxy launches a fresh runtime and restores its latest snapshot + // before forwarding the request. Preserve the fail-fast guard for every + // other non-running node state. + if (workspace.nodeStatus !== 'running' && workspace.nodeStatus !== 'sleeping') { throw errors.conflict( 'The workspace node is no longer running. Start a new chat to create a fresh workspace.' ); @@ -128,7 +126,7 @@ export async function resolveLiveAgentSessionForChat( and( eq(schema.agentSessions.workspaceId, workspace.id), eq(schema.agentSessions.userId, userId), - eq(schema.agentSessions.status, 'running') + inArray(schema.agentSessions.status, ['running', 'sleeping']) ) ) .limit(1); @@ -137,5 +135,5 @@ export async function resolveLiveAgentSessionForChat( throw errors.notFound('No running agent session found'); } - return { workspace: { id: workspace.id, nodeId: workspace.nodeId }, agentSession }; + return { workspace: { id: workspace.id, nodeId: workspace.nodeId, nodeStatus: workspace.nodeStatus }, agentSession }; } diff --git a/apps/api/src/routes/chat.ts b/apps/api/src/routes/chat.ts index a941c0ccd..7ec9af872 100644 --- a/apps/api/src/routes/chat.ts +++ b/apps/api/src/routes/chat.ts @@ -517,14 +517,18 @@ chatRoutes.post('/:sessionId/prompt', async (c) => { const { enrichedMessage } = await enrichMessageWithMentions(content, db, projectId, userId, c.env); // Forward the prompt to the VM agent - const { sendPromptToAgentOnNode } = await import('../services/node-agent'); + const { getCfContainerWakeTimeoutMs, sendPromptToAgentOnNode } = await import('../services/node-agent'); const result = await sendPromptToAgentOnNode( workspace.nodeId, workspace.id, agentSession.id, enrichedMessage, c.env, - userId + userId, + undefined, + workspace.nodeStatus === 'sleeping' + ? { requestTimeoutMs: getCfContainerWakeTimeoutMs(c.env) } + : undefined ); return c.json(expectJsonRecord(result, 'chat.agent_prompt_result')); diff --git a/apps/api/src/routes/deployment-environments.ts b/apps/api/src/routes/deployment-environments.ts index e73c22ff5..11a7c8945 100644 --- a/apps/api/src/routes/deployment-environments.ts +++ b/apps/api/src/routes/deployment-environments.ts @@ -37,7 +37,7 @@ import { getNodeLogsFromNode, getNodeSystemInfoFromNode, listNodeContainersFromNode, -} from '../services/node-agent'; +} from '../services/node-agent-diagnostics'; import { deleteNodeResources } from '../services/nodes'; import { registerDeploymentEnvironmentLifecycleRoutes } from './deployment-environment-lifecycle'; diff --git a/apps/api/src/routes/mcp/deployment-tools.ts b/apps/api/src/routes/mcp/deployment-tools.ts index 55d7ba2d4..2ee0c5564 100644 --- a/apps/api/src/routes/mcp/deployment-tools.ts +++ b/apps/api/src/routes/mcp/deployment-tools.ts @@ -27,7 +27,7 @@ import { type DeploymentRouteTargetOptions, } from '../../services/deployment-routing'; import { getRuntimeLimits } from '../../services/limits'; -import { getNodeLogsFromNode } from '../../services/node-agent'; +import { getNodeLogsFromNode } from '../../services/node-agent-diagnostics'; import { byteLength, PROJECT_ENV_KEY_PATTERN } from '../projects/_helpers'; import { INTERNAL_ERROR, diff --git a/apps/api/src/routes/nodes.ts b/apps/api/src/routes/nodes.ts index c1374f03b..93ceedd4b 100644 --- a/apps/api/src/routes/nodes.ts +++ b/apps/api/src/routes/nodes.ts @@ -23,13 +23,15 @@ import { getRuntimeLimits } from '../services/limits'; import { fetchNodeAgent, getNodeAgentRequestTimeoutMs, - getNodeLogsFromNode, - getNodeSystemInfoFromNode, - listNodeContainersFromNode, listNodeEventsOnNode, nodeAgentRawRequest, stopWorkspaceOnNode, } from '../services/node-agent'; +import { + getNodeLogsFromNode, + getNodeSystemInfoFromNode, + listNodeContainersFromNode, +} from '../services/node-agent-diagnostics'; import { createNodeRecord, deleteNodeResources, provisionNode, stopNodeResources } from '../services/nodes'; import { recordNodeRoutingMetric } from '../services/telemetry'; diff --git a/apps/api/src/routes/projects/agent-activity-callback.ts b/apps/api/src/routes/projects/agent-activity-callback.ts index 8c6bd2669..59e01d956 100644 --- a/apps/api/src/routes/projects/agent-activity-callback.ts +++ b/apps/api/src/routes/projects/agent-activity-callback.ts @@ -1,11 +1,15 @@ +import { eq } from 'drizzle-orm'; +import { drizzle } from 'drizzle-orm/d1'; import { Hono } from 'hono'; +import * as schema from '../../db/schema'; import type { Env } from '../../env'; import { extractBearerToken } from '../../lib/auth-helpers'; import { log } from '../../lib/logger'; import { errors } from '../../middleware/error'; import { AcpSessionActivityReportSchema, jsonValidator } from '../../schemas'; import { verifyCallbackToken } from '../../services/jwt'; +import { hibernateAgentSessionOnNode } from '../../services/node-agent'; import * as projectDataService from '../../services/project-data'; import { markVmAgentContainerActiveWorkEndedBestEffort } from '../../services/vm-agent-container'; @@ -70,6 +74,41 @@ agentActivityCallbackRoute.post( statusError: body.statusError, }); if (body.activity === 'idle' || body.activity === 'error') { + if (body.activity === 'idle' && existing.workspaceId && existing.nodeId && existing.acpSdkSessionId) { + const db = drizzle(c.env.DATABASE, { schema }); + const workspace = await db + .select({ + id: schema.workspaces.id, + userId: schema.workspaces.userId, + chatSessionId: schema.workspaces.chatSessionId, + runtime: schema.nodes.runtime, + }) + .from(schema.workspaces) + .leftJoin(schema.nodes, eq(schema.nodes.id, schema.workspaces.nodeId)) + .where(eq(schema.workspaces.id, existing.workspaceId)) + .get(); + if (workspace?.runtime === 'cf-container' && workspace.chatSessionId) { + await hibernateAgentSessionOnNode( + existing.nodeId, + existing.workspaceId, + existing.acpSdkSessionId, + c.env, + workspace.userId, + { + chatSessionId: workspace.chatSessionId, + runtime: 'cf-container', + } + ).catch((err) => { + log.warn('acp_activity.session_snapshot_failed', { + projectId, + sessionId, + workspaceId: existing.workspaceId, + nodeId: existing.nodeId, + error: err instanceof Error ? err.message : String(err), + }); + }); + } + } await markVmAgentContainerActiveWorkEndedBestEffort( c.env, existing.nodeId, diff --git a/apps/api/src/routes/workspaces/index.ts b/apps/api/src/routes/workspaces/index.ts index a1577ec1e..b4e0eca1f 100644 --- a/apps/api/src/routes/workspaces/index.ts +++ b/apps/api/src/routes/workspaces/index.ts @@ -6,6 +6,7 @@ import { crudRoutes } from './crud'; import { lifecycleRoutes } from './lifecycle'; import { localForwardRoutes } from './local-forward'; import { runtimeRoutes } from './runtime'; +import { sessionSnapshotRoutes } from './session-snapshots'; const workspacesRoutes = new Hono<{ Bindings: Env }>(); workspacesRoutes.route('/', crudRoutes); @@ -13,5 +14,6 @@ workspacesRoutes.route('/', localForwardRoutes); workspacesRoutes.route('/', lifecycleRoutes); workspacesRoutes.route('/', agentSessionRoutes); workspacesRoutes.route('/', runtimeRoutes); +workspacesRoutes.route('/', sessionSnapshotRoutes); export { workspacesRoutes }; diff --git a/apps/api/src/routes/workspaces/session-snapshots.ts b/apps/api/src/routes/workspaces/session-snapshots.ts new file mode 100644 index 000000000..ac4b72eae --- /dev/null +++ b/apps/api/src/routes/workspaces/session-snapshots.ts @@ -0,0 +1,285 @@ +import { eq } from 'drizzle-orm'; +import { drizzle } from 'drizzle-orm/d1'; +import type { Context } from 'hono'; +import { Hono } from 'hono'; + +import * as schema from '../../db/schema'; +import type { Env } from '../../env'; +import { expectJsonRecord } from '../../lib/runtime-validation'; +import { errors } from '../../middleware/error'; +import { + buildSessionSnapshotR2Key, + completeSessionSnapshot, + getRestorableSessionSnapshot, + getSessionSnapshotConfig, + prepareSessionSnapshot, + recordSessionSnapshotRestoreResult, + type SessionSnapshotArtifact, + type SessionSnapshotDegradation, + type SessionSnapshotManifest, + type SessionSnapshotStatus, +} from '../../services/session-snapshots'; +import { verifyWorkspaceCallbackAuth } from './_helpers'; + +const sessionSnapshotRoutes = new Hono<{ Bindings: Env }>(); +type SnapshotRouteContext = Context<{ Bindings: Env }>; + +const ARTIFACTS = new Set(['home', 'wip', 'manifest']); +const COMPLETE_STATUSES = new Set(['available', 'degraded', 'failed']); +const DEGRADATIONS = new Set([ + 'none', + 'home-skipped', + 'wip-only', + 'transcript-only', +]); + +async function readJsonBody(c: SnapshotRouteContext) { + const raw = await c.req.raw.text(); + if (new TextEncoder().encode(raw).byteLength > getSessionSnapshotConfig(c.env).jsonBodyMaxBytes) { + throw errors.badRequest('Snapshot request body is too large'); + } + try { + return expectJsonRecord(JSON.parse(raw), 'session snapshot request'); + } catch { + throw errors.badRequest('Invalid JSON request body'); + } +} + +function stringField(body: Record, key: string, required = true): string | null { + const value = body[key]; + if (typeof value !== 'string' || value.trim() === '') { + if (required) throw errors.badRequest(`${key} is required`); + return null; + } + return value.trim(); +} + +function requiredStringField(body: Record, key: string): string { + const value = stringField(body, key, true); + if (value === null) throw errors.badRequest(`${key} is required`); + return value; +} + +function artifactFromParam(value: string): SessionSnapshotArtifact { + if (!ARTIFACTS.has(value as SessionSnapshotArtifact)) { + throw errors.badRequest('Unknown snapshot artifact'); + } + return value as SessionSnapshotArtifact; +} + +async function requireWorkspace(c: SnapshotRouteContext, workspaceId: string) { + const db = drizzle(c.env.DATABASE, { schema }); + const rows = await db + .select() + .from(schema.workspaces) + .where(eq(schema.workspaces.id, workspaceId)) + .limit(1); + const workspace = rows[0]; + if (!workspace) throw errors.notFound('Workspace'); + return { db, workspace }; +} + +sessionSnapshotRoutes.post('/:id/session-snapshot/prepare', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const { db, workspace } = await requireWorkspace(c, workspaceId); + const body = await readJsonBody(c); + const chatSessionId = requiredStringField(body, 'chatSessionId'); + const agentSessionId = stringField(body, 'agentSessionId', false); + const runtime = stringField(body, 'runtime', false) || 'runtime-neutral'; + + if (!workspace.chatSessionId || workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + + const prepared = await prepareSessionSnapshot(db, c.env, { + workspaceId, + nodeId: workspace.nodeId, + projectId: workspace.projectId, + userId: workspace.userId, + chatSessionId, + agentSessionId, + runtime, + }); + + return c.json({ + snapshotId: prepared.snapshotId, + expiresAt: prepared.expiresAt, + config: prepared.config, + upload: { + home: `/api/workspaces/${workspaceId}/session-snapshot/artifacts/home?chatSessionId=${encodeURIComponent(chatSessionId)}`, + wip: `/api/workspaces/${workspaceId}/session-snapshot/artifacts/wip?chatSessionId=${encodeURIComponent(chatSessionId)}`, + }, + }); +}); + +sessionSnapshotRoutes.put('/:id/session-snapshot/artifacts/:artifact', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const artifact = artifactFromParam(c.req.param('artifact')); + if (artifact === 'manifest') { + throw errors.badRequest('Manifest is written by the complete endpoint'); + } + const chatSessionId = c.req.query('chatSessionId')?.trim(); + if (!chatSessionId) throw errors.badRequest('chatSessionId is required'); + const { workspace } = await requireWorkspace(c, workspaceId); + if (!workspace.chatSessionId || workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + const contentLength = c.req.header('content-length'); + if (!contentLength) throw errors.badRequest('Snapshot artifact Content-Length is required'); + { + const parsed = Number.parseInt(contentLength, 10); + if (!Number.isFinite(parsed) || parsed < 0) throw errors.badRequest('Invalid Content-Length'); + if (parsed > getSessionSnapshotConfig(c.env).totalBudgetBytes) { + throw errors.badRequest('Snapshot artifact exceeds configured total budget'); + } + } + if (!c.req.raw.body) throw errors.badRequest('Snapshot artifact body is required'); + + const key = buildSessionSnapshotR2Key(c.env, chatSessionId, artifact); + await c.env.R2.put(key, c.req.raw.body, { + httpMetadata: { + contentType: artifact === 'home' ? 'application/x-tar' : 'application/octet-stream', + }, + }); + return c.json({ key, artifact }); +}); + +sessionSnapshotRoutes.post('/:id/session-snapshot/complete', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const { db, workspace } = await requireWorkspace(c, workspaceId); + const body = await readJsonBody(c); + const chatSessionId = stringField(body, 'chatSessionId'); + if (!workspace.chatSessionId || workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + + const status = requiredStringField(body, 'status') as SessionSnapshotStatus; + const degradation = requiredStringField(body, 'degradation') as SessionSnapshotDegradation; + if (!COMPLETE_STATUSES.has(status)) throw errors.badRequest('Invalid snapshot status'); + if (!DEGRADATIONS.has(degradation)) throw errors.badRequest('Invalid snapshot degradation'); + const manifestRecord = expectJsonRecord(body.manifest, 'snapshot manifest'); + const manifest = manifestRecord as unknown as SessionSnapshotManifest; + if ( + manifest.version !== 1 || + manifest.chatSessionId !== chatSessionId || + manifest.workspaceId !== workspaceId + ) { + throw errors.badRequest('Snapshot manifest identity does not match request'); + } + const manifestArtifacts = expectJsonRecord( + manifestRecord.artifacts ?? {}, + 'snapshot manifest artifacts' + ); + const artifactSizes: { homeBytes?: number; wipBytes?: number } = {}; + for (const [artifact, sizeKey] of [ + ['home', 'homeBytes'], + ['wip', 'wipBytes'], + ] as const) { + if (!(artifact in manifestArtifacts)) continue; + const artifactRecord = expectJsonRecord( + manifestArtifacts[artifact], + 'snapshot ' + artifact + ' artifact' + ); + const claimedSize = artifactRecord.sizeBytes; + if (typeof claimedSize !== 'number' || !Number.isSafeInteger(claimedSize) || claimedSize < 0) { + throw errors.badRequest('Snapshot ' + artifact + ' size is invalid'); + } + const object = await c.env.R2.head(buildSessionSnapshotR2Key(c.env, chatSessionId, artifact)); + if (!object) throw errors.badRequest('Snapshot ' + artifact + ' artifact is missing'); + if (object.size !== claimedSize) + throw errors.badRequest('Snapshot ' + artifact + ' size does not match upload'); + artifactSizes[sizeKey] = object.size; + } + const totalArtifactBytes = (artifactSizes.homeBytes ?? 0) + (artifactSizes.wipBytes ?? 0); + if (totalArtifactBytes > getSessionSnapshotConfig(c.env).totalBudgetBytes) { + throw errors.badRequest('Snapshot artifacts exceed configured total budget'); + } + + await completeSessionSnapshot(db, c.env, { + workspaceId, + chatSessionId, + agentSessionId: stringField(body, 'agentSessionId', false), + runtime: stringField(body, 'runtime', false) || 'runtime-neutral', + baseCommit: stringField(body, 'baseCommit', false), + status, + degradation, + manifest, + artifactSizes, + }); + + return c.json({ status, degradation }); +}); + +sessionSnapshotRoutes.get('/:id/session-snapshot/restore', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const { db, workspace } = await requireWorkspace(c, workspaceId); + const chatSessionId = c.req.query('chatSessionId')?.trim() || workspace.chatSessionId; + if (!chatSessionId) throw errors.badRequest('chatSessionId is required'); + if (workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + + const snapshot = await getRestorableSessionSnapshot(db, chatSessionId); + if (!snapshot) { + return c.json({ available: false, reason: 'snapshot_missing_or_unavailable' }); + } + + return c.json({ + available: true, + status: snapshot.status, + degradation: snapshot.degradation, + baseCommit: snapshot.baseCommit, + manifest: snapshot.manifestJson ? JSON.parse(snapshot.manifestJson) : null, + config: getSessionSnapshotConfig(c.env), + download: { + home: snapshot.homeR2Key + ? `/api/workspaces/${workspaceId}/session-snapshot/artifacts/home?chatSessionId=${encodeURIComponent(chatSessionId)}` + : null, + wip: snapshot.wipR2Key + ? `/api/workspaces/${workspaceId}/session-snapshot/artifacts/wip?chatSessionId=${encodeURIComponent(chatSessionId)}` + : null, + manifest: `/api/workspaces/${workspaceId}/session-snapshot/artifacts/manifest?chatSessionId=${encodeURIComponent(chatSessionId)}`, + }, + }); +}); + +sessionSnapshotRoutes.get('/:id/session-snapshot/artifacts/:artifact', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const artifact = artifactFromParam(c.req.param('artifact')); + const chatSessionId = c.req.query('chatSessionId')?.trim(); + if (!chatSessionId) throw errors.badRequest('chatSessionId is required'); + const { workspace } = await requireWorkspace(c, workspaceId); + if (workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + const object = await c.env.R2.get(buildSessionSnapshotR2Key(c.env, chatSessionId, artifact)); + if (!object) throw errors.notFound('Snapshot artifact'); + const headers = new Headers(); + object.writeHttpMetadata(headers); + headers.set('etag', object.httpEtag); + return new Response(object.body, { headers }); +}); + +sessionSnapshotRoutes.post('/:id/session-snapshot/restore-result', async (c) => { + const workspaceId = c.req.param('id'); + await verifyWorkspaceCallbackAuth(c, workspaceId); + const { db, workspace } = await requireWorkspace(c, workspaceId); + const body = await readJsonBody(c); + const chatSessionId = requiredStringField(body, 'chatSessionId'); + if (workspace.chatSessionId !== chatSessionId) { + throw errors.forbidden('Snapshot chat session does not match workspace'); + } + await recordSessionSnapshotRestoreResult(db, { + chatSessionId, + status: requiredStringField(body, 'status'), + message: stringField(body, 'message', false), + }); + return c.body(null, 204); +}); + +export { sessionSnapshotRoutes }; diff --git a/apps/api/src/services/node-agent-diagnostics.ts b/apps/api/src/services/node-agent-diagnostics.ts new file mode 100644 index 000000000..a8b27f6a6 --- /dev/null +++ b/apps/api/src/services/node-agent-diagnostics.ts @@ -0,0 +1,37 @@ +import type { Env } from '../env'; +import { nodeAgentRequest } from './node-agent'; + +export async function getNodeSystemInfoFromNode( + nodeId: string, + env: Env, + userId: string +): Promise { + return nodeAgentRequest(nodeId, env, '/system-info', { + method: 'GET', + userId, + }); +} + +export async function getNodeLogsFromNode( + nodeId: string, + env: Env, + userId: string, + queryString: string +): Promise { + const path = queryString ? `/logs?${queryString}` : '/logs'; + return nodeAgentRequest(nodeId, env, path, { + method: 'GET', + userId, + }); +} + +export async function listNodeContainersFromNode( + nodeId: string, + env: Env, + userId: string +): Promise { + return nodeAgentRequest(nodeId, env, '/containers', { + method: 'GET', + userId, + }); +} diff --git a/apps/api/src/services/node-agent-session-snapshots.ts b/apps/api/src/services/node-agent-session-snapshots.ts new file mode 100644 index 000000000..c94944a7e --- /dev/null +++ b/apps/api/src/services/node-agent-session-snapshots.ts @@ -0,0 +1,47 @@ +import type { Env } from '../env'; +import { getNodeAgentRequestTimeoutMs, nodeAgentRequest } from './node-agent'; + +interface SessionSnapshotRequest { + chatSessionId: string; + runtime: string; +} + +function requestSessionSnapshot( + action: 'hibernate' | 'restore', + nodeId: string, + workspaceId: string, + sessionId: string, + env: Env, + userId: string, + input: SessionSnapshotRequest +): Promise { + return nodeAgentRequest(nodeId, env, `/workspaces/${workspaceId}/agent-sessions/${sessionId}/${action}`, { + method: 'POST', + userId, + workspaceId, + requestTimeoutMs: getNodeAgentRequestTimeoutMs(env), + body: JSON.stringify(input), + }); +} + +export function hibernateAgentSessionOnNode( + nodeId: string, + workspaceId: string, + sessionId: string, + env: Env, + userId: string, + input: SessionSnapshotRequest +): Promise { + return requestSessionSnapshot('hibernate', nodeId, workspaceId, sessionId, env, userId, input); +} + +export function restoreAgentSessionOnNode( + nodeId: string, + workspaceId: string, + sessionId: string, + env: Env, + userId: string, + input: SessionSnapshotRequest +): Promise { + return requestSessionSnapshot('restore', nodeId, workspaceId, sessionId, env, userId, input); +} diff --git a/apps/api/src/services/node-agent.ts b/apps/api/src/services/node-agent.ts index 401bd5f87..4e5c76389 100644 --- a/apps/api/src/services/node-agent.ts +++ b/apps/api/src/services/node-agent.ts @@ -15,6 +15,7 @@ import { } from './vm-agent-container'; const DEFAULT_NODE_AGENT_REQUEST_TIMEOUT_MS = 30_000; +const DEFAULT_CF_CONTAINER_WAKE_TIMEOUT_MS = 120_000; const DEFAULT_NODE_AGENT_READY_TIMEOUT_MS = 900_000; // 15 min — cloud-init takes 8-12 min on Hetzner const DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS = 5000; @@ -66,6 +67,10 @@ export function getNodeAgentRequestTimeoutMs(env: { NODE_AGENT_REQUEST_TIMEOUT_M return getTimeoutMs(env.NODE_AGENT_REQUEST_TIMEOUT_MS, DEFAULT_NODE_AGENT_REQUEST_TIMEOUT_MS); } +export function getCfContainerWakeTimeoutMs(env: { CF_CONTAINER_WAKE_TIMEOUT_MS?: string }): number { + return getTimeoutMs(env.CF_CONTAINER_WAKE_TIMEOUT_MS, DEFAULT_CF_CONTAINER_WAKE_TIMEOUT_MS); +} + export async function waitForNodeAgentReady(nodeId: string, env: Env): Promise { const timeoutMs = getNodeAgentReadyTimeoutMs(env); const pollIntervalMs = getNodeAgentReadyPollIntervalMs(env); @@ -125,7 +130,7 @@ export async function waitForNodeAgentReady(nodeId: string, env: Env): Promise { - return nodeAgentRequest(nodeId, env, '/system-info', { - method: 'GET', - userId, - }); -} - -export async function getNodeLogsFromNode( - nodeId: string, - env: Env, - userId: string, - queryString: string -): Promise { - const path = queryString ? `/logs?${queryString}` : '/logs'; - return nodeAgentRequest(nodeId, env, path, { - method: 'GET', - userId, - }); -} - -export async function listNodeContainersFromNode( - nodeId: string, - env: Env, - userId: string -): Promise { - return nodeAgentRequest(nodeId, env, '/containers', { - method: 'GET', - userId, - }); -} - /** * Raw binary proxy to a VM agent endpoint. * Returns the raw Response (not parsed as JSON) so callers can stream the body. diff --git a/apps/api/src/services/session-snapshots.ts b/apps/api/src/services/session-snapshots.ts new file mode 100644 index 000000000..393209b7d --- /dev/null +++ b/apps/api/src/services/session-snapshots.ts @@ -0,0 +1,279 @@ +import { eq } from 'drizzle-orm'; +import { type drizzle } from 'drizzle-orm/d1'; + +import * as schema from '../db/schema'; +import type { Env } from '../env'; +import { log } from '../lib/logger'; +import { parsePositiveInt } from '../lib/route-helpers'; +import { ulid } from '../lib/ulid'; + +type Db = ReturnType>; + +export const DEFAULT_SESSION_SNAPSHOT_TTL_DAYS = 7; +export const DEFAULT_SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES = 100 * 1024 * 1024; +export const DEFAULT_SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES = 50 * 1024 * 1024; +export const DEFAULT_SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS = 30_000; +export const DEFAULT_SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES = 256 * 1024; +export const DEFAULT_SESSION_SNAPSHOT_R2_PREFIX = 'session-snapshots'; + +export type SessionSnapshotArtifact = 'home' | 'wip' | 'manifest'; +export type SessionSnapshotStatus = 'pending' | 'available' | 'degraded' | 'failed' | 'expired'; +export type SessionSnapshotDegradation = 'none' | 'home-skipped' | 'wip-only' | 'transcript-only'; + +export interface SessionSnapshotManifest { + version: 1; + chatSessionId: string; + workspaceId: string; + agentSessionId?: string; + baseCommit?: string; + status: SessionSnapshotStatus; + degradation: SessionSnapshotDegradation; + skipped: Array<{ + path: string; + reason: string; + sizeBytes?: number; + }>; + artifacts: { + home?: { sizeBytes: number; sha256?: string }; + wip?: { sizeBytes: number; sha256?: string }; + }; + createdAt: string; +} + +export interface SessionSnapshotConfig { + ttlDays: number; + totalBudgetBytes: number; + entryThresholdBytes: number; + transferIdleTimeoutMs: number; + jsonBodyMaxBytes: number; + r2Prefix: string; +} + +export interface PrepareSessionSnapshotInput { + workspaceId: string; + nodeId: string | null; + projectId: string | null; + userId: string; + chatSessionId: string; + agentSessionId: string | null; + runtime: string; +} + +export interface CompleteSessionSnapshotInput { + workspaceId: string; + chatSessionId: string; + agentSessionId: string | null; + runtime: string; + baseCommit: string | null; + status: SessionSnapshotStatus; + degradation: SessionSnapshotDegradation; + manifest: SessionSnapshotManifest; + artifactSizes: { + homeBytes?: number; + wipBytes?: number; + }; +} + +export function getSessionSnapshotConfig(env: Env): SessionSnapshotConfig { + return { + ttlDays: parsePositiveInt(env.SESSION_SNAPSHOT_TTL_DAYS, DEFAULT_SESSION_SNAPSHOT_TTL_DAYS), + totalBudgetBytes: parsePositiveInt( + env.SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES, + DEFAULT_SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES + ), + entryThresholdBytes: parsePositiveInt( + env.SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES, + DEFAULT_SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES + ), + transferIdleTimeoutMs: parsePositiveInt( + env.SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS, + DEFAULT_SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS + ), + jsonBodyMaxBytes: parsePositiveInt( + env.SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES, + DEFAULT_SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES + ), + r2Prefix: sanitizeSnapshotPrefix(env.SESSION_SNAPSHOT_R2_PREFIX || DEFAULT_SESSION_SNAPSHOT_R2_PREFIX), + }; +} + +function sanitizeSnapshotPrefix(prefix: string): string { + const normalized = prefix + .trim() + .replace(/^\/+|\/+$/g, '') + .replace(/[^A-Za-z0-9._/-]/g, '-'); + return normalized || DEFAULT_SESSION_SNAPSHOT_R2_PREFIX; +} + +function sanitizeKeySegment(value: string): string { + return value.trim().replace(/[^A-Za-z0-9._-]/g, '-'); +} + +export function buildSessionSnapshotR2Key( + env: Env, + chatSessionId: string, + artifact: SessionSnapshotArtifact +): string { + const { r2Prefix } = getSessionSnapshotConfig(env); + const sessionKey = sanitizeKeySegment(chatSessionId); + const suffix = artifact === 'home' ? 'home.tar' : artifact === 'wip' ? 'wip.bundle' : 'manifest.json'; + return `${r2Prefix}/${sessionKey}/${suffix}`; +} + +function snapshotExpiry(now: Date, ttlDays: number): string { + return new Date(now.getTime() + ttlDays * 24 * 60 * 60 * 1000).toISOString(); +} + +export async function prepareSessionSnapshot( + db: Db, + env: Env, + input: PrepareSessionSnapshotInput +): Promise<{ + snapshotId: string; + expiresAt: string; + keys: Record; + config: SessionSnapshotConfig; +}> { + const config = getSessionSnapshotConfig(env); + const now = new Date(); + const expiresAt = snapshotExpiry(now, config.ttlDays); + const keys = { + home: buildSessionSnapshotR2Key(env, input.chatSessionId, 'home'), + wip: buildSessionSnapshotR2Key(env, input.chatSessionId, 'wip'), + manifest: buildSessionSnapshotR2Key(env, input.chatSessionId, 'manifest'), + }; + + const existing = await db + .select({ id: schema.sessionSnapshots.id }) + .from(schema.sessionSnapshots) + .where(eq(schema.sessionSnapshots.chatSessionId, input.chatSessionId)) + .limit(1); + + const snapshotId = existing[0]?.id || ulid(); + const row = { + id: snapshotId, + projectId: input.projectId, + workspaceId: input.workspaceId, + nodeId: input.nodeId, + userId: input.userId, + chatSessionId: input.chatSessionId, + agentSessionId: input.agentSessionId, + runtime: input.runtime, + status: 'pending', + degradation: 'none', + homeR2Key: keys.home, + wipR2Key: keys.wip, + manifestR2Key: keys.manifest, + baseCommit: null, + expiresAt, + manifestJson: null, + restoreStatus: null, + restoreMessage: null, + restoredAt: null, + updatedAt: now.toISOString(), + } satisfies schema.NewSessionSnapshot; + + if (existing[0]) { + await db.update(schema.sessionSnapshots).set(row).where(eq(schema.sessionSnapshots.id, snapshotId)); + } else { + await db.insert(schema.sessionSnapshots).values({ ...row, createdAt: now.toISOString() }); + } + + return { snapshotId, expiresAt, keys, config }; +} + +export async function completeSessionSnapshot( + db: Db, + env: Env, + input: CompleteSessionSnapshotInput +): Promise { + const homeR2Key = + input.artifactSizes.homeBytes == null + ? null + : buildSessionSnapshotR2Key(env, input.chatSessionId, 'home'); + const wipR2Key = + input.artifactSizes.wipBytes == null + ? null + : buildSessionSnapshotR2Key(env, input.chatSessionId, 'wip'); + const manifestR2Key = buildSessionSnapshotR2Key(env, input.chatSessionId, 'manifest'); + const manifestJson = JSON.stringify(input.manifest); + await env.R2.put(manifestR2Key, manifestJson, { + httpMetadata: { contentType: 'application/json' }, + }); + + await db + .update(schema.sessionSnapshots) + .set({ + agentSessionId: input.agentSessionId, + runtime: input.runtime, + status: input.status, + degradation: input.degradation, + homeR2Key, + wipR2Key, + manifestR2Key, + baseCommit: input.baseCommit, + manifestJson, + expiresAt: + (await db + .select({ expiresAt: schema.sessionSnapshots.expiresAt }) + .from(schema.sessionSnapshots) + .where(eq(schema.sessionSnapshots.chatSessionId, input.chatSessionId)) + .limit(1))[0]?.expiresAt || snapshotExpiry(new Date(), getSessionSnapshotConfig(env).ttlDays), + updatedAt: new Date().toISOString(), + }) + .where(eq(schema.sessionSnapshots.chatSessionId, input.chatSessionId)); +} + +export async function getRestorableSessionSnapshot( + db: Db, + chatSessionId: string, + now = new Date() +): Promise { + const rows = await db + .select() + .from(schema.sessionSnapshots) + .where(eq(schema.sessionSnapshots.chatSessionId, chatSessionId)) + .limit(1); + const snapshot = rows[0]; + if (!snapshot) return null; + if (Date.parse(snapshot.expiresAt) <= now.getTime()) { + return { ...snapshot, status: 'expired' }; + } + if (snapshot.status !== 'available' && snapshot.status !== 'degraded') { + return null; + } + return snapshot; +} + +export async function recordSessionSnapshotRestoreResult( + db: Db, + input: { + chatSessionId: string; + status: string; + message: string | null; + } +): Promise { + await db + .update(schema.sessionSnapshots) + .set({ + restoreStatus: input.status, + restoreMessage: input.message, + restoredAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }) + .where(eq(schema.sessionSnapshots.chatSessionId, input.chatSessionId)); +} + +export async function deleteSessionSnapshotArtifacts(env: Env, chatSessionId: string): Promise { + await Promise.all( + (['home', 'wip', 'manifest'] as const).map((artifact) => + env.R2.delete(buildSessionSnapshotR2Key(env, chatSessionId, artifact)).catch((err) => { + log.warn('session_snapshot.r2_delete_failed', { + chatSessionId, + artifact, + error: err instanceof Error ? err.message : String(err), + }); + }) + ) + ); +} diff --git a/apps/api/tests/unit/cf-container-runtime-contract.test.ts b/apps/api/tests/unit/cf-container-runtime-contract.test.ts index cc7a8fe96..6e3a86943 100644 --- a/apps/api/tests/unit/cf-container-runtime-contract.test.ts +++ b/apps/api/tests/unit/cf-container-runtime-contract.test.ts @@ -132,13 +132,15 @@ describe('cf-container runtime spike contracts', () => { expect(containerDo).toContain("await this.markRuntimeSleeping('Container idle timeout expired; container is sleeping.')"); expect(containerDo).toContain("await this.ctx.storage.put('lifecycleStatus', 'sleeping' satisfies LifecycleStatus)"); expect(containerDo).toContain("status: 'sleeping'"); - expect(containerDo).toContain('Container is asleep; wake/rehydrate is not implemented yet.'); - expect(containerDo).toContain('Phase 3 of idea 01KX4KSXEXQMP41KS34TW9EN01'); - expect(containerDo).toContain("return new Response(SLEEPING_RESPONSE, { status: 503 })"); + expect(containerDo).toContain("if (lifecycleStatus === 'sleeping')"); + expect(containerDo).toContain('const wake = await this.ensureAwake()'); + expect(containerDo).toContain('WAKE_DEGRADED_RESPONSE'); + expect(containerDo).toContain("await this.ctx.storage.put('lifecycleStatus', 'launching' satisfies LifecycleStatus)"); + expect(containerDo).toContain("await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus)"); expect(chatResolver).toContain("inArray(schema.workspaces.status, ['running', 'recovery', 'sleeping'])"); - expect(chatResolver).toContain("workspace.nodeStatus === 'sleeping'"); - expect(chatResolver).toContain('The workspace container is asleep.'); - expect(chatResolver).toContain('Phase 3 of idea 01KX4KSXEXQMP41KS34TW9EN01'); + expect(chatResolver).toContain("workspace.nodeStatus !== 'running' && workspace.nodeStatus !== 'sleeping'"); + expect(chatResolver).toContain("inArray(schema.agentSessions.status, ['running', 'sleeping'])"); + expect(chatResolver).not.toContain('The workspace container is asleep.'); expect(containerDo).toContain("return new Response('Container is stopped; create a new instant session.', { status: 410 })"); expect(containerDo).toContain("await this.ctx.storage.put('launchConfig', config)"); expect(containerDo).toContain('nodeCallbackToken: string'); diff --git a/apps/api/tests/unit/durable-objects/vm-agent-container-wake-state.test.ts b/apps/api/tests/unit/durable-objects/vm-agent-container-wake-state.test.ts new file mode 100644 index 000000000..2c3028a55 --- /dev/null +++ b/apps/api/tests/unit/durable-objects/vm-agent-container-wake-state.test.ts @@ -0,0 +1,141 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { VmAgentContainer } from '../../../src/durable-objects/vm-agent-container'; + +// Regression test for the restored-session prompt failure: `proxyHttp` read the +// container state ONCE before `wakeFromSnapshot()`, then applied the +// `stopped`/`stopped_with_code` -> 410 guard using that stale pre-wake state. +// A freshly-woken, restored container was therefore rejected with 410 (surfaced +// by the Worker as a generic 500), even though restore succeeded. The fix +// re-reads the container state after a successful wake. + +interface FakeState { + status: string; +} + +function makeFake(opts: { + statuses: string[]; // sequence returned by getState() + lifecycleStatus: string; + wakeOk: boolean; +}) { + const getState = vi.fn<[], Promise>(); + for (const s of opts.statuses) { + getState.mockResolvedValueOnce({ status: s }); + } + getState.mockResolvedValue({ status: opts.statuses[opts.statuses.length - 1] }); + + const containerFetch = vi + .fn() + .mockResolvedValue(new Response('proxied', { status: 200 })); + const wakeFromSnapshot = vi + .fn() + .mockResolvedValue(opts.wakeOk ? { ok: true } : { ok: false, message: 'degraded' }); + + const fake = { + getState, + containerFetch, + wakeFromSnapshot, + defaultPort: 8080, + wakeChain: Promise.resolve(), + ensureAwake: (VmAgentContainer.prototype as unknown as { ensureAwake: unknown }) + .ensureAwake, + ctx: { storage: { get: vi.fn().mockResolvedValue(opts.lifecycleStatus) } }, + }; + return { fake, getState, containerFetch, wakeFromSnapshot }; +} + +function callProxyHttp(fake: unknown, request: Request): Promise { + return (VmAgentContainer.prototype as unknown as { + proxyHttp: (this: unknown, request: Request, port?: number) => Promise; + }).proxyHttp.call(fake, request); +} + +describe('VmAgentContainer.proxyHttp wake state re-read', () => { + it('proxies the prompt after a successful wake even though the pre-wake state was stopped', async () => { + // Pre-wake getState -> stopped; post-wake getState -> running (fresh container). + const { fake, getState, containerFetch, wakeFromSnapshot } = makeFake({ + statuses: ['stopped', 'running'], + lifecycleStatus: 'sleeping', + wakeOk: true, + }); + + const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + + expect(wakeFromSnapshot).toHaveBeenCalledTimes(1); + // State must be re-read after wake (once before, once after) so the stopped + // guard sees the now-running container. + expect(getState).toHaveBeenCalledTimes(2); + // The request is proxied to the running container, NOT rejected with 410. + expect(containerFetch).toHaveBeenCalledTimes(1); + expect(res.status).toBe(200); + }); + + it('returns 503 (not 410/proxy) when wake fails', async () => { + const { fake, containerFetch } = makeFake({ + statuses: ['stopped', 'stopped'], + lifecycleStatus: 'sleeping', + wakeOk: false, + }); + + const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + + expect(res.status).toBe(503); + expect(containerFetch).not.toHaveBeenCalled(); + }); + + it('still returns 410 for a genuinely stopped, non-sleeping container', async () => { + const { fake, containerFetch, wakeFromSnapshot } = makeFake({ + statuses: ['stopped'], + lifecycleStatus: 'running', + wakeOk: true, + }); + + const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + + expect(wakeFromSnapshot).not.toHaveBeenCalled(); + expect(containerFetch).not.toHaveBeenCalled(); + expect(res.status).toBe(410); + }); +}); + +describe('VmAgentContainer.ensureAwake concurrency (rule 45)', () => { + it('wakes a sleeping container exactly once under two concurrent requests', async () => { + // Shared, mutable container state so the mock models a real wake: the first + // wake flips lifecycleStatus to running, and getState follows it. + const shared = { lifecycle: 'sleeping' as string }; + + const getState = vi.fn(async () => ({ + status: shared.lifecycle === 'running' ? 'running' : 'stopped', + })); + const containerFetch = vi.fn().mockResolvedValue(new Response('proxied', { status: 200 })); + const wakeFromSnapshot = vi.fn(async () => { + // Simulate the async launch+restore so the two requests interleave across + // this await; the second must observe the running state and NOT re-wake. + await new Promise((r) => setTimeout(r, 20)); + shared.lifecycle = 'running'; + return { ok: true }; + }); + + const fake = { + getState, + containerFetch, + wakeFromSnapshot, + defaultPort: 8080, + wakeChain: Promise.resolve(), + ensureAwake: (VmAgentContainer.prototype as unknown as { ensureAwake: unknown }).ensureAwake, + ctx: { storage: { get: vi.fn(async () => shared.lifecycle) } }, + }; + + const [a, b] = await Promise.all([ + callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })), + callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })), + ]); + + // The one-time launch+restore fired exactly once despite two concurrent + // requests; both requests were proxied to the now-running container. + expect(wakeFromSnapshot).toHaveBeenCalledTimes(1); + expect(containerFetch).toHaveBeenCalledTimes(2); + expect(a.status).toBe(200); + expect(b.status).toBe(200); + }); +}); diff --git a/apps/api/tests/unit/mcp-deployment-tools.test.ts b/apps/api/tests/unit/mcp-deployment-tools.test.ts index 6c7b11c77..3895bdff4 100644 --- a/apps/api/tests/unit/mcp-deployment-tools.test.ts +++ b/apps/api/tests/unit/mcp-deployment-tools.test.ts @@ -31,6 +31,12 @@ vi.mock('../../src/services/node-agent', () => ({ getNodeLogsFromNode: (...args: unknown[]) => mockGetNodeLogsFromNode(...args), })); +vi.mock('../../src/services/node-agent-diagnostics', () => ({ + getNodeLogsFromNode: (...args: unknown[]) => mockGetNodeLogsFromNode(...args), + getNodeSystemInfoFromNode: vi.fn(), + listNodeContainersFromNode: vi.fn(), +})); + const schema = await import('../../src/db/schema'); const { handleCreateDeploymentEnvironment, diff --git a/apps/api/tests/unit/node-agent-health.test.ts b/apps/api/tests/unit/node-agent-health.test.ts index 9ecbc759e..1beb54496 100644 --- a/apps/api/tests/unit/node-agent-health.test.ts +++ b/apps/api/tests/unit/node-agent-health.test.ts @@ -12,10 +12,23 @@ import { afterEach,beforeEach, describe, expect, it, vi } from 'vitest'; import { + getCfContainerWakeTimeoutMs, getNodeAgentReadyPollIntervalMs, getNodeAgentReadyTimeoutMs, } from '../../src/services/node-agent'; +describe('getCfContainerWakeTimeoutMs', () => { + it('defaults to two minutes and accepts a positive override', () => { + expect(getCfContainerWakeTimeoutMs({})).toBe(120_000); + expect(getCfContainerWakeTimeoutMs({ CF_CONTAINER_WAKE_TIMEOUT_MS: '180000' })).toBe(180_000); + }); + + it('rejects non-positive and invalid overrides', () => { + expect(getCfContainerWakeTimeoutMs({ CF_CONTAINER_WAKE_TIMEOUT_MS: '0' })).toBe(120_000); + expect(getCfContainerWakeTimeoutMs({ CF_CONTAINER_WAKE_TIMEOUT_MS: 'invalid' })).toBe(120_000); + }); +}); + // ============================================================================= // getNodeAgentReadyTimeoutMs — env var parsing // ============================================================================= diff --git a/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts b/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts index 461ae9ced..b8e88bece 100644 --- a/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts +++ b/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts @@ -11,6 +11,7 @@ const mocks = vi.hoisted(() => ({ requireProjectCapability: vi.fn(), sendPromptToAgentOnNode: vi.fn(), cancelAgentSessionOnNode: vi.fn(), + getCfContainerWakeTimeoutMs: vi.fn(() => 120_000), enrichMessageWithMentions: vi.fn(), parseOptionalBody: vi.fn(), })); @@ -80,6 +81,7 @@ vi.mock('../../../src/schemas', () => ({ vi.mock('../../../src/services/node-agent', () => ({ sendPromptToAgentOnNode: mocks.sendPromptToAgentOnNode, cancelAgentSessionOnNode: mocks.cancelAgentSessionOnNode, + getCfContainerWakeTimeoutMs: mocks.getCfContainerWakeTimeoutMs, })); vi.mock('../../../src/services/mention-enrichment', () => ({ @@ -236,7 +238,24 @@ describe('POST /sessions/:sessionId/prompt', () => { expect(response.status).toBe(200); expect(mocks.sendPromptToAgentOnNode).toHaveBeenCalledWith( 'node-1', 'ws-1', 'agent-sess-1', - 'hello agent', expect.anything(), 'user-1', + 'hello agent', expect.anything(), 'user-1', undefined, undefined, + ); + }); + + it('uses the extended wake budget for a sleeping session', async () => { + setupDrizzle({ + workspace: { id: 'ws-1', nodeId: 'node-1', nodeStatus: 'sleeping' }, + agentSession: { id: 'agent-sess-1' }, + }); + mocks.sendPromptToAgentOnNode.mockResolvedValue({ ok: true }); + + const response = await postPrompt(); + + expect(response.status).toBe(200); + expect(mocks.sendPromptToAgentOnNode).toHaveBeenCalledWith( + 'node-1', 'ws-1', 'agent-sess-1', + 'hello agent', expect.anything(), 'user-1', undefined, + { requestTimeoutMs: 120_000 }, ); }); diff --git a/apps/api/tests/unit/routes/deployment-environment-observability.test.ts b/apps/api/tests/unit/routes/deployment-environment-observability.test.ts index 6ed76e1ff..1cfe3279a 100644 --- a/apps/api/tests/unit/routes/deployment-environment-observability.test.ts +++ b/apps/api/tests/unit/routes/deployment-environment-observability.test.ts @@ -76,6 +76,12 @@ vi.mock('../../../src/services/node-agent', () => ({ mockTeardownDeploymentEnvironmentOnNode(...args), })); +vi.mock('../../../src/services/node-agent-diagnostics', () => ({ + getNodeLogsFromNode: (...args: unknown[]) => mockGetNodeLogsFromNode(...args), + getNodeSystemInfoFromNode: (...args: unknown[]) => mockGetNodeSystemInfoFromNode(...args), + listNodeContainersFromNode: (...args: unknown[]) => mockListNodeContainersFromNode(...args), +})); + vi.mock('../../../src/services/deployment-control', () => ({ DEPLOYMENT_ENVIRONMENT_NAME_RE: /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$/, encodeAllowedDeployProfileIds: vi.fn(() => null), diff --git a/apps/api/tests/unit/routes/deployment-membership-auth.test.ts b/apps/api/tests/unit/routes/deployment-membership-auth.test.ts index 659c7b930..ba6fc364b 100644 --- a/apps/api/tests/unit/routes/deployment-membership-auth.test.ts +++ b/apps/api/tests/unit/routes/deployment-membership-auth.test.ts @@ -143,6 +143,12 @@ vi.mock('../../../src/services/node-agent', () => ({ teardownDeploymentEnvironmentOnNode: vi.fn(), })); +vi.mock('../../../src/services/node-agent-diagnostics', () => ({ + getNodeLogsFromNode: (...args: unknown[]) => mockGetNodeLogsFromNode(...args), + getNodeSystemInfoFromNode: vi.fn(), + listNodeContainersFromNode: vi.fn(), +})); + vi.mock('../../../src/services/deployment-provisioning', () => ({ provisionDeploymentNode: vi.fn(), })); diff --git a/apps/api/tests/unit/routes/node-observability-logs.test.ts b/apps/api/tests/unit/routes/node-observability-logs.test.ts index 83635c204..54ce1e2e4 100644 --- a/apps/api/tests/unit/routes/node-observability-logs.test.ts +++ b/apps/api/tests/unit/routes/node-observability-logs.test.ts @@ -26,6 +26,12 @@ vi.mock('../../../src/services/node-agent', () => ({ stopWorkspaceOnNode: vi.fn(), })); +vi.mock('../../../src/services/node-agent-diagnostics', () => ({ + getNodeLogsFromNode: (...args: unknown[]) => mockGetNodeLogsFromNode(...args), + listNodeContainersFromNode: (...args: unknown[]) => mockListNodeContainersFromNode(...args), + getNodeSystemInfoFromNode: vi.fn(), +})); + vi.mock('../../../src/services/nodes', () => ({ createNodeRecord: vi.fn(), deleteNodeResources: vi.fn(), diff --git a/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts b/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts new file mode 100644 index 000000000..86ed800e2 --- /dev/null +++ b/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts @@ -0,0 +1,265 @@ +import { drizzle } from 'drizzle-orm/d1'; +import type { Hono } from 'hono'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { Env } from '../../../src/env'; +import { workspacesRoutes } from '../../../src/routes/workspaces'; +import { createRouteTestApp } from './route-test-app'; + +const mocks = vi.hoisted(() => ({ + completeSessionSnapshot: vi.fn(), + getRestorableSessionSnapshot: vi.fn(), + prepareSessionSnapshot: vi.fn(), + recordSessionSnapshotRestoreResult: vi.fn(), + verifyCallbackToken: vi.fn(), +})); + +vi.mock('drizzle-orm/d1'); +vi.mock('../../../src/middleware/auth', () => ({ + requireAuth: () => vi.fn((c: any, next: any) => next()), + requireApproved: () => vi.fn((c: any, next: any) => next()), + getUserId: () => 'user-1', + getAuth: () => ({ user: { id: 'user-1', name: 'User', email: 'user@example.com' } }), +})); +vi.mock('../../../src/services/jwt', () => ({ + verifyCallbackToken: mocks.verifyCallbackToken, + signCallbackToken: vi.fn(), +})); +vi.mock('../../../src/services/session-snapshots', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + completeSessionSnapshot: mocks.completeSessionSnapshot, + getRestorableSessionSnapshot: mocks.getRestorableSessionSnapshot, + prepareSessionSnapshot: mocks.prepareSessionSnapshot, + recordSessionSnapshotRestoreResult: mocks.recordSessionSnapshotRestoreResult, + }; +}); + +function makeDb(workspace: Record) { + return { + select: vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(() => ({ + limit: vi.fn(async () => [workspace]), + })), + })), + })), + }; +} + +describe('workspaces session snapshot callback routes', () => { + let app: Hono<{ Bindings: Env }>; + const r2 = { + put: vi.fn(), + get: vi.fn(), + head: vi.fn(), + }; + const runtimeBindings = { + DATABASE: {} as any, + R2: r2, + ENCRYPTION_KEY: 'enc-key', + SESSION_SNAPSHOT_R2_PREFIX: 'test-snapshots', + SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES: '1024', + } as unknown as Env; + const workspace = { + id: 'WS_1', + nodeId: 'node-1', + projectId: 'project-1', + userId: 'user-1', + chatSessionId: 'chat-1', + }; + + beforeEach(() => { + vi.clearAllMocks(); + (drizzle as any).mockReturnValue(makeDb(workspace)); + mocks.verifyCallbackToken.mockResolvedValue({ + workspace: 'WS_1', + type: 'callback', + scope: 'workspace', + }); + mocks.prepareSessionSnapshot.mockResolvedValue({ + snapshotId: 'snapshot-1', + expiresAt: '2026-07-18T00:00:00.000Z', + config: { + ttlDays: 7, + totalBudgetBytes: 1024, + entryThresholdBytes: 512, + transferIdleTimeoutMs: 30000, + jsonBodyMaxBytes: 262144, + r2Prefix: 'test-snapshots', + }, + keys: { + home: 'test-snapshots/chat-1/home.tar', + wip: 'test-snapshots/chat-1/wip.bundle', + manifest: 'test-snapshots/chat-1/manifest.json', + }, + }); + + app = createRouteTestApp('/api/workspaces', workspacesRoutes); + }); + + it('prepares deterministic upload URLs for the workspace chat session', async () => { + const res = await app.request( + '/api/workspaces/WS_1/session-snapshot/prepare', + { + method: 'POST', + headers: { + Authorization: 'Bearer callback-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + chatSessionId: 'chat-1', + agentSessionId: 'agent-session-1', + runtime: 'cf-container', + }), + }, + runtimeBindings + ); + + expect(res.status).toBe(200); + await expect(res.json()).resolves.toMatchObject({ + snapshotId: 'snapshot-1', + upload: { + home: '/api/workspaces/WS_1/session-snapshot/artifacts/home?chatSessionId=chat-1', + wip: '/api/workspaces/WS_1/session-snapshot/artifacts/wip?chatSessionId=chat-1', + }, + }); + expect(mocks.prepareSessionSnapshot).toHaveBeenCalledWith(expect.anything(), runtimeBindings, { + workspaceId: 'WS_1', + nodeId: 'node-1', + projectId: 'project-1', + userId: 'user-1', + chatSessionId: 'chat-1', + agentSessionId: 'agent-session-1', + runtime: 'cf-container', + }); + }); + + it('uploads artifacts to server-derived R2 keys and rejects oversized content-lengths', async () => { + const ok = await app.request( + '/api/workspaces/WS_1/session-snapshot/artifacts/home?chatSessionId=chat-1', + { + method: 'PUT', + headers: { + Authorization: 'Bearer callback-token', + 'Content-Type': 'application/octet-stream', + 'Content-Length': '4', + }, + body: 'home', + }, + runtimeBindings + ); + + expect(ok.status).toBe(200); + expect(r2.put).toHaveBeenCalledWith( + 'test-snapshots/chat-1/home.tar', + expect.anything(), + expect.objectContaining({ httpMetadata: { contentType: 'application/x-tar' } }) + ); + + const tooLarge = await app.request( + '/api/workspaces/WS_1/session-snapshot/artifacts/wip?chatSessionId=chat-1', + { + method: 'PUT', + headers: { + Authorization: 'Bearer callback-token', + 'Content-Length': '1025', + }, + body: 'wip', + }, + runtimeBindings + ); + expect(tooLarge.status).toBe(400); + }); + + it('rejects node-scoped callback tokens before snapshot service access', async () => { + mocks.verifyCallbackToken.mockResolvedValueOnce({ + workspace: 'WS_1', + type: 'callback', + scope: 'node', + }); + + const res = await app.request( + '/api/workspaces/WS_1/session-snapshot/prepare', + { + method: 'POST', + headers: { + Authorization: 'Bearer callback-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ chatSessionId: 'chat-1' }), + }, + runtimeBindings + ); + + expect(res.status).toBe(403); + expect(mocks.prepareSessionSnapshot).not.toHaveBeenCalled(); + }); + + it('rejects artifact uploads without an authoritative Content-Length', async () => { + const res = await app.request( + '/api/workspaces/WS_1/session-snapshot/artifacts/home?chatSessionId=chat-1', + { + method: 'PUT', + headers: { Authorization: 'Bearer callback-token' }, + body: 'home', + }, + runtimeBindings + ); + expect(res.status).toBe(400); + expect(r2.put).not.toHaveBeenCalled(); + }); + + it('derives completion sizes from R2 and rejects manifest identity mismatches', async () => { + r2.head.mockImplementation(async (key: string) => + key.endsWith('home.tar') ? { size: 4 } : key.endsWith('wip.bundle') ? { size: 3 } : null + ); + const body = { + chatSessionId: 'chat-1', + agentSessionId: 'agent-session-1', + runtime: 'cf-container', + status: 'available', + degradation: 'none', + baseCommit: 'abc123', + artifactSizes: { homeBytes: 999, wipBytes: 999 }, + manifest: { + version: 1, + chatSessionId: 'chat-1', + workspaceId: 'WS_1', + agentSessionId: 'agent-session-1', + status: 'available', + degradation: 'none', + skipped: [], + artifacts: { home: { sizeBytes: 4 }, wip: { sizeBytes: 3 } }, + createdAt: '2026-07-11T00:00:00.000Z', + }, + }; + const ok = await app.request( + '/api/workspaces/WS_1/session-snapshot/complete', + { + method: 'POST', + headers: { Authorization: 'Bearer callback-token', 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }, + runtimeBindings + ); + expect(ok.status).toBe(200); + expect(mocks.completeSessionSnapshot).toHaveBeenCalledWith( + expect.anything(), + runtimeBindings, + expect.objectContaining({ artifactSizes: { homeBytes: 4, wipBytes: 3 } }) + ); + + const mismatch = await app.request( + '/api/workspaces/WS_1/session-snapshot/complete', + { + method: 'POST', + headers: { Authorization: 'Bearer callback-token', 'Content-Type': 'application/json' }, + body: JSON.stringify({ ...body, manifest: { ...body.manifest, workspaceId: 'WS_OTHER' } }), + }, + runtimeBindings + ); + expect(mismatch.status).toBe(400); + }); +}); diff --git a/apps/api/tests/unit/session-snapshots.test.ts b/apps/api/tests/unit/session-snapshots.test.ts new file mode 100644 index 000000000..a2d3e1ef3 --- /dev/null +++ b/apps/api/tests/unit/session-snapshots.test.ts @@ -0,0 +1,66 @@ +import { describe, expect, it } from 'vitest'; + +import type { Env } from '../../src/env'; +import { + buildSessionSnapshotR2Key, + DEFAULT_SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES, + DEFAULT_SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES, + DEFAULT_SESSION_SNAPSHOT_R2_PREFIX, + DEFAULT_SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES, + DEFAULT_SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS, + DEFAULT_SESSION_SNAPSHOT_TTL_DAYS, + getSessionSnapshotConfig, +} from '../../src/services/session-snapshots'; + +function env(overrides: Partial = {}): Env { + return overrides as Env; +} + +describe('session snapshot config', () => { + it('uses constitution-compliant defaults when env vars are absent', () => { + const config = getSessionSnapshotConfig(env()); + + expect(config).toEqual({ + ttlDays: DEFAULT_SESSION_SNAPSHOT_TTL_DAYS, + totalBudgetBytes: DEFAULT_SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES, + entryThresholdBytes: DEFAULT_SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES, + transferIdleTimeoutMs: DEFAULT_SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS, + jsonBodyMaxBytes: DEFAULT_SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES, + r2Prefix: DEFAULT_SESSION_SNAPSHOT_R2_PREFIX, + }); + }); + + it('uses positive env overrides and sanitizes the R2 prefix', () => { + const config = getSessionSnapshotConfig(env({ + SESSION_SNAPSHOT_TTL_DAYS: '3', + SESSION_SNAPSHOT_TOTAL_BUDGET_BYTES: '1234', + SESSION_SNAPSHOT_ENTRY_THRESHOLD_BYTES: '567', + SESSION_SNAPSHOT_TRANSFER_IDLE_TIMEOUT_MS: '890', + SESSION_SNAPSHOT_JSON_BODY_MAX_BYTES: '1024', + SESSION_SNAPSHOT_R2_PREFIX: '/tenant snapshots/private/', + })); + + expect(config.ttlDays).toBe(3); + expect(config.totalBudgetBytes).toBe(1234); + expect(config.entryThresholdBytes).toBe(567); + expect(config.transferIdleTimeoutMs).toBe(890); + expect(config.jsonBodyMaxBytes).toBe(1024); + expect(config.r2Prefix).toBe('tenant-snapshots/private'); + }); +}); + +describe('session snapshot R2 keys', () => { + it('derives deterministic one-snapshot-per-session artifact keys', () => { + const input = env({ SESSION_SNAPSHOT_R2_PREFIX: 'snapshots' }); + + expect(buildSessionSnapshotR2Key(input, 'chat/../session 1', 'home')).toBe( + 'snapshots/chat-..-session-1/home.tar' + ); + expect(buildSessionSnapshotR2Key(input, 'chat/../session 1', 'wip')).toBe( + 'snapshots/chat-..-session-1/wip.bundle' + ); + expect(buildSessionSnapshotR2Key(input, 'chat/../session 1', 'manifest')).toBe( + 'snapshots/chat-..-session-1/manifest.json' + ); + }); +}); diff --git a/apps/api/tests/workers/session-snapshot-wiring.test.ts b/apps/api/tests/workers/session-snapshot-wiring.test.ts new file mode 100644 index 000000000..48fb1e39f --- /dev/null +++ b/apps/api/tests/workers/session-snapshot-wiring.test.ts @@ -0,0 +1,125 @@ +import { env } from 'cloudflare:test'; +import { drizzle } from 'drizzle-orm/d1'; +import { describe, expect, it } from 'vitest'; + +import * as schema from '../../src/db/schema'; +import type { Env } from '../../src/env'; +import { + completeSessionSnapshot, + getRestorableSessionSnapshot, + prepareSessionSnapshot, + type SessionSnapshotManifest, +} from '../../src/services/session-snapshots'; + +const TEST_PREFIX = `snapshot-${Date.now()}`; + +describe('session snapshot D1/R2 worker wiring', () => { + it('creates one restorable snapshot per chat session and writes the manifest to R2', async () => { + const userId = `${TEST_PREFIX}-user`; + const nodeId = `${TEST_PREFIX}-node`; + const workspaceId = `${TEST_PREFIX}-workspace`; + const chatSessionId = `${TEST_PREFIX}-chat`; + const agentSessionId = `${TEST_PREFIX}-agent-session`; + const bindings = { + ...env, + SESSION_SNAPSHOT_R2_PREFIX: `${TEST_PREFIX}/snapshots`, + SESSION_SNAPSHOT_TTL_DAYS: '7', + } as unknown as Env; + const db = drizzle(env.DATABASE, { schema }); + + await env.DATABASE.prepare( + `INSERT INTO users (id, email, name, created_at, updated_at) + VALUES (?, ?, ?, ?, ?)` + ) + .bind(userId, `${userId}@example.com`, 'Snapshot User', Date.now(), Date.now()) + .run(); + await env.DATABASE.prepare( + `INSERT INTO nodes (id, user_id, name, status, vm_size, vm_location, runtime, created_at, updated_at) + VALUES (?, ?, ?, 'running', 'small', 'local', 'cf-container', datetime('now'), datetime('now'))` + ) + .bind(nodeId, userId, 'snapshot-node') + .run(); + await env.DATABASE.prepare( + `INSERT INTO workspaces (id, node_id, user_id, name, repository, branch, status, vm_size, vm_location, chat_session_id, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, 'main', 'sleeping', 'small', 'local', ?, datetime('now'), datetime('now'))` + ) + .bind(workspaceId, nodeId, userId, 'snapshot-workspace', 'owner/repo', chatSessionId) + .run(); + await env.DATABASE.prepare( + `INSERT INTO agent_sessions (id, workspace_id, user_id, status, created_at, updated_at) + VALUES (?, ?, ?, 'sleeping', datetime('now'), datetime('now'))` + ) + .bind(agentSessionId, workspaceId, userId) + .run(); + + const prepared = await prepareSessionSnapshot(db, bindings, { + workspaceId, + nodeId, + projectId: null, + userId, + chatSessionId, + agentSessionId, + runtime: 'cf-container', + }); + + expect(prepared.keys.home).toBe(`${TEST_PREFIX}/snapshots/${chatSessionId}/home.tar`); + expect(prepared.config.ttlDays).toBe(7); + + const manifest: SessionSnapshotManifest = { + version: 1, + chatSessionId, + workspaceId, + agentSessionId, + baseCommit: 'base-commit', + status: 'available', + degradation: 'none', + skipped: [], + artifacts: { + home: { sizeBytes: 4 }, + wip: { sizeBytes: 3 }, + }, + createdAt: new Date().toISOString(), + }; + + await env.R2.put(prepared.keys.home, 'home'); + await env.R2.put(prepared.keys.wip, 'wip'); + await completeSessionSnapshot(db, bindings, { + workspaceId, + chatSessionId, + agentSessionId, + runtime: 'cf-container', + baseCommit: 'base-commit', + status: 'available', + degradation: 'none', + manifest, + artifactSizes: { + homeBytes: 4, + wipBytes: 3, + }, + }); + + const restorable = await getRestorableSessionSnapshot(db, chatSessionId); + expect(restorable).toMatchObject({ + chatSessionId, + workspaceId, + status: 'available', + degradation: 'none', + homeR2Key: prepared.keys.home, + wipR2Key: prepared.keys.wip, + manifestR2Key: prepared.keys.manifest, + baseCommit: 'base-commit', + }); + + const manifestObject = await env.R2.get(prepared.keys.manifest); + expect(manifestObject).not.toBeNull(); + if (!manifestObject) throw new Error('snapshot manifest was not written to R2'); + await expect(manifestObject.json()).resolves.toMatchObject({ + chatSessionId, + workspaceId, + artifacts: { + home: { sizeBytes: 4 }, + wip: { sizeBytes: 3 }, + }, + }); + }); +}); diff --git a/apps/www/src/content/docs/docs/reference/configuration.md b/apps/www/src/content/docs/docs/reference/configuration.md index f847053d2..8ad846f30 100644 --- a/apps/www/src/content/docs/docs/reference/configuration.md +++ b/apps/www/src/content/docs/docs/reference/configuration.md @@ -73,6 +73,7 @@ GitHub App secrets use `GH_*` prefix (e.g., `GH_CLIENT_ID`, `GH_WEBHOOK_SECRET`) | Variable | Default | Description | | -------------------------------- | ------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `CF_CONTAINER_ENABLED` | `true` | Enables Cloudflare Container instant sessions for matching profiles and zero-config runtime selection. Set `false` to force cloud VM runtime. | +| `CF_CONTAINER_WAKE_TIMEOUT_MS` | `120000` | Maximum time for a sleeping container to launch, restore its snapshot, and accept the triggering request. | | `REQUIRE_APPROVAL` | _(unset)_ | Default signup approval gate. Superadmins can override it at runtime in Admin → Users without redeploying; when no runtime override exists, this value is used. The first genuine human becomes superadmin regardless of this flag — see [First Login & Admin Access](/docs/guides/self-hosting/#first-login--admin-access). | | `TRIAL_ANONYMOUS_USER_ID` | `system_anonymous_trials` | Id of the internal anonymous-trial sentinel user, excluded from first-user superadmin checks. Override only if your deployment uses a different sentinel id. | | `CAPACITY_SIZE_FALLBACK_ENABLED` | `true` | When a new node's VM size is exhausted on transient capacity, descend the size chain (large→medium→small). Only applies to default-derived sizes (project/platform default), never user-requested sizes. Set `false` to disable. | diff --git a/infra/__tests__/config.test.ts b/infra/__tests__/config.test.ts index 33a123a16..70f01fce7 100644 --- a/infra/__tests__/config.test.ts +++ b/infra/__tests__/config.test.ts @@ -46,6 +46,7 @@ describe('infra config parsing', () => { stack: 'staging', r2Location: configModule.DEFAULT_R2_LOCATION, pagesProductionBranch: configModule.DEFAULT_PAGES_PRODUCTION_BRANCH, + sessionSnapshotTtlDays: configModule.DEFAULT_SESSION_SNAPSHOT_TTL_DAYS, }); expect(parsed.prefix).toBe(configModule.derivePrefix('example.com')); }); @@ -56,6 +57,7 @@ describe('infra config parsing', () => { resourcePrefix: 'fork', r2Location: 'WEUR', pagesProductionBranch: 'release/2026-06', + sessionSnapshotTtlDays: '14', }), 'prod' ); @@ -64,6 +66,7 @@ describe('infra config parsing', () => { prefix: 'fork', r2Location: 'WEUR', pagesProductionBranch: 'release/2026-06', + sessionSnapshotTtlDays: 14, }); }); @@ -76,6 +79,8 @@ describe('infra config parsing', () => { ['baseDomain', ' ', 'must not be empty'], ['resourcePrefix', ' ', 'must not be empty'], ['pagesProductionBranch', ' ', 'must not be empty'], + ['sessionSnapshotTtlDays', '0', 'must be a positive integer'], + ['sessionSnapshotTtlDays', '1.5', 'must be a positive integer'], ])('fails fast for invalid %s config', (key, value, expectedMessage) => { const beforeResourceCount = getRegisteredResources().length; diff --git a/infra/__tests__/setup.ts b/infra/__tests__/setup.ts index 7b095d739..2f7c3701f 100644 --- a/infra/__tests__/setup.ts +++ b/infra/__tests__/setup.ts @@ -69,6 +69,17 @@ vi.mock('@pulumi/cloudflare', async (importOriginal) => { super(name, args, opts); } }, + R2BucketLifecycle: class extends actual.R2BucketLifecycle { + constructor(name: string, args: ResourceInputs, opts?: pulumi.CustomResourceOptions) { + resourceRecorder.record( + 'cloudflare:index/r2BucketLifecycle:R2BucketLifecycle', + name, + args, + opts + ); + super(name, args, opts); + } + }, Record: class extends actual.Record { constructor(name: string, args: ResourceInputs, opts?: pulumi.CustomResourceOptions) { resourceRecorder.record('cloudflare:index/record:Record', name, args, opts); diff --git a/infra/__tests__/storage.test.ts b/infra/__tests__/storage.test.ts index 98a17b16c..4b695501b 100644 --- a/infra/__tests__/storage.test.ts +++ b/infra/__tests__/storage.test.ts @@ -27,4 +27,31 @@ describe('R2 Bucket Resource', () => { location: configModule.DEFAULT_R2_LOCATION, }); }); + + it('provisions prefix-scoped session snapshot expiration with the configured TTL', async () => { + const lifecycle = findRegisteredResource( + `${configModule.prefix}-r2-lifecycle`, + 'cloudflare:index/r2BucketLifecycle:R2BucketLifecycle' + ); + + expect(await getOutputValue(lifecycle.inputs.bucketName as never)).toBe( + configModule.prefix + '-' + configModule.stack + '-assets' + ); + expect(lifecycle.inputs).toMatchObject({ + accountId: 'test-account-id-00000000000000000000', + rules: [ + { + id: storageModule.SESSION_SNAPSHOT_LIFECYCLE_RULE_ID, + conditions: { prefix: storageModule.SESSION_SNAPSHOT_R2_PREFIX }, + enabled: true, + deleteObjectsTransition: { + condition: { + maxAge: configModule.DEFAULT_SESSION_SNAPSHOT_TTL_DAYS * 24 * 60 * 60, + type: 'Age', + }, + }, + }, + ], + }); + }); }); diff --git a/infra/__tests__/validate-pulumi-outputs.test.ts b/infra/__tests__/validate-pulumi-outputs.test.ts index 62aad3d14..ee9f469bf 100644 --- a/infra/__tests__/validate-pulumi-outputs.test.ts +++ b/infra/__tests__/validate-pulumi-outputs.test.ts @@ -11,6 +11,7 @@ function makeValidOutputs(): PulumiOutputs { kvId: 'kv-789', kvName: 'sa379a6-prod-sessions', r2Name: 'sa379a6-prod-assets', + sessionSnapshotTtlDays: 7, cloudflareAccountId: 'cf-account-abc', pagesName: 'sa379a6-web-prod', dnsIds: { api: 'dns-1', app: 'dns-2', wildcard: 'dns-3' }, diff --git a/infra/index.ts b/infra/index.ts index 64978e3cc..e61518bab 100644 --- a/infra/index.ts +++ b/infra/index.ts @@ -10,7 +10,7 @@ import { observabilityDatabaseName, } from './resources/database'; import { kvNamespace, kvNamespaceId, kvNamespaceName } from './resources/kv'; -import { r2Bucket, r2BucketName } from './resources/storage'; +import { r2Bucket, r2BucketLifecycle, r2BucketName } from './resources/storage'; import { apiDnsRecord, appDnsRecord, @@ -29,6 +29,7 @@ export { observabilityDatabase, kvNamespace, r2Bucket, + r2BucketLifecycle, pagesProject, pagesCustomDomain, apiDnsRecord, @@ -45,6 +46,7 @@ export const observabilityD1DatabaseName = observabilityDatabaseName; export const kvId = kvNamespaceId; export const kvName = kvNamespaceName; export const r2Name = r2BucketName; +export { sessionSnapshotTtlDays } from './resources/config'; export const pagesName = pagesProjectName; export const dnsIds = dnsRecordIds; export const hostnames = dnsHostnames; diff --git a/infra/resources/config.ts b/infra/resources/config.ts index e552c0a4c..917a0f01e 100644 --- a/infra/resources/config.ts +++ b/infra/resources/config.ts @@ -6,6 +6,7 @@ export const SUPPORTED_R2_LOCATIONS = ['WNAM', 'ENAM', 'WEUR', 'EEUR', 'APAC', ' export type R2Location = (typeof SUPPORTED_R2_LOCATIONS)[number]; export const DEFAULT_PAGES_PRODUCTION_BRANCH = 'main'; +export const DEFAULT_SESSION_SNAPSHOT_TTL_DAYS = 7; export interface ConfigReader { get(key: string): string | undefined; @@ -20,6 +21,7 @@ export interface InfraConfig { prefix: string; r2Location: R2Location; pagesProductionBranch: string; + sessionSnapshotTtlDays: number; } const pulumiConfig = new pulumi.Config(); @@ -37,6 +39,11 @@ export function parseInfraConfig(config: ConfigReader, currentStack: string): In r2Location: parseR2Location(optionalNonEmptyConfig(config, 'r2Location')), pagesProductionBranch: optionalNonEmptyConfig(config, 'pagesProductionBranch') ?? DEFAULT_PAGES_PRODUCTION_BRANCH, + sessionSnapshotTtlDays: parsePositiveInteger( + optionalNonEmptyConfig(config, 'sessionSnapshotTtlDays'), + DEFAULT_SESSION_SNAPSHOT_TTL_DAYS, + 'sessionSnapshotTtlDays' + ), }; } @@ -58,6 +65,7 @@ export const stack = infraConfig.stack; export const prefix = infraConfig.prefix; export const r2Location = infraConfig.r2Location; export const pagesProductionBranch = infraConfig.pagesProductionBranch; +export const sessionSnapshotTtlDays = infraConfig.sessionSnapshotTtlDays; export function derivePrefix(domain: string): string { const hash = crypto.createHash('sha256').update(domain).digest('hex'); @@ -87,6 +95,15 @@ function optionalNonEmptyConfig(config: ConfigReader, key: string): string | und return value; } +function parsePositiveInteger(value: string | undefined, fallback: number, key: string): number { + if (value === undefined) return fallback; + const parsed = Number(value); + if (!Number.isSafeInteger(parsed) || parsed <= 0) { + throw new Error('Pulumi config "' + key + '" must be a positive integer'); + } + return parsed; +} + function parseR2Location(value: string | undefined): R2Location { if (value === undefined) { return DEFAULT_R2_LOCATION; diff --git a/infra/resources/storage.ts b/infra/resources/storage.ts index 3d141dd5f..d49dc93bf 100644 --- a/infra/resources/storage.ts +++ b/infra/resources/storage.ts @@ -1,10 +1,36 @@ import * as cloudflare from '@pulumi/cloudflare'; -import { accountId, prefix, r2Location, stack } from './config'; +import { accountId, prefix, r2Location, sessionSnapshotTtlDays, stack } from './config'; + +export const SESSION_SNAPSHOT_LIFECYCLE_RULE_ID = 'expire-session-snapshots'; +export const SESSION_SNAPSHOT_R2_PREFIX = 'session-snapshots/'; +const SECONDS_PER_DAY = 24 * 60 * 60; export const r2Bucket = new cloudflare.R2Bucket(`${prefix}-r2`, { - accountId: accountId, + accountId, name: `${prefix}-${stack}-assets`, location: r2Location, }); +export const r2BucketLifecycle = new cloudflare.R2BucketLifecycle( + `${prefix}-r2-lifecycle`, + { + accountId, + bucketName: r2Bucket.name, + rules: [ + { + id: SESSION_SNAPSHOT_LIFECYCLE_RULE_ID, + conditions: { prefix: SESSION_SNAPSHOT_R2_PREFIX }, + enabled: true, + deleteObjectsTransition: { + condition: { + maxAge: sessionSnapshotTtlDays * SECONDS_PER_DAY, + type: 'Age', + }, + }, + }, + ], + }, + { dependsOn: r2Bucket } +); + export const r2BucketName = r2Bucket.name; diff --git a/packages/acp-client/tests/unit/components/AgentCrashReportView.test.tsx b/packages/acp-client/tests/unit/components/AgentCrashReportView.test.tsx index 7a926c16e..15b9e5621 100644 --- a/packages/acp-client/tests/unit/components/AgentCrashReportView.test.tsx +++ b/packages/acp-client/tests/unit/components/AgentCrashReportView.test.tsx @@ -53,7 +53,7 @@ describe('AgentCrashReportView', () => { expect(copiedReport).toContain('stderr truncated: yes'); expect(copiedReport).toContain('fatal: peer disconnected before response'); expect(copiedReport).not.toContain('sk-secret'); - expect(screen.getByRole('button', { name: 'Copied' })).not.toBeNull(); + expect(await screen.findByRole('button', { name: 'Copied' })).not.toBeNull(); }); it('shows copy failure feedback when clipboard is unavailable', async () => { diff --git a/packages/vm-agent/internal/server/server.go b/packages/vm-agent/internal/server/server.go index 1e4febf35..6ab3266e9 100644 --- a/packages/vm-agent/internal/server/server.go +++ b/packages/vm-agent/internal/server/server.go @@ -1027,6 +1027,8 @@ func (s *Server) setupRoutes(mux *http.ServeMux) { mux.HandleFunc("POST /workspaces/{workspaceId}/agent-sessions/{sessionId}/suspend", s.handleSuspendAgentSession) mux.HandleFunc("POST /workspaces/{workspaceId}/agent-sessions/{sessionId}/resume", s.handleResumeAgentSession) mux.HandleFunc("POST /workspaces/{workspaceId}/agent-sessions/{sessionId}/prompt", s.handleSendPrompt) + mux.HandleFunc("POST /workspaces/{workspaceId}/agent-sessions/{sessionId}/hibernate", s.handleHibernateAgentSession) + mux.HandleFunc("POST /workspaces/{workspaceId}/agent-sessions/{sessionId}/restore", s.handleRestoreAgentSession) mux.HandleFunc("GET /workspaces/{workspaceId}/tabs", s.handleListTabs) // Git integration (browser-authenticated via workspace session/token) diff --git a/packages/vm-agent/internal/server/session_snapshot.go b/packages/vm-agent/internal/server/session_snapshot.go new file mode 100644 index 000000000..72c9e381e --- /dev/null +++ b/packages/vm-agent/internal/server/session_snapshot.go @@ -0,0 +1,639 @@ +package server + +import ( + "archive/tar" + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "github.com/workspace/vm-agent/internal/acp" +) + +const ( + defaultSnapshotTotalBudgetBytes int64 = 100 * 1024 * 1024 + defaultSnapshotEntryThresholdBytes int64 = 50 * 1024 * 1024 + defaultSnapshotTransferIdleTimeout time.Duration = 30 * time.Second +) + +type snapshotPrepareResponse struct { + ExpiresAt string `json:"expiresAt"` + Config struct { + TotalBudgetBytes int64 `json:"totalBudgetBytes"` + EntryThresholdBytes int64 `json:"entryThresholdBytes"` + TransferIdleTimeoutMs int64 `json:"transferIdleTimeoutMs"` + } `json:"config"` + Upload struct { + Home string `json:"home"` + WIP string `json:"wip"` + } `json:"upload"` +} + +type snapshotRestoreResponse struct { + Available bool `json:"available"` + Reason string `json:"reason,omitempty"` + Status string `json:"status,omitempty"` + Degradation string `json:"degradation,omitempty"` + BaseCommit string `json:"baseCommit,omitempty"` + Manifest map[string]interface{} `json:"manifest,omitempty"` + Config struct { + TransferIdleTimeoutMs int64 `json:"transferIdleTimeoutMs"` + } `json:"config"` + Download struct { + Home string `json:"home"` + WIP string `json:"wip"` + Manifest string `json:"manifest"` + } `json:"download"` +} + +type snapshotManifest struct { + Version int `json:"version"` + ChatSessionID string `json:"chatSessionId"` + WorkspaceID string `json:"workspaceId"` + AgentSessionID string `json:"agentSessionId,omitempty"` + BaseCommit string `json:"baseCommit,omitempty"` + Status string `json:"status"` + Degradation string `json:"degradation"` + Skipped []snapshotSkippedEntry `json:"skipped"` + Artifacts map[string]snapshotArtifact `json:"artifacts"` + CreatedAt string `json:"createdAt"` +} + +type snapshotSkippedEntry struct { + Path string `json:"path"` + Reason string `json:"reason"` + SizeBytes int64 `json:"sizeBytes,omitempty"` +} + +type snapshotArtifact struct { + SizeBytes int64 `json:"sizeBytes"` + SHA256 string `json:"sha256,omitempty"` +} + +type countingReader struct { + r io.Reader + n int64 + hash hashWriter +} + +type hashWriter interface { + Write([]byte) (int, error) + Sum([]byte) []byte +} + +func (r *countingReader) Read(p []byte) (int, error) { + n, err := r.r.Read(p) + if n > 0 { + r.n += int64(n) + _, _ = r.hash.Write(p[:n]) + } + return n, err +} + +type sessionSnapshotHandlerInput struct { + workspaceID string + sessionID string + chatSessionID string + runtimeName string + runtime *WorkspaceRuntime + callbackToken string + agentType string +} + +func (s *Server) sessionSnapshotHandlerInput(w http.ResponseWriter, r *http.Request) (*sessionSnapshotHandlerInput, bool) { + workspaceID := r.PathValue("workspaceId") + sessionID := r.PathValue("sessionId") + if workspaceID == "" || sessionID == "" { + writeError(w, http.StatusBadRequest, "workspaceId and sessionId are required") + return nil, false + } + if !s.requireNodeManagementAuth(w, r, workspaceID) { + return nil, false + } + var body struct { + ChatSessionID string `json:"chatSessionId"` + Runtime string `json:"runtime"` + AgentType string `json:"agentType"` + WorkspaceCallbackToken string `json:"workspaceCallbackToken"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeError(w, http.StatusBadRequest, "invalid request body") + return nil, false + } + body.ChatSessionID = strings.TrimSpace(body.ChatSessionID) + if body.ChatSessionID == "" { + writeError(w, http.StatusBadRequest, "chatSessionId is required") + return nil, false + } + // A freshly-woken container never ran create-workspace, so its + // runtime.CallbackToken (the workspace-scoped token used by the message + // reporter and the snapshot callbacks) is unset. Persist the token the + // control plane provides on the restore request so chat replies and + // snapshot callbacks can authenticate after a wake. + if wsToken := strings.TrimSpace(body.WorkspaceCallbackToken); wsToken != "" { + s.upsertWorkspaceRuntime(workspaceID, "", "", "", wsToken) + } + runtime, ok := s.getWorkspaceRuntime(workspaceID) + if !ok { + writeError(w, http.StatusNotFound, "workspace not found") + return nil, false + } + callbackToken := s.callbackTokenForWorkspace(workspaceID) + if callbackToken == "" { + writeError(w, http.StatusConflict, "workspace callback token unavailable") + return nil, false + } + return &sessionSnapshotHandlerInput{ + workspaceID: workspaceID, + sessionID: sessionID, + chatSessionID: body.ChatSessionID, + runtimeName: body.Runtime, + runtime: runtime, + callbackToken: callbackToken, + agentType: strings.TrimSpace(body.AgentType), + }, true +} + +func (s *Server) handleHibernateAgentSession(w http.ResponseWriter, r *http.Request) { + input, ok := s.sessionSnapshotHandlerInput(w, r) + if !ok { + return + } + result, err := s.hibernateSessionSnapshot(r.Context(), input.runtime, input.sessionID, input.chatSessionID, input.runtimeName, input.callbackToken) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + writeJSON(w, http.StatusOK, result) +} + +func (s *Server) handleRestoreAgentSession(w http.ResponseWriter, r *http.Request) { + input, ok := s.sessionSnapshotHandlerInput(w, r) + if !ok { + return + } + result, err := s.restoreSessionSnapshot(r.Context(), input.runtime, input.sessionID, input.chatSessionID, input.agentType, input.callbackToken) + if err != nil { + _ = s.reportSnapshotRestoreResult(context.Background(), input.workspaceID, input.chatSessionID, "degraded", err.Error(), input.callbackToken) + writeJSON(w, http.StatusOK, map[string]interface{}{"status": "degraded", "message": err.Error()}) + return + } + writeJSON(w, http.StatusOK, result) +} + +func (s *Server) hibernateSessionSnapshot(ctx context.Context, runtime *WorkspaceRuntime, sessionID, chatSessionID, runtimeName, callbackToken string) (map[string]interface{}, error) { + prepare, err := s.prepareSnapshot(ctx, runtime.ID, sessionID, chatSessionID, runtimeName, callbackToken) + if err != nil { + return nil, err + } + totalBudget := choosePositiveInt64(prepare.Config.TotalBudgetBytes, defaultSnapshotTotalBudgetBytes) + entryThreshold := choosePositiveInt64(prepare.Config.EntryThresholdBytes, defaultSnapshotEntryThresholdBytes) + idleTimeout := choosePositiveDurationMs(prepare.Config.TransferIdleTimeoutMs, defaultSnapshotTransferIdleTimeout) + manifest := snapshotManifest{ + Version: 1, + ChatSessionID: chatSessionID, + WorkspaceID: runtime.ID, + AgentSessionID: sessionID, + Status: "available", + Degradation: "none", + Skipped: []snapshotSkippedEntry{}, + Artifacts: map[string]snapshotArtifact{}, + CreatedAt: time.Now().UTC().Format(time.RFC3339), + } + workDir := standaloneWorkspaceWorkDir(runtime, s.config.WorkspaceDir, s.config.ContainerWorkDir) + baseCommit, wipPath, wipSkipped, err := createWIPBundle(ctx, workDir, entryThreshold) + manifest.BaseCommit = baseCommit + manifest.Skipped = append(manifest.Skipped, wipSkipped...) + if err != nil { + manifest.Skipped = append(manifest.Skipped, snapshotSkippedEntry{Path: workDir, Reason: err.Error()}) + } + + remaining := totalBudget + if wipPath != "" { + size, sha, uploadErr := s.uploadSnapshotFile(ctx, prepare.Upload.WIP, wipPath, callbackToken, idleTimeout) + _ = os.Remove(wipPath) + if uploadErr != nil { + manifest.Skipped = append(manifest.Skipped, snapshotSkippedEntry{Path: workDir, Reason: uploadErr.Error()}) + } else { + manifest.Artifacts["wip"] = snapshotArtifact{SizeBytes: size, SHA256: sha} + remaining -= size + } + } + homePath, homeSkipped, err := createHomeTar(os.UserHomeDir, entryThreshold, remaining) + manifest.Skipped = append(manifest.Skipped, homeSkipped...) + if err != nil { + manifest.Skipped = append(manifest.Skipped, snapshotSkippedEntry{Path: "$HOME", Reason: err.Error()}) + } + if homePath != "" { + size, sha, uploadErr := s.uploadSnapshotFile(ctx, prepare.Upload.Home, homePath, callbackToken, idleTimeout) + _ = os.Remove(homePath) + if uploadErr != nil { + manifest.Skipped = append(manifest.Skipped, snapshotSkippedEntry{Path: "$HOME", Reason: uploadErr.Error()}) + } else { + manifest.Artifacts["home"] = snapshotArtifact{SizeBytes: size, SHA256: sha} + } + } + if _, ok := manifest.Artifacts["home"]; !ok { + manifest.Degradation = "wip-only" + } + if _, ok := manifest.Artifacts["wip"]; !ok { + if manifest.Degradation == "wip-only" { + manifest.Degradation = "transcript-only" + manifest.Status = "degraded" + } else { + manifest.Degradation = "home-skipped" + manifest.Status = "degraded" + } + } + if len(manifest.Skipped) > 0 && manifest.Status == "available" { + manifest.Status = "degraded" + } + err = s.completeSnapshot(ctx, runtime.ID, sessionID, chatSessionID, runtimeName, callbackToken, manifest) + if err != nil { + return nil, err + } + return map[string]interface{}{"status": manifest.Status, "degradation": manifest.Degradation, "skipped": manifest.Skipped}, nil +} + +func (s *Server) restoreSessionSnapshot(ctx context.Context, runtime *WorkspaceRuntime, sessionID, chatSessionID, agentType, callbackToken string) (map[string]interface{}, error) { + restore, err := s.fetchSnapshotRestore(ctx, runtime.ID, chatSessionID, callbackToken) + if err != nil { + return nil, err + } + if !restore.Available { + _ = s.reportSnapshotRestoreResult(ctx, runtime.ID, chatSessionID, "missing", restore.Reason, callbackToken) + return map[string]interface{}{"status": "transcript-replay", "reason": restore.Reason}, nil + } + idleTimeout := choosePositiveDurationMs(restore.Config.TransferIdleTimeoutMs, defaultSnapshotTransferIdleTimeout) + if restore.Download.Home != "" { + if err := s.downloadAndExtractTar(ctx, restore.Download.Home, callbackToken, idleTimeout); err != nil { + _ = s.reportSnapshotRestoreResult(ctx, runtime.ID, chatSessionID, "home_failed", err.Error(), callbackToken) + return nil, err + } + } + // A freshly launched runtime has no repository yet. Provision it after HOME + // extraction so current runtime assets and credentials overwrite stale + // snapshot copies, but before applying the Git WIP bundle, which requires a + // materialized repository. + var provisionErr error + if s.config.IsStandaloneMode() { + provisionErr = s.prepareStandaloneWorkspaceRuntime(ctx, runtime) + } else { + _, provisionErr = s.provisionWorkspaceRuntime(ctx, runtime) + } + if provisionErr != nil { + _ = s.reportSnapshotRestoreResult(ctx, runtime.ID, chatSessionID, "fresh_injection_failed", provisionErr.Error(), callbackToken) + return nil, provisionErr + } + if restore.Download.WIP != "" { + workDir := standaloneWorkspaceWorkDir(runtime, s.config.WorkspaceDir, s.config.ContainerWorkDir) + if err := s.downloadAndRestoreWIP(ctx, restore.Download.WIP, callbackToken, idleTimeout, workDir, restore.BaseCommit); err != nil { + _ = s.reportSnapshotRestoreResult(ctx, runtime.ID, chatSessionID, "wip_failed", err.Error(), callbackToken) + return nil, err + } + } + if s.config.IsStandaloneMode() { + if strings.TrimSpace(agentType) == "" { + return nil, fmt.Errorf("agent type is required to restore a standalone session") + } + // Prime the per-workspace message reporter before the agent starts. + // handleCreateAgentSession does this on the normal path; the restore path + // skipped it, so the restored agent's output had no reporter to enqueue + // to and chat replies were silently dropped after a wake. + s.primeRestoredMessageReporter(runtime, chatSessionID) + session, _, createErr := s.agentSessions.Create(runtime.ID, sessionID, "Restored session", "restore:"+sessionID) + if createErr != nil { + return nil, fmt.Errorf("recreate restored agent session: %w", createErr) + } + hostKey := runtime.ID + ":" + sessionID + host := s.getOrCreateSessionHost(hostKey, runtime.ID, sessionID, session, runtime, "") + host.SelectAgent(ctx, agentType) + if host.Status() != acp.HostReady { + return nil, fmt.Errorf("restored agent failed to become ready: %s", host.Status()) + } + } else if session, exists := s.agentSessions.Get(runtime.ID, sessionID); exists { + hostKey := runtime.ID + ":" + sessionID + _ = s.getOrCreateSessionHost(hostKey, runtime.ID, sessionID, session, runtime, "") + } + _ = s.reportSnapshotRestoreResult(ctx, runtime.ID, chatSessionID, "restored", "", callbackToken) + return map[string]interface{}{"status": "restored", "degradation": restore.Degradation}, nil +} + +// primeRestoredMessageReporter ensures the per-workspace message reporter exists +// and is bound to the restored chat session before the agent starts producing +// output. handleCreateAgentSession does this on the normal path; the restore +// path must replicate it or the restored agent's replies are never enqueued and +// are silently dropped after a wake. +func (s *Server) primeRestoredMessageReporter(runtime *WorkspaceRuntime, chatSessionID string) { + if runtime == nil { + return + } + chatSessionID = strings.TrimSpace(chatSessionID) + projectID := strings.TrimSpace(runtime.ProjectID) + if projectID == "" { + projectID = strings.TrimSpace(s.config.ProjectID) + } + if projectID == "" || chatSessionID == "" { + // Without a project + chat session the reporter cannot be created, so + // the restored agent's output would be silently dropped. Log loudly so + // this failure mode is diagnosable instead of invisible. + slog.Warn("Restored session message reporter not primed: missing project or chat session", + "workspaceId", runtime.ID, "hasProjectID", projectID != "", "hasChatSessionID", chatSessionID != "") + return + } + s.workspaceMu.Lock() + if rt, ok := s.workspaces[runtime.ID]; ok && strings.TrimSpace(rt.ProjectID) == "" { + rt.ProjectID = projectID + } + s.workspaceMu.Unlock() + if reporter := s.getOrCreateReporter(runtime.ID, projectID, chatSessionID); reporter != nil { + reporter.SetSessionID(chatSessionID) + } +} + +func (s *Server) prepareSnapshot(ctx context.Context, workspaceID, sessionID, chatSessionID, runtimeName, token string) (*snapshotPrepareResponse, error) { + payload := map[string]string{"chatSessionId": chatSessionID, "agentSessionId": sessionID, "runtime": runtimeName} + var out snapshotPrepareResponse + err := s.doSnapshotJSON(ctx, http.MethodPost, workspaceID, "/session-snapshot/prepare", token, payload, &out) + return &out, err +} + +func (s *Server) completeSnapshot(ctx context.Context, workspaceID, sessionID, chatSessionID, runtimeName, token string, manifest snapshotManifest) error { + artifactSizes := map[string]int64{} + if artifact, ok := manifest.Artifacts["home"]; ok { + artifactSizes["homeBytes"] = artifact.SizeBytes + } + if artifact, ok := manifest.Artifacts["wip"]; ok { + artifactSizes["wipBytes"] = artifact.SizeBytes + } + payload := map[string]interface{}{ + "chatSessionId": chatSessionID, + "agentSessionId": sessionID, + "runtime": runtimeName, + "baseCommit": manifest.BaseCommit, + "status": manifest.Status, + "degradation": manifest.Degradation, + "manifest": manifest, + "artifactSizes": artifactSizes, + } + var out map[string]interface{} + return s.doSnapshotJSON(ctx, http.MethodPost, workspaceID, "/session-snapshot/complete", token, payload, &out) +} + +func (s *Server) fetchSnapshotRestore(ctx context.Context, workspaceID, chatSessionID, token string) (*snapshotRestoreResponse, error) { + path := "/session-snapshot/restore?chatSessionId=" + url.QueryEscape(chatSessionID) + var out snapshotRestoreResponse + err := s.doSnapshotJSON(ctx, http.MethodGet, workspaceID, path, token, nil, &out) + return &out, err +} + +func (s *Server) reportSnapshotRestoreResult(ctx context.Context, workspaceID, chatSessionID, status, message, token string) error { + payload := map[string]string{"chatSessionId": chatSessionID, "status": status, "message": message} + var out map[string]interface{} + return s.doSnapshotJSON(ctx, http.MethodPost, workspaceID, "/session-snapshot/restore-result", token, payload, &out) +} + +func (s *Server) doSnapshotJSON(ctx context.Context, method, workspaceID, path, token string, payload interface{}, out interface{}) error { + endpoint := strings.TrimRight(s.config.ControlPlaneURL, "/") + "/api/workspaces/" + url.PathEscape(workspaceID) + path + var body io.Reader + if payload != nil { + data, err := json.Marshal(payload) + if err != nil { + return err + } + body = bytes.NewReader(data) + } + req, err := http.NewRequestWithContext(ctx, method, endpoint, body) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+token) + if payload != nil { + req.Header.Set("Content-Type", "application/json") + } + res, err := s.controlPlaneHTTPClient(0).Do(req) + if err != nil { + return err + } + defer res.Body.Close() + data, _ := io.ReadAll(io.LimitReader(res.Body, 1024*1024)) + if res.StatusCode < 200 || res.StatusCode >= 300 { + return fmt.Errorf("snapshot control plane returned HTTP %d: %s", res.StatusCode, strings.TrimSpace(string(data))) + } + if out != nil && len(data) > 0 { + if err := json.Unmarshal(data, out); err != nil { + return err + } + } + return nil +} + +func (s *Server) uploadSnapshotFile(ctx context.Context, uploadPath, filePath, token string, idleTimeout time.Duration) (int64, string, error) { + target := absoluteControlPlaneURL(s.config.ControlPlaneURL, uploadPath) + file, err := os.Open(filePath) + if err != nil { + return 0, "", err + } + defer file.Close() + fileInfo, err := file.Stat() + if err != nil { + return 0, "", err + } + h := sha256.New() + reader := &countingReader{r: file, hash: h} + req, err := http.NewRequestWithContext(ctx, http.MethodPut, target, newIdleReader(reader, idleTimeout)) + if err != nil { + return 0, "", err + } + req.Header.Set("Authorization", "Bearer "+token) + req.ContentLength = fileInfo.Size() + res, err := s.controlPlaneHTTPClient(0).Do(req) + if err != nil { + return 0, "", err + } + defer res.Body.Close() + body, _ := io.ReadAll(io.LimitReader(res.Body, 64*1024)) + if res.StatusCode < 200 || res.StatusCode >= 300 { + return 0, "", fmt.Errorf("artifact upload failed HTTP %d: %s", res.StatusCode, strings.TrimSpace(string(body))) + } + return reader.n, hex.EncodeToString(h.Sum(nil)), nil +} + +func (s *Server) downloadAndExtractTar(ctx context.Context, downloadPath, token string, idleTimeout time.Duration) error { + res, err := s.snapshotDownload(ctx, downloadPath, token) + if err != nil { + return err + } + defer res.Body.Close() + home, err := os.UserHomeDir() + if err != nil { + return err + } + home = filepath.Clean(home) + tr := tar.NewReader(newIdleReader(res.Body, idleTimeout)) + for { + header, err := tr.Next() + if err == io.EOF { + return nil + } + if err != nil { + return err + } + cleanName := filepath.Clean(header.Name) + if filepath.IsAbs(cleanName) || cleanName == "." || cleanName == ".." || strings.HasPrefix(cleanName, ".."+string(filepath.Separator)) { + continue + } + target := filepath.Join(home, cleanName) + relTarget, relErr := filepath.Rel(home, target) + if relErr != nil || relTarget == ".." || strings.HasPrefix(relTarget, ".."+string(filepath.Separator)) { + continue + } + if err := rejectSymlinkPath(home, target); err != nil { + return err + } + if header.Typeflag != tar.TypeDir && header.Typeflag != tar.TypeReg && header.Typeflag != tar.TypeRegA { + continue + } + if header.FileInfo().IsDir() { + if err := os.MkdirAll(target, header.FileInfo().Mode()); err != nil { + return err + } + continue + } + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return err + } + f, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, header.FileInfo().Mode()) + if err != nil { + return err + } + _, copyErr := io.Copy(f, tr) + closeErr := f.Close() + if copyErr != nil { + return copyErr + } + if closeErr != nil { + return closeErr + } + } +} + +func (s *Server) downloadAndRestoreWIP(ctx context.Context, downloadPath, token string, idleTimeout time.Duration, workDir, baseCommit string) error { + res, err := s.snapshotDownload(ctx, downloadPath, token) + if err != nil { + return err + } + defer res.Body.Close() + tmp, err := os.CreateTemp("", "sam-session-restore-*.bundle") + if err != nil { + return err + } + tmpPath := tmp.Name() + _, copyErr := io.Copy(tmp, newIdleReader(res.Body, idleTimeout)) + closeErr := tmp.Close() + if copyErr != nil { + _ = os.Remove(tmpPath) + return copyErr + } + if closeErr != nil { + _ = os.Remove(tmpPath) + return closeErr + } + defer os.Remove(tmpPath) + heads, err := runStandaloneGitCommand(ctx, workDir, nil, "bundle", "list-heads", tmpPath) + if err != nil { + return fmt.Errorf("list snapshot bundle heads: %w: %s", err, heads) + } + fields := strings.Fields(heads) + if len(fields) < 2 { + return fmt.Errorf("snapshot bundle has no restorable ref") + } + if output, err := runStandaloneGitCommand(ctx, workDir, nil, "fetch", tmpPath, fields[1]); err != nil { + return fmt.Errorf("fetch snapshot bundle: %w: %s", err, output) + } + if output, err := runStandaloneGitCommand(ctx, workDir, nil, "read-tree", "--reset", "-u", "FETCH_HEAD"); err != nil { + return fmt.Errorf("materialize snapshot tree: %w: %s", err, output) + } + if strings.TrimSpace(baseCommit) != "" { + if output, err := runStandaloneGitCommand(ctx, workDir, nil, "reset", "--mixed", baseCommit); err != nil { + return fmt.Errorf("restore snapshot base commit: %w: %s", err, output) + } + } + return nil +} + +func (s *Server) snapshotDownload(ctx context.Context, downloadPath, token string) (*http.Response, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, absoluteControlPlaneURL(s.config.ControlPlaneURL, downloadPath), nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+token) + res, err := s.controlPlaneHTTPClient(0).Do(req) + if err != nil { + return nil, err + } + if res.StatusCode < 200 || res.StatusCode >= 300 { + body, _ := io.ReadAll(io.LimitReader(res.Body, 64*1024)) + _ = res.Body.Close() + return nil, fmt.Errorf("artifact download failed HTTP %d: %s", res.StatusCode, strings.TrimSpace(string(body))) + } + return res, nil +} + +func absoluteControlPlaneURL(base, path string) string { + if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") { + return path + } + return strings.TrimRight(base, "/") + path +} + +type idleReader struct { + reader io.Reader + idleTimeout time.Duration +} + +func newIdleReader(reader io.Reader, idleTimeout time.Duration) io.Reader { + return &idleReader{reader: reader, idleTimeout: idleTimeout} +} + +func (r *idleReader) Read(p []byte) (int, error) { + type result struct { + n int + err error + } + ch := make(chan result, 1) + go func() { + n, err := r.reader.Read(p) + ch <- result{n: n, err: err} + }() + select { + case res := <-ch: + return res.n, res.err + case <-time.After(r.idleTimeout): + return 0, fmt.Errorf("snapshot transfer stalled for %s", r.idleTimeout) + } +} + +func choosePositiveInt64(value, fallback int64) int64 { + if value > 0 { + return value + } + return fallback +} + +func choosePositiveDurationMs(value int64, fallback time.Duration) time.Duration { + if value > 0 { + return time.Duration(value) * time.Millisecond + } + return fallback +} diff --git a/packages/vm-agent/internal/server/session_snapshot_archive.go b/packages/vm-agent/internal/server/session_snapshot_archive.go new file mode 100644 index 000000000..ccf1763e1 --- /dev/null +++ b/packages/vm-agent/internal/server/session_snapshot_archive.go @@ -0,0 +1,233 @@ +package server + +import ( + "archive/tar" + "context" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +func createWIPBundle(ctx context.Context, workDir string, entryThreshold int64) (string, string, []snapshotSkippedEntry, error) { + if ok, err := standaloneRepositoryPresent(workDir); err != nil || !ok { + if err != nil { + return "", "", nil, err + } + return "", "", nil, nil + } + if gitOperationInProgress(workDir) { + return "", "", []snapshotSkippedEntry{{Path: workDir, Reason: "git operation in progress"}}, nil + } + base, err := runStandaloneGitCommand(ctx, workDir, nil, "rev-parse", "HEAD") + if err != nil { + return "", "", nil, fmt.Errorf("resolve base commit: %w", err) + } + status, err := runStandaloneGitCommand(ctx, workDir, nil, "status", "--porcelain") + if err != nil { + return base, "", nil, fmt.Errorf("git status: %w", err) + } + if strings.TrimSpace(status) == "" { + return base, "", nil, nil + } + + indexFile, err := os.CreateTemp("", "sam-session-index-*") + if err != nil { + return base, "", nil, err + } + indexPath := indexFile.Name() + _ = indexFile.Close() + _ = os.Remove(indexPath) + defer os.Remove(indexPath) + gitEnv := []string{"GIT_INDEX_FILE=" + indexPath} + if _, err := runStandaloneGitCommand(ctx, workDir, gitEnv, "read-tree", "HEAD"); err != nil { + return base, "", nil, fmt.Errorf("initialize snapshot index: %w", err) + } + if _, err := runStandaloneGitCommand(ctx, workDir, gitEnv, "add", "-A"); err != nil { + return base, "", nil, fmt.Errorf("stage snapshot index: %w", err) + } + skipped := skipOversizedUntracked(workDir, entryThreshold) + for _, entry := range skipped { + if entry.Path != "" { + _, _ = runStandaloneGitCommand(ctx, workDir, gitEnv, "reset", "--", entry.Path) + } + } + tree, err := runStandaloneGitCommand(ctx, workDir, gitEnv, "write-tree") + if err != nil { + return base, "", skipped, fmt.Errorf("write snapshot tree: %w", err) + } + commitEnv := append(gitEnv, "GIT_AUTHOR_NAME=SAM Snapshot", "GIT_AUTHOR_EMAIL=snapshot@localhost", "GIT_COMMITTER_NAME=SAM Snapshot", "GIT_COMMITTER_EMAIL=snapshot@localhost") + commit, err := runStandaloneGitCommand(ctx, workDir, commitEnv, "commit-tree", tree, "-p", base, "-m", "SAM session snapshot") + if err != nil { + return base, "", skipped, fmt.Errorf("create snapshot commit: %w", err) + } + bundle, err := os.CreateTemp("", "sam-session-wip-*.bundle") + if err != nil { + return base, "", skipped, err + } + bundlePath := bundle.Name() + _ = bundle.Close() + snapshotRef := "refs/sam/session-snapshot/" + strings.TrimSuffix(filepath.Base(bundlePath), ".bundle") + if _, err := runStandaloneGitCommand(ctx, workDir, nil, "update-ref", snapshotRef, commit); err != nil { + _ = os.Remove(bundlePath) + return base, "", skipped, fmt.Errorf("create snapshot ref: %w", err) + } + defer func() { + _, _ = runStandaloneGitCommand(context.Background(), workDir, nil, "update-ref", "-d", snapshotRef) + }() + if _, err := runStandaloneGitCommand(ctx, workDir, nil, "bundle", "create", bundlePath, snapshotRef); err != nil { + _ = os.Remove(bundlePath) + return base, "", skipped, fmt.Errorf("create git bundle: %w", err) + } + return base, bundlePath, skipped, nil +} + +func gitOperationInProgress(workDir string) bool { + gitDir := filepath.Join(workDir, ".git") + for _, marker := range []string{"MERGE_HEAD", "CHERRY_PICK_HEAD", "REVERT_HEAD", "rebase-merge", "rebase-apply"} { + if _, err := os.Stat(filepath.Join(gitDir, marker)); err == nil { + return true + } + } + return false +} + +func skipOversizedUntracked(workDir string, threshold int64) []snapshotSkippedEntry { + var skipped []snapshotSkippedEntry + _ = filepath.WalkDir(workDir, func(path string, d os.DirEntry, err error) error { + if err != nil || path == workDir { + return nil + } + if d.IsDir() && d.Name() == ".git" { + return filepath.SkipDir + } + if d.IsDir() { + return nil + } + info, statErr := d.Info() + if statErr != nil || info.Size() <= threshold { + return nil + } + rel, _ := filepath.Rel(workDir, path) + if out, gitErr := runStandaloneGitCommand(context.Background(), workDir, nil, "check-ignore", "-q", rel); gitErr == nil && strings.TrimSpace(out) == "" { + return nil + } + skipped = append(skipped, snapshotSkippedEntry{Path: rel, Reason: "entry exceeds size threshold", SizeBytes: info.Size()}) + return nil + }) + return skipped +} + +func createHomeTar(homeDirFn func() (string, error), entryThreshold, totalBudget int64) (string, []snapshotSkippedEntry, error) { + home, err := homeDirFn() + if err != nil { + return "", nil, err + } + home = filepath.Clean(home) + out, err := os.CreateTemp("", "sam-session-home-*.tar") + if err != nil { + return "", nil, err + } + path := out.Name() + tw := tar.NewWriter(out) + var written int64 + var skipped []snapshotSkippedEntry + walkErr := filepath.WalkDir(home, func(path string, d os.DirEntry, err error) error { + if err != nil || path == home { + return nil + } + rel, _ := filepath.Rel(home, path) + if shouldExcludeHomePath(rel) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + info, statErr := d.Info() + if statErr != nil { + return nil + } + if info.Size() > entryThreshold { + skipped = append(skipped, snapshotSkippedEntry{Path: "~/" + rel, Reason: "entry exceeds size threshold", SizeBytes: info.Size()}) + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if !info.Mode().IsRegular() && !info.IsDir() { + skipped = append(skipped, snapshotSkippedEntry{Path: "~/" + rel, Reason: "unsupported home entry type"}) + return nil + } + if !info.IsDir() && written+info.Size() > totalBudget { + skipped = append(skipped, snapshotSkippedEntry{Path: "~/" + rel, Reason: "snapshot budget exhausted", SizeBytes: info.Size()}) + return nil + } + header, headerErr := tar.FileInfoHeader(info, "") + if headerErr != nil { + return nil + } + header.Name = rel + if err := tw.WriteHeader(header); err != nil { + return err + } + if info.Mode().IsRegular() { + f, openErr := os.Open(path) + if openErr != nil { + return nil + } + n, copyErr := io.Copy(tw, f) + _ = f.Close() + written += n + if copyErr != nil { + return copyErr + } + } + return nil + }) + closeErr := tw.Close() + fileCloseErr := out.Close() + if walkErr != nil || closeErr != nil || fileCloseErr != nil { + _ = os.Remove(path) + if walkErr != nil { + return "", skipped, walkErr + } + if closeErr != nil { + return "", skipped, closeErr + } + return "", skipped, fileCloseErr + } + return path, skipped, nil +} + +func shouldExcludeHomePath(rel string) bool { + first := strings.Split(filepath.ToSlash(rel), "/")[0] + switch first { + case ".cache", ".npm", ".cargo", ".rustup", ".local", "node_modules", ".docker": + return true + default: + return false + } +} + +func rejectSymlinkPath(root, target string) error { + rel, err := filepath.Rel(root, target) + if err != nil || rel == "." || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return fmt.Errorf("snapshot target is outside home") + } + current := root + for _, part := range strings.Split(rel, string(filepath.Separator)) { + current = filepath.Join(current, part) + info, statErr := os.Lstat(current) + if os.IsNotExist(statErr) { + continue + } + if statErr != nil { + return statErr + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("snapshot target traverses symlink: %s", rel) + } + } + return nil +} diff --git a/packages/vm-agent/internal/server/session_snapshot_git_test.go b/packages/vm-agent/internal/server/session_snapshot_git_test.go new file mode 100644 index 000000000..70b5dd400 --- /dev/null +++ b/packages/vm-agent/internal/server/session_snapshot_git_test.go @@ -0,0 +1,141 @@ +package server + +import ( + "archive/tar" + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/workspace/vm-agent/internal/config" +) + +func TestCreateWIPBundlePreservesBranchAndIndex(t *testing.T) { + repo := initSnapshotTestRepo(t) + if err := os.WriteFile(filepath.Join(repo, "README.md"), []byte("staged"), 0o600); err != nil { + t.Fatal(err) + } + runGit(t, repo, "add", "README.md") + if err := os.WriteFile(filepath.Join(repo, "README.md"), []byte("staged and unstaged"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(repo, "untracked.txt"), []byte("untracked"), 0o600); err != nil { + t.Fatal(err) + } + beforeStatus := gitOutput(t, repo, "status", "--porcelain=v1") + beforeHead := gitOutput(t, repo, "rev-parse", "HEAD") + beforeBranch := gitOutput(t, repo, "branch", "--show-current") + + _, bundlePath, _, err := createWIPBundle(context.Background(), repo, 1024) + if err != nil { + t.Fatal(err) + } + defer os.Remove(bundlePath) + + if got := gitOutput(t, repo, "status", "--porcelain=v1"); got != beforeStatus { + t.Fatalf("status changed by snapshot:\nwant %q\n got %q", beforeStatus, got) + } + if got := gitOutput(t, repo, "rev-parse", "HEAD"); got != beforeHead { + t.Fatalf("HEAD changed by snapshot: want %s, got %s", beforeHead, got) + } + if got := gitOutput(t, repo, "branch", "--show-current"); got != beforeBranch { + t.Fatalf("branch changed by snapshot: want %s, got %s", beforeBranch, got) + } +} + +func TestDownloadAndRestoreWIPKeepsOriginalBranch(t *testing.T) { + repo := initSnapshotTestRepo(t) + base := gitOutput(t, repo, "rev-parse", "HEAD") + branch := gitOutput(t, repo, "branch", "--show-current") + if err := os.WriteFile(filepath.Join(repo, "README.md"), []byte("restored change"), 0o600); err != nil { + t.Fatal(err) + } + _, bundlePath, _, err := createWIPBundle(context.Background(), repo, 1024) + if err != nil { + t.Fatal(err) + } + defer os.Remove(bundlePath) + runGit(t, repo, "reset", "--hard", base) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + f, openErr := os.Open(bundlePath) + if openErr != nil { + http.Error(w, openErr.Error(), http.StatusInternalServerError) + return + } + defer f.Close() + _, _ = io.Copy(w, f) + })) + defer server.Close() + s := &Server{config: &config.Config{ControlPlaneURL: server.URL}} + if err := s.downloadAndRestoreWIP(context.Background(), server.URL, "token", time.Second, repo, base); err != nil { + t.Fatal(err) + } + if got := gitOutput(t, repo, "branch", "--show-current"); got != branch { + t.Fatalf("branch changed by restore: want %s, got %s", branch, got) + } + if got := strings.TrimSpace(gitOutput(t, repo, "diff", "--", "README.md")); got == "" { + t.Fatal("restored WIP is missing") + } +} + +func initSnapshotTestRepo(t *testing.T) string { + t.Helper() + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git not available") + } + repo := t.TempDir() + runGit(t, repo, "init") + runGit(t, repo, "config", "user.email", "sam@example.test") + runGit(t, repo, "config", "user.name", "SAM") + if err := os.WriteFile(filepath.Join(repo, "README.md"), []byte("base"), 0o600); err != nil { + t.Fatal(err) + } + runGit(t, repo, "add", "README.md") + runGit(t, repo, "commit", "-m", "base") + return repo +} + +func gitOutput(t *testing.T, dir string, args ...string) string { + t.Helper() + cmd := exec.Command("git", args...) + cmd.Dir = dir + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("git %v failed: %v\n%s", args, err, out) + } + return strings.TrimSpace(string(out)) +} + +func TestDownloadAndExtractTarRejectsExistingHomeSymlink(t *testing.T) { + home := t.TempDir() + outside := t.TempDir() + t.Setenv("HOME", home) + if err := os.Symlink(outside, filepath.Join(home, "linked")); err != nil { + t.Fatal(err) + } + var tarBody bytes.Buffer + tw := tar.NewWriter(&tarBody) + writeTarFile(t, tw, "linked/credential", "secret") + if err := tw.Close(); err != nil { + t.Fatal(err) + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write(tarBody.Bytes()) + })) + defer server.Close() + s := &Server{config: &config.Config{ControlPlaneURL: server.URL}} + if err := s.downloadAndExtractTar(context.Background(), server.URL, "token", time.Second); err == nil || !strings.Contains(err.Error(), "symlink") { + t.Fatalf("error = %v, want symlink rejection", err) + } + if _, err := os.Stat(filepath.Join(outside, "credential")); !os.IsNotExist(err) { + t.Fatalf("outside credential stat err = %v, want not exist", err) + } +} diff --git a/packages/vm-agent/internal/server/session_snapshot_reporter_test.go b/packages/vm-agent/internal/server/session_snapshot_reporter_test.go new file mode 100644 index 000000000..14aaaba1f --- /dev/null +++ b/packages/vm-agent/internal/server/session_snapshot_reporter_test.go @@ -0,0 +1,75 @@ +package server + +import "testing" + +// TestUpsertWorkspaceRuntime_PersistsWorkspaceCallbackToken covers the fix that +// lets a freshly-woken container authenticate its message + snapshot callbacks: +// the restore handler persists the workspace-scoped token from the request body +// via upsertWorkspaceRuntime, and it must then be returned by both +// callbackTokenForWorkspace (with node fallback) and workspaceCallbackToken +// (no fallback — used by the message reporter). +func TestUpsertWorkspaceRuntime_PersistsWorkspaceCallbackToken(t *testing.T) { + s, _ := newServerWithoutReporter(t) + s.config.CallbackToken = "node-scoped-token" + ws := "ws-token-persist-1" + + // A fresh wake container starts with no workspace-scoped token; the snapshot + // handler falls back to the node-scoped CALLBACK_TOKEN, and the no-fallback + // workspaceCallbackToken is empty (which is why message replies would drop). + s.upsertWorkspaceRuntime(ws, "org/repo", "main", "running", "") + if got := s.callbackTokenForWorkspace(ws); got != "node-scoped-token" { + t.Fatalf("callbackTokenForWorkspace before restore: expected node fallback, got %q", got) + } + if got := s.workspaceCallbackToken(ws); got != "" { + t.Fatalf("workspaceCallbackToken before restore: expected empty, got %q", got) + } + + // The restore handler persists the workspace-scoped token from the body. + s.upsertWorkspaceRuntime(ws, "", "", "", "workspace-scoped-token") + + if got := s.callbackTokenForWorkspace(ws); got != "workspace-scoped-token" { + t.Fatalf("callbackTokenForWorkspace after restore: expected workspace token, got %q", got) + } + if got := s.workspaceCallbackToken(ws); got != "workspace-scoped-token" { + t.Fatalf("workspaceCallbackToken after restore: expected workspace token, got %q", got) + } +} + +// TestPrimeRestoredMessageReporter_CreatesReporter covers the fix that makes a +// restored agent's replies actually persist: the restore path must create/bind +// the per-workspace message reporter (the normal create-agent-session path does +// this; restore previously skipped it and replies were silently dropped). +func TestPrimeRestoredMessageReporter_CreatesReporter(t *testing.T) { + s, _ := newServerWithoutReporter(t) + s.config.ProjectID = "proj-restore-1" // set from PROJECT_ID env at launch + ws := "ws-prime-1" + rt := s.upsertWorkspaceRuntime(ws, "org/repo", "main", "running", "ws-tok") + + if len(s.messageReporters) != 0 { + t.Fatalf("expected no reporter before priming, got %d", len(s.messageReporters)) + } + + s.primeRestoredMessageReporter(rt, "chat-sess-1") + + if _, ok := s.messageReporters[ws]; !ok { + t.Fatal("expected message reporter to be primed after restore") + } + if rt.ProjectID != "proj-restore-1" { + t.Fatalf("expected runtime.ProjectID to be set from config fallback, got %q", rt.ProjectID) + } +} + +// TestPrimeRestoredMessageReporter_NoOpWithoutProject verifies the guard: with +// no project context available (neither runtime nor config), the reporter is +// not created (and a warning is logged) rather than panicking. +func TestPrimeRestoredMessageReporter_NoOpWithoutProject(t *testing.T) { + s, _ := newServerWithoutReporter(t) // config.ProjectID is empty + ws := "ws-prime-2" + rt := s.upsertWorkspaceRuntime(ws, "org/repo", "main", "running", "ws-tok") + + s.primeRestoredMessageReporter(rt, "chat-sess-2") + + if _, ok := s.messageReporters[ws]; ok { + t.Fatal("expected no reporter when project context is unavailable") + } +} diff --git a/packages/vm-agent/internal/server/session_snapshot_test.go b/packages/vm-agent/internal/server/session_snapshot_test.go new file mode 100644 index 000000000..63fcefd5d --- /dev/null +++ b/packages/vm-agent/internal/server/session_snapshot_test.go @@ -0,0 +1,175 @@ +package server + +import ( + "archive/tar" + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/workspace/vm-agent/internal/config" +) + +func TestCreateHomeTarExcludesCachesAndRecordsOversizedFiles(t *testing.T) { + home := t.TempDir() + if err := os.WriteFile(filepath.Join(home, ".codex-session.jsonl"), []byte("session"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(filepath.Join(home, ".cache"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(home, ".cache", "ignored"), []byte("cache"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(home, "large.bin"), []byte("0123456789"), 0o600); err != nil { + t.Fatal(err) + } + + tarPath, skipped, err := createHomeTar(func() (string, error) { return home, nil }, 8, 1024) + if err != nil { + t.Fatal(err) + } + defer os.Remove(tarPath) + + if len(skipped) != 1 { + t.Fatalf("skipped len = %d, want 1: %#v", len(skipped), skipped) + } + if skipped[0].Path != "~/large.bin" { + t.Fatalf("skipped path = %q, want ~/large.bin", skipped[0].Path) + } + + names := tarNames(t, tarPath) + if !containsString(names, ".codex-session.jsonl") { + t.Fatalf("tar names missing session file: %#v", names) + } + if containsString(names, ".cache/ignored") { + t.Fatalf("tar names included cache file: %#v", names) + } + if containsString(names, "large.bin") { + t.Fatalf("tar names included oversized file: %#v", names) + } +} + +func TestCreateWIPBundleDegradesDuringMerge(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git not available") + } + repo := t.TempDir() + runGit(t, repo, "init") + runGit(t, repo, "config", "user.email", "sam@example.test") + runGit(t, repo, "config", "user.name", "SAM") + if err := os.WriteFile(filepath.Join(repo, "README.md"), []byte("base"), 0o600); err != nil { + t.Fatal(err) + } + runGit(t, repo, "add", "README.md") + runGit(t, repo, "commit", "-m", "base") + if err := os.WriteFile(filepath.Join(repo, ".git", "MERGE_HEAD"), []byte("deadbeef"), 0o600); err != nil { + t.Fatal(err) + } + + _, bundlePath, skipped, err := createWIPBundle(context.Background(), repo, 1024) + if err != nil { + t.Fatal(err) + } + if bundlePath != "" { + t.Fatalf("bundlePath = %q, want empty during merge", bundlePath) + } + if len(skipped) != 1 || skipped[0].Reason != "git operation in progress" { + t.Fatalf("skipped = %#v, want git operation degradation", skipped) + } +} + +func TestDownloadAndExtractTarRejectsPathTraversal(t *testing.T) { + home := t.TempDir() + outside := filepath.Join(t.TempDir(), "outside.txt") + absolute := filepath.Join(t.TempDir(), "absolute.txt") + t.Setenv("HOME", home) + + var tarBody bytes.Buffer + tw := tar.NewWriter(&tarBody) + writeTarFile(t, tw, "session.jsonl", "inside") + writeTarFile(t, tw, "../outside.txt", "outside") + writeTarFile(t, tw, absolute, "absolute") + if err := tw.Close(); err != nil { + t.Fatal(err) + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(tarBody.Bytes()) + })) + defer server.Close() + + s := &Server{config: &config.Config{ControlPlaneURL: server.URL}} + if err := s.downloadAndExtractTar(context.Background(), "/artifact", "token", time.Second); err != nil { + t.Fatal(err) + } + + if got, err := os.ReadFile(filepath.Join(home, "session.jsonl")); err != nil || string(got) != "inside" { + t.Fatalf("inside file = %q, %v; want inside", got, err) + } + if _, err := os.Stat(outside); !os.IsNotExist(err) { + t.Fatalf("outside file stat err = %v, want not exist", err) + } + if _, err := os.Stat(absolute); !os.IsNotExist(err) { + t.Fatalf("absolute file stat err = %v, want not exist", err) + } +} + +func tarNames(t *testing.T, path string) []string { + t.Helper() + f, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + tr := tar.NewReader(f) + var names []string + for { + header, err := tr.Next() + if err == io.EOF { + return names + } + if err != nil { + t.Fatal(err) + } + names = append(names, header.Name) + } +} + +func runGit(t *testing.T, dir string, args ...string) { + t.Helper() + cmd := exec.Command("git", args...) + cmd.Dir = dir + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("git %v failed: %v\n%s", args, err, out) + } +} + +func writeTarFile(t *testing.T, tw *tar.Writer, name, body string) { + t.Helper() + if err := tw.WriteHeader(&tar.Header{ + Name: name, + Mode: 0o600, + Size: int64(len(body)), + }); err != nil { + t.Fatal(err) + } + if _, err := tw.Write([]byte(body)); err != nil { + t.Fatal(err) + } +} + +func containsString(values []string, needle string) bool { + for _, value := range values { + if value == needle { + return true + } + } + return false +} diff --git a/scripts/deploy/sync-wrangler-config.ts b/scripts/deploy/sync-wrangler-config.ts index 7cb111e7f..e2074d5bb 100644 --- a/scripts/deploy/sync-wrangler-config.ts +++ b/scripts/deploy/sync-wrangler-config.ts @@ -116,12 +116,16 @@ export function validatePulumiOutputs(outputs: unknown): asserts outputs is Pulu { key: 'observabilityD1DatabaseName', label: 'Observability D1 Database Name' }, { key: 'kvId', label: 'KV Namespace ID' }, { key: 'r2Name', label: 'R2 Bucket Name' }, + { key: 'sessionSnapshotTtlDays', label: 'Session Snapshot TTL Days' }, { key: 'cloudflareAccountId', label: 'Cloudflare Account ID' }, { key: 'pagesName', label: 'Pages Project Name' }, ]; const missing = required.filter(({ key }) => { const value = record[key]; + if (key === 'sessionSnapshotTtlDays') { + return typeof value !== 'number' || !Number.isSafeInteger(value) || value <= 0; + } return typeof value !== 'string' || value.length === 0; }); @@ -333,11 +337,14 @@ function getApiWorkerVars( VERSION: DEPLOYMENT_CONFIG.version, PAGES_PROJECT_NAME: outputs.pagesName, R2_BUCKET_NAME: outputs.r2Name, + SESSION_SNAPSHOT_TTL_DAYS: String(outputs.sessionSnapshotTtlDays), ...getOptionalProcessEnvVars([ 'REQUIRE_APPROVAL', 'HETZNER_BASE_IMAGE', 'CF_CONTAINER_ENABLED', + 'CF_CONTAINER_SLEEP_AFTER', 'CF_CONTAINER_PORT_READY_TIMEOUT_MS', + 'CF_CONTAINER_WAKE_TIMEOUT_MS', 'CF_CONTAINER_VM_AGENT_PORT', 'SANDBOX_ENABLED', 'SANDBOX_EXEC_TIMEOUT_MS', diff --git a/scripts/deploy/types.ts b/scripts/deploy/types.ts index e13d912a1..c713c9949 100644 --- a/scripts/deploy/types.ts +++ b/scripts/deploy/types.ts @@ -348,6 +348,7 @@ export interface PulumiOutputs { kvId: string; kvName: string; r2Name: string; + sessionSnapshotTtlDays: number; dnsIds: { api: string; app: string; diff --git a/tasks/archive/2026-07-11-runtime-neutral-session-hibernate-wake.md b/tasks/archive/2026-07-11-runtime-neutral-session-hibernate-wake.md new file mode 100644 index 000000000..132aa10e6 --- /dev/null +++ b/tasks/archive/2026-07-11-runtime-neutral-session-hibernate-wake.md @@ -0,0 +1,84 @@ +# Runtime-neutral session hibernate/wake for cf-container workspaces + +## Problem + +Cloudflare Container instant workspaces lose their local disk after sleep. That currently means harness-native session state under `$HOME` and uncommitted repository work disappear when a sleeping cf-container workspace wakes. Phase 3a of idea `01KX4KSXEXQMP41KS34TW9EN01` must add a runtime-neutral hibernate/restore contract and wire the first consumer to the cf-container runtime. + +Updated authorization from Raphaël, 2026-07-11: + +- The earlier draft-only and no-staging hold is superseded. +- Complete the full `/do` workflow on existing PR #1562. +- Staging end-to-end verification is required before merge. +- Merge only after every required gate passes, then monitor production deployment. + +## Research Findings + +- `apps/api/src/durable-objects/vm-agent-container.ts` owns cf-container launch, active-work keepalive, `sleeping` lifecycle, and currently returns `503` for sleeping containers. +- `apps/api/src/services/vm-agent-container.ts` wraps the DO and already filters active-work calls to `runtime = 'cf-container'`. +- `apps/api/src/services/node-agent.ts` dispatches initial and follow-up prompts, marking cf-container active work started/ended. +- `apps/api/src/routes/projects/agent-activity-callback.ts` receives VM-agent activity with callback JWT auth and ends active work when the harness reports idle/error. +- `packages/vm-agent/internal/server/workspaces.go` creates `SessionHost` instances. Follow-up prompt handling currently requires an in-memory host, so a fresh container after sleep needs a rebuild path. +- `packages/vm-agent/internal/server/standalone_workspace.go` prepares raw cf-container standalone workspaces by cloning the repo. Restore must happen after clone but before fresh runtime-asset injection/harness start. +- R2 access is via Worker binding for metadata/object ownership. VM-agent needs control-plane endpoints or signed URLs for artifact transfer; do not push WIP refs to the user's remote. +- Relevant rules: migration safety, vertical-slice testing, VM-agent callback auth scoping, long-running progress/idle watchdogs, no hardcoded values, DO concurrency, and credential snapshot resilience. + +## Implementation Checklist + +- [x] Add additive D1 schema/migration for runtime-neutral session snapshots: canonical chat session key, R2 keys, manifest/degradation, expiry, and restore diagnostics. +- [x] Add env-configurable defaults for snapshot TTL, size budgets, per-entry thresholds, progress idle watchdogs, and R2 key prefix. +- [x] Add Worker snapshot service for deterministic one-snapshot-per-session R2 keys, metadata persistence, retention checks, cleanup hooks, and visible degradation state. +- [x] Add VM-agent callback/API endpoints with callback JWT auth for snapshot upload/download coordination, mounted outside browser-session middleware. +- [x] Add VM-agent hibernate implementation: git WIP bundle, `$HOME` tar with generic cache exclusions, size/skip manifest, progress idle watchdogs, and no remote push. +- [x] Add VM-agent restore implementation: restore `$HOME`, restore WIP bundle with soft reset to base, re-run runtime injection afterward, and report visible degraded results. +- [x] Wire cf-container hibernate on control handback before sleep is allowed. +- [x] Wire cf-container wake before proxying a sleeping container and before follow-up prompt dispatch; fall back visibly when no snapshot exists or restore fails. +- [x] Update follow-up prompt handling so a fresh vm-agent can rebuild a `SessionHost` and attempt native harness resume from restored `$HOME`. +- [x] Add local tests: Worker service/unit tests, Miniflare/DO integration tests, VM-agent Go tests including `go test -race`, and a vertical-slice test covering Worker → DO → vm-agent → R2 state. +- [x] Run migration safety checks: `pnpm quality:migration-safety` and `pnpm quality:do-migration-safety`. +- [x] Run local quality gates: `pnpm lint`, `pnpm typecheck`, `pnpm test`, `pnpm build`, and package-specific tests. +- [x] Run local specialist review skills and address findings. +- [x] Opened draft PR with the earlier local-only evidence. +- [ ] Replace draft-only evidence with completed specialist, staging, and preflight evidence; remove `needs-human-review` only after all reviews pass. +- [ ] Update idea `01KX4KSXEXQMP41KS34TW9EN01` with final PR, staging, merge, and deployment evidence. + +## Verification Notes + +- `pnpm quality:migration-safety` passed. +- `pnpm quality:do-migration-safety` passed. +- `pnpm typecheck` passed. +- `pnpm lint` passed with existing warnings and no errors. +- `pnpm build` passed. +- `pnpm test` passed. +- `pnpm --filter @simple-agent-manager/api test -- tests/unit/session-snapshots.test.ts tests/unit/routes/workspaces-session-snapshots.test.ts tests/unit/cf-container-runtime-contract.test.ts` passed. +- `pnpm --filter @simple-agent-manager/api typecheck` passed. +- `go test ./internal/server` passed in `packages/vm-agent`. +- `go test -race ./...` passed in `packages/vm-agent`. +- `pnpm --filter @simple-agent-manager/api test:workers -- tests/workers/session-snapshot-wiring.test.ts` was attempted locally, but the Cloudflare `workerd` binary (`@cloudflare/workerd-linux-64@1.20260329.1`) segfaulted during startup before the test body ran. Staging verification remains intentionally deferred by instruction. + +## Local Specialist Review Notes + +- `$cloudflare-specialist`: additive D1 migration and R2-backed snapshot metadata are in place; app enforces snapshot expiry and avoids cron sweeping. R2 lifecycle policy provisioning remains an operational follow-up before merge. +- `$go-specialist`: vm-agent hibernate/restore uses local git bundles, no remote push, progress/idle transfer watchdogs, tar extraction hardening, and `go test -race ./...` passed. +- `$security-auditor`: callback auth is workspace scoped, R2 keys are server derived, fresh secret/env/file injection reruns after restore, and snapshot errors avoid logging secret material. +- `$constitution-validator` and `$env-validator`: snapshot TTL, budgets, thresholds, watchdogs, JSON body limit, and R2 prefix are env-configurable with `DEFAULT_*` constants. +- `$test-engineer`: unit, route, worker wiring, and Go edge tests cover deterministic keys, auth scoping, oversized entries, git operation degradation, and tar traversal rejection. + +## Acceptance Criteria + +- A cf-container session snapshots harness `$HOME` state and uncommitted WIP on control handback without pushing anything to origin. +- Snapshot storage is deterministic, private, R2-backed, one object set per canonical chat session, and retention is env-configurable with a 7-day default. +- Oversized or unsafe snapshot content degrades visibly with a manifest; no files are silently dropped. +- Waking a sleeping cf-container provisions a fresh container, restores `$HOME`, restores WIP as uncommitted work, reruns fresh runtime asset injection, and attempts native resume before visible transcript fallback. +- Ungraceful kills or expired/missing snapshots degrade visibly to existing transcript/fork behavior. +- Timeouts for transfer work are progress/idle watchdogs, not fixed wall-clock caps. +- Tests cover edge cases for public-repo privacy, huge files, repo operation states, missing/corrupt artifacts, retention expiry, and no secret leakage in logs/events. +- PR #1562 passes local, specialist, CI, staging E2E, merge, and production-deploy gates under the superseding authorization. + +## Continuation Findings (PR #1562 landing pass) + +- The original Git WIP bundle path mutated the user's real index and branch; remediation now builds the snapshot commit from a temporary Git index and tests exact status/HEAD/branch preservation. +- Restore checked out a temporary branch and attempted to delete it while current; remediation now imports the advertised bundle ref, materializes its tree, and leaves WIP uncommitted on the original branch. +- R2 cleanup was not provisioned. Pulumi now owns a prefix-scoped lifecycle rule, and its positive TTL configuration is injected into the Worker so application expiry and object deletion stay aligned. +- The initial PR exceeded the API file-size gate and omitted preflight/specialist evidence markers; the node snapshot wrappers are extracted and durable PR evidence must be added after final reviewers complete. +- Final security pass requires upload `Content-Length`, derives artifact sizes from R2, enforces the aggregate budget, validates manifest workspace/chat identity, and rejects restore paths that traverse pre-existing home symlinks. +- PR evidence now includes the required Agent Preflight and eight-reviewer table; staging remains the only intentionally pending merge gate. diff --git a/tasks/backlog/2026-07-11-runtime-neutral-session-hibernate-wake.md b/tasks/backlog/2026-07-11-runtime-neutral-session-hibernate-wake.md deleted file mode 100644 index 07955c062..000000000 --- a/tasks/backlog/2026-07-11-runtime-neutral-session-hibernate-wake.md +++ /dev/null @@ -1,51 +0,0 @@ -# Runtime-neutral session hibernate/wake for cf-container workspaces - -## Problem - -Cloudflare Container instant workspaces lose their local disk after sleep. That currently means harness-native session state under `$HOME` and uncommitted repository work disappear when a sleeping cf-container workspace wakes. Phase 3a of idea `01KX4KSXEXQMP41KS34TW9EN01` must add a runtime-neutral hibernate/restore contract and wire the first consumer to the cf-container runtime. - -Hard constraints from Raphaël, 2026-07-11: -- Open a draft PR only. -- Do not merge. -- Do not deploy to staging or mutate staging. -- Verification is local only. - -## Research Findings - -- `apps/api/src/durable-objects/vm-agent-container.ts` owns cf-container launch, active-work keepalive, `sleeping` lifecycle, and currently returns `503` for sleeping containers. -- `apps/api/src/services/vm-agent-container.ts` wraps the DO and already filters active-work calls to `runtime = 'cf-container'`. -- `apps/api/src/services/node-agent.ts` dispatches initial and follow-up prompts, marking cf-container active work started/ended. -- `apps/api/src/routes/projects/agent-activity-callback.ts` receives VM-agent activity with callback JWT auth and ends active work when the harness reports idle/error. -- `packages/vm-agent/internal/server/workspaces.go` creates `SessionHost` instances. Follow-up prompt handling currently requires an in-memory host, so a fresh container after sleep needs a rebuild path. -- `packages/vm-agent/internal/server/standalone_workspace.go` prepares raw cf-container standalone workspaces by cloning the repo. Restore must happen after clone but before fresh runtime-asset injection/harness start. -- R2 access is via Worker binding for metadata/object ownership. VM-agent needs control-plane endpoints or signed URLs for artifact transfer; do not push WIP refs to the user's remote. -- Relevant rules: migration safety, vertical-slice testing, VM-agent callback auth scoping, long-running progress/idle watchdogs, no hardcoded values, DO concurrency, and credential snapshot resilience. - -## Implementation Checklist - -- [ ] Add additive D1 schema/migration for runtime-neutral session snapshots: canonical chat session key, R2 keys, manifest/degradation, expiry, and restore diagnostics. -- [ ] Add env-configurable defaults for snapshot TTL, size budgets, per-entry thresholds, progress idle watchdogs, and R2 key prefix. -- [ ] Add Worker snapshot service for deterministic one-snapshot-per-session R2 keys, metadata persistence, retention checks, cleanup hooks, and visible degradation state. -- [ ] Add VM-agent callback/API endpoints with callback JWT auth for snapshot upload/download coordination, mounted outside browser-session middleware. -- [ ] Add VM-agent hibernate implementation: git WIP bundle, `$HOME` tar with generic cache exclusions, size/skip manifest, progress idle watchdogs, and no remote push. -- [ ] Add VM-agent restore implementation: restore `$HOME`, restore WIP bundle with soft reset to base, re-run runtime injection afterward, and report visible degraded results. -- [ ] Wire cf-container hibernate on control handback before sleep is allowed. -- [ ] Wire cf-container wake before proxying a sleeping container and before follow-up prompt dispatch; fall back visibly when no snapshot exists or restore fails. -- [ ] Update follow-up prompt handling so a fresh vm-agent can rebuild a `SessionHost` and attempt native harness resume from restored `$HOME`. -- [ ] Add local tests: Worker service/unit tests, Miniflare/DO integration tests, VM-agent Go tests including `go test -race`, and a vertical-slice test covering Worker → DO → vm-agent → R2 state. -- [ ] Run migration safety checks: `pnpm quality:migration-safety` and `pnpm quality:do-migration-safety`. -- [ ] Run local quality gates: `pnpm lint`, `pnpm typecheck`, `pnpm test`, `pnpm build`, and package-specific tests. -- [ ] Run local specialist review skills and address findings. -- [ ] Open draft PR with `needs-human-review`, explicit no-merge/no-staging language, and local-only verification evidence. -- [ ] Append idea `01KX4KSXEXQMP41KS34TW9EN01` with branch, draft PR, and deferred staging/merge status. Leave the idea open. - -## Acceptance Criteria - -- A cf-container session snapshots harness `$HOME` state and uncommitted WIP on control handback without pushing anything to origin. -- Snapshot storage is deterministic, private, R2-backed, one object set per canonical chat session, and retention is env-configurable with a 7-day default. -- Oversized or unsafe snapshot content degrades visibly with a manifest; no files are silently dropped. -- Waking a sleeping cf-container provisions a fresh container, restores `$HOME`, restores WIP as uncommitted work, reruns fresh runtime asset injection, and attempts native resume before visible transcript fallback. -- Ungraceful kills or expired/missing snapshots degrade visibly to existing transcript/fork behavior. -- Timeouts for transfer work are progress/idle watchdogs, not fixed wall-clock caps. -- Tests cover edge cases for public-repo privacy, huge files, repo operation states, missing/corrupt artifacts, retention expiry, and no secret leakage in logs/events. -- The PR is draft-only, not merged, and staging verification is explicitly deferred per instruction. diff --git a/tasks/backlog/2026-07-12-cf-container-wake-restore-hardening.md b/tasks/backlog/2026-07-12-cf-container-wake-restore-hardening.md new file mode 100644 index 000000000..ab456623a --- /dev/null +++ b/tasks/backlog/2026-07-12-cf-container-wake-restore-hardening.md @@ -0,0 +1,28 @@ +# cf-container wake/restore hardening follow-ups + +Deferred, non-blocking hardening items raised by the specialist reviewers during +PR #1562 (cf-container session hibernate/wake/restore). None are regressions +introduced by #1562 — the feature is verified working end-to-end on both the +instant-container and VM paths. The CRITICAL/HIGH items from that review (wake +concurrency mutex per rule 45, and Go unit tests for the token-persist + +reporter-prime invariants) were fixed in #1562 itself. These remaining LOW / +pre-existing items are tracked here. + +## Context +- Origin: PR #1562 specialist review (go-specialist, cloudflare-specialist, security-auditor, test-engineer). +- Feature files: `apps/api/src/durable-objects/vm-agent-container.ts`, `packages/vm-agent/internal/server/session_snapshot.go`, `.../session_snapshot_archive.go`, `apps/api/src/services/session-snapshots.ts`. + +## Acceptance Criteria +- [ ] **Bound the DO-internal restore call with a wall-clock budget.** `wakeFromSnapshot`'s `containerFetch` to `/restore` is not bounded by an `AbortSignal`/timeout inside the DO (the Worker request carries `getCfContainerWakeTimeoutMs`, but a hung restore inside the DO can still block). Add an env-configurable deadline (`.claude/rules/47`). +- [ ] **Decouple restore agent-start from the HTTP request context.** `restoreSessionSnapshot` passes the request `ctx` into `SelectAgent`; a proxy timeout / client disconnect could abort a cold agent install mid-flight. Use a job-owned `context.WithTimeout(context.Background(), )` (`.claude/rules/43`). +- [ ] **Add a decompression size limit** (`io.LimitReader` with configurable max) around the tar stream in `downloadAndExtractTar` to prevent decompression bombs. +- [ ] **Filter/truncate `restoreBody`** before it is passed to `markWakeDegraded` (persisted to D1 `errorMessage`) and returned in the 503 response, so future vm-agent error-message changes cannot leak sensitive content. +- [ ] **Decouple node-management token TTL** from `TERMINAL_TOKEN_EXPIRY_MS`: introduce `NODE_MANAGEMENT_TOKEN_EXPIRY_MS` / `getNodeManagementTokenExpiry` (default 1h). +- [ ] **Validate hostname in `absoluteControlPlaneURL`** to restrict absolute-URL passthrough to the expected SAM domain. +- [ ] **Break the pre-existing circular import**: `node-agent.ts` re-exports `hibernateAgentSessionOnNode`/`restoreAgentSessionOnNode` from `node-agent-session-snapshots`; have callers import directly (mirrors the `node-agent-diagnostics.ts` split done in #1562). +- [ ] **`markWakeDegraded` node status accuracy**: if the container launched before restore failed, reflect the real container state rather than leaving it stale. +- [ ] **Thread `ctx` through `skipOversizedUntracked`** (currently uses `context.Background()` for `git check-ignore`, ignoring cancellation). Pre-existing. +- [ ] **Add an HTTP-level Go test** for `sessionSnapshotHandlerInput` (with a valid node-management JWT) asserting a restore request with `workspaceCallbackToken` in the body persists it — complements the invariant-level test added in #1562. + +## Notes +- The security review's checklist (token scopes, no-leak, restore-endpoint auth, no scope escalation, TTLs) all PASSED for #1562; these are additive defense-in-depth.