diff --git a/.claude/rules/02-quality-gates.md b/.claude/rules/02-quality-gates.md index 5f6b36fa7..8bdd0c8c2 100644 --- a/.claude/rules/02-quality-gates.md +++ b/.claude/rules/02-quality-gates.md @@ -91,6 +91,7 @@ Ask: "What test, if it existed before the breaking change was introduced, would - **If the bug was a missing propagation** (value set in A but never forwarded to B), write a test that constructs the real lifecycle (A then B) and asserts the value arrives. - **If the bug involves streamed UI data that is later reconstructed from durable storage**, write a parity regression test for the persisted representation, not only the live stream. The test MUST include a partial/status-only update event and assert omitted fields do not clear previously visible metadata. See the retained incident lesson in this rule. - **If the bug involves lifecycle control across a runtime boundary** (agent/session/workspace/node stop, cancel, retry, replacement, suspend, or resume), the regression test MUST assert the runtime command is invoked before accepting the terminal state or dispatching replacement work. Database state changes and successful JSON responses are insufficient; the test must prove the external agent/node/workspace control side effect. +- **If runtime liveness is represented in more than one control plane** (for example D1, a session Durable Object, and a runtime Durable Object), timeout and cleanup tests MUST cross those boundaries with deliberately stale secondary state. A heartbeat timeout or sweep may terminalize work only after the runtime owner reports a conclusively terminal lifecycle; sleep, wake, restore, replacement, probe failure, and unknown state are inconclusive. Use one shared lifecycle classifier for every cleanup path so a stale replica cannot strand recoverable work. - **If the bug involves shell or process execution lifecycle** (process groups, child processes, cancellation, timeout, or cleanup after command completion), the regression test MUST cover the success path as well as failure/cancellation paths and prove spawned children are not left alive after the tool or command returns. - **If the bug involves a utility LLM call through a provider-compatible API**, the regression test MUST assert the exact provider payload controls that make the response contract reliable, not just the returned parsed text. For reasoning-capable models, this includes any explicit thinking/reasoning-disable parameters or response-format controls required for the utility to receive text in the field it reads. - **For utility-model defaults and capability changes**, validate the exact production-shaped payload against the real provider before merge, record the accepted and rejected field matrix without prompts or credentials, and test omission of redundant nullable fields. Non-2xx handling must preserve bounded sanitized provider codes/parameter names so later schema drift is diagnosable without logging raw response bodies. diff --git a/.claude/skills/env-reference/SKILL.md b/.claude/skills/env-reference/SKILL.md index c1c35cddd..348a4f32c 100644 --- a/.claude/skills/env-reference/SKILL.md +++ b/.claude/skills/env-reference/SKILL.md @@ -63,10 +63,16 @@ See `apps/api/.env.example` for the full list. Key variables: - `WRANGLER_PORT` — Local dev port (default: 8787) - `BASE_DOMAIN` — Set automatically by sync scripts - `CF_CONTAINER_ENABLED` — Enables Cloudflare Container instant-session runtime in generated deployment envs (default: `true`; set `false` to force VM runtime) -- `CF_CONTAINER_SLEEP_AFTER` — Container idle sleep duration for instant-session runtime (default: `10m`) +- `CF_CONTAINER_SLEEP_AFTER` — Container idle sleep duration for instant-session runtime (default: `1h`) +- `CF_CONTAINER_ACTIVE_WORK_MAX_MS` — Defensive maximum active-work keepalive duration (default: `7200000`) +- `CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS` — Active-work keepalive renewal interval (default: `300000`) - `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`) -- `$1 +- `CF_CONTAINER_WAKE_TIMEOUT_MS` — Max wait for launch, snapshot restore, and request readiness (default: `120000`) +- `CF_CONTAINER_RECOVERY_MAX_ATTEMPTS` — Max restore attempts before terminal status reconciliation (default: `2`) +- `INSTANT_STALE_CALLBACK_MARGIN_MS` — Freshness margin for rejecting destructive callbacks from superseded Instant containers (default: `60000`) +- `CF_CONTAINER_CREATE_WORKSPACE_TIMEOUT_MS` — Synchronous workspace creation and clone budget (default: `120000`) +- `CF_CONTAINER_CLONE_FILTER` — Git partial-clone filter (default: `blob:none`; `off` disables partial clone) - `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`) diff --git a/apps/api/.env.example b/apps/api/.env.example index 78a72cfd2..c81627d15 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -74,9 +74,19 @@ BASE_DOMAIN=workspaces.example.com # Generated deployment environments default this to true. Set false to force # profile/runtime resolution back to cloud VMs. # CF_CONTAINER_ENABLED=true -# CF_CONTAINER_SLEEP_AFTER=10m +# CF_CONTAINER_SLEEP_AFTER=1h +# Maximum active-work keepalive lifetime and renewal interval. +# CF_CONTAINER_ACTIVE_WORK_MAX_MS=7200000 +# CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS=300000 # CF_CONTAINER_VM_AGENT_PORT=8080 # CF_CONTAINER_PORT_READY_TIMEOUT_MS=30000 +# Budget for a sleeping/recovering runtime to launch, restore, and accept work. +# CF_CONTAINER_WAKE_TIMEOUT_MS=120000 +# Maximum snapshot restore attempts before terminal status reconciliation. +# CF_CONTAINER_RECOVERY_MAX_ATTEMPTS=2 +# Freshness margin for rejecting destructive (error/failed) callbacks from a +# superseded Instant container generation after a completed recovery. +# INSTANT_STALE_CALLBACK_MARGIN_MS=60000 # CF_CONTAINER_WORKSPACE_BASE_DIR=/workspaces # Budget for the synchronous standalone create-workspace request (includes the # repository clone inside the container). Default 120000 (120s). diff --git a/apps/api/src/db/migrations/0099_tasks_workspace_id_index.sql b/apps/api/src/db/migrations/0099_tasks_workspace_id_index.sql new file mode 100644 index 000000000..306f409d6 --- /dev/null +++ b/apps/api/src/db/migrations/0099_tasks_workspace_id_index.sql @@ -0,0 +1,8 @@ +-- Index tasks by workspace_id for the mass-outage Instant recovery hot path. +-- persistRuntimeRecoveryFailed and other workspace-scoped task lookups run +-- `SELECT ... FROM tasks WHERE workspace_id = ? AND status IN (...)`, which +-- previously required a full table scan. Partial (workspace_id IS NOT NULL) +-- because most tasks never bind a workspace; the predicate still serves every +-- `workspace_id = ?` lookup since the bound value is always non-null. +CREATE INDEX IF NOT EXISTS idx_tasks_workspace_id + ON tasks(workspace_id) WHERE workspace_id IS NOT NULL; diff --git a/apps/api/src/db/schema.ts b/apps/api/src/db/schema.ts index f77fd8fd8..f74857088 100644 --- a/apps/api/src/db/schema.ts +++ b/apps/api/src/db/schema.ts @@ -909,6 +909,12 @@ export const tasks = sqliteTable( chatSessionIdUnique: uniqueIndex('idx_tasks_chat_session_id_unique') .on(table.chatSessionId) .where(sql`chat_session_id IS NOT NULL`), + // Supports the `WHERE workspace_id = ? AND status IN (...)` lookups on the + // mass-outage recovery hot path (persistRuntimeRecoveryFailed) and other + // workspace-scoped task queries. Partial: most tasks never bind a workspace. + workspaceIdIdx: index('idx_tasks_workspace_id') + .on(table.workspaceId) + .where(sql`workspace_id IS NOT NULL`), triggerExecutionIdIdx: index('idx_tasks_trigger_execution_id') .on(table.triggerExecutionId) .where(sql`trigger_execution_id IS NOT NULL`), diff --git a/apps/api/src/durable-objects/project-data/acp-sessions.ts b/apps/api/src/durable-objects/project-data/acp-sessions.ts index 7ef3772ad..665640548 100644 --- a/apps/api/src/durable-objects/project-data/acp-sessions.ts +++ b/apps/api/src/durable-objects/project-data/acp-sessions.ts @@ -396,7 +396,16 @@ export function checkHeartbeatTimeouts( reason?: string | null; errorMessage?: string; metadata?: Record | null; - }) => Promise + }) => Promise, + options: { + shouldDeferTimeout?: (session: { + id: string; + chatSessionId: string; + workspaceId: string | null; + nodeId: string | null; + lastHeartbeatAt: number | null; + }) => Promise<{ defer: boolean; reason: string }>; + } = {} ): Promise> { const detectionWindow = parseInt( env.ACP_SESSION_DETECTION_WINDOW_MS || String(ACP_SESSION_DEFAULTS.DETECTION_WINDOW_MS), @@ -419,6 +428,28 @@ export function checkHeartbeatTimeouts( const timedOut: Array<{ sessionId: string; workspaceId: string | null }> = []; const promises = staleSessions.map(async (session) => { try { + if (options.shouldDeferTimeout) { + try { + const decision = await options.shouldDeferTimeout(session); + if (decision.defer) { + log.info('acp_session.heartbeat_timeout_deferred', { + sessionId: session.id, + workspaceId: session.workspaceId, + nodeId: session.nodeId, + reason: decision.reason, + }); + return; + } + } catch (err) { + log.warn('acp_session.heartbeat_timeout_policy_failed', { + sessionId: session.id, + workspaceId: session.workspaceId, + nodeId: session.nodeId, + error: err instanceof Error ? err.message : String(err), + }); + return; + } + } await transitionFn(session.id, 'interrupted', { actorType: 'alarm', reason: 'Heartbeat timeout exceeded detection window', errorMessage: `Heartbeat timeout: last heartbeat at ${session.lastHeartbeatAt}, cutoff was ${cutoff}`, @@ -476,7 +507,8 @@ export function computeHeartbeatAlarmTime(sql: SqlStorage, env: Env): number | n const earliestHeartbeat = parseMinEarliest(earliestRow, 'acp_sessions.earliest_heartbeat'); if (earliestHeartbeat === null) return null; - return earliestHeartbeat + detectionWindow; + const dueAt = earliestHeartbeat + detectionWindow; + return dueAt <= Date.now() ? Date.now() + detectionWindow : dueAt; } /** diff --git a/apps/api/src/durable-objects/project-data/index.ts b/apps/api/src/durable-objects/project-data/index.ts index cf8d5c0a2..3395fbdc7 100644 --- a/apps/api/src/durable-objects/project-data/index.ts +++ b/apps/api/src/durable-objects/project-data/index.ts @@ -30,6 +30,7 @@ import * as missionState from './missions'; import * as policies from './policies'; import * as reconciliation from './reconciliation'; import { parseCountCnt, parseMaxLatest, parseMetaValue } from './row-schemas'; +import { checkRuntimeHeartbeatTimeouts } from './runtime-heartbeat-policy'; import * as sessionState from './session-state'; import * as sessionSummarySync from './session-summary-sync'; import * as sessions from './sessions'; @@ -341,9 +342,8 @@ export class ProjectData extends DurableObject { // --- DO Alarm Handler --- async alarm(): Promise { - const timedOut = await acpSessions.checkHeartbeatTimeouts(this.sql, this.env, async (sessionId, toStatus, opts) => { - await this.transitionAcpSession(sessionId, toStatus, opts); - }); + const timedOut = await checkRuntimeHeartbeatTimeouts( + this.sql, this.env, this.transitionAcpSession.bind(this)); // For conversation-mode sessions, couple agent death to workspace death. // Stop workspaces whose ACP sessions timed out to prevent zombie state. diff --git a/apps/api/src/durable-objects/project-data/runtime-heartbeat-policy.ts b/apps/api/src/durable-objects/project-data/runtime-heartbeat-policy.ts new file mode 100644 index 000000000..b4f5e84b1 --- /dev/null +++ b/apps/api/src/durable-objects/project-data/runtime-heartbeat-policy.ts @@ -0,0 +1,79 @@ +import { DEFAULT_TASK_LIVENESS_PROBE_TIMEOUT_MS } from '@simple-agent-manager/shared'; + +import type { VmAgentContainer } from '../vm-agent-container'; +import { isVmAgentContainerLifecycleTerminal } from '../vm-agent-container-lifecycle'; +import { checkHeartbeatTimeouts } from './acp-sessions'; +import type { Env } from './types'; + +interface HeartbeatTimeoutCandidate { + workspaceId: string | null; + nodeId: string | null; +} + +const TERMINAL_WORKSPACE_STATUSES = new Set(['stopping', 'stopped', 'error', 'deleted']); + +export function checkRuntimeHeartbeatTimeouts( + sql: SqlStorage, + env: Env, + transitionFn: Parameters[2] +) { + return checkHeartbeatTimeouts(sql, env, transitionFn, { + shouldDeferTimeout: (session) => shouldDeferRuntimeHeartbeatTimeout(env, session), + }); +} + +function probeTimeoutMs(env: Env): number { + const configured = Number.parseInt(env.TASK_LIVENESS_PROBE_TIMEOUT_MS ?? '', 10); + return Number.isFinite(configured) && configured > 0 + ? configured + : DEFAULT_TASK_LIVENESS_PROBE_TIMEOUT_MS; +} + +export async function shouldDeferRuntimeHeartbeatTimeout( + env: Env, + candidate: HeartbeatTimeoutCandidate +): Promise<{ defer: boolean; reason: string }> { + if (!candidate.workspaceId || !candidate.nodeId) { + return { defer: false, reason: 'runtime_identity_incomplete' }; + } + + const row = await env.DATABASE.prepare( + `SELECT w.status AS workspace_status, n.runtime AS node_runtime + FROM workspaces w + LEFT JOIN nodes n ON n.id = w.node_id + WHERE w.id = ? AND w.node_id = ? + LIMIT 1` + ) + .bind(candidate.workspaceId, candidate.nodeId) + .first<{ workspace_status: string; node_runtime: string | null }>(); + if (!row || row.node_runtime !== 'cf-container') { + return { defer: false, reason: row ? 'non_container_runtime' : 'workspace_missing' }; + } + if (TERMINAL_WORKSPACE_STATUSES.has(row.workspace_status)) { + return { defer: false, reason: `workspace_${row.workspace_status}` }; + } + if (!env.VM_AGENT_CONTAINER) { + return { defer: true, reason: 'cf_container_lifecycle_binding_unavailable' }; + } + + const id = env.VM_AGENT_CONTAINER.idFromName(candidate.nodeId.toLowerCase()); + const stub = env.VM_AGENT_CONTAINER.get(id) as DurableObjectStub; + const timeoutMs = probeTimeoutMs(env); + const timeout = Symbol('container_lifecycle_probe_timeout'); + let timer: ReturnType | undefined; + const lifecycle = await Promise.race([ + stub.inspectLifecycle(), + new Promise((resolve) => { + timer = setTimeout(() => resolve(timeout), timeoutMs); + }), + ]).finally(() => { + if (timer) clearTimeout(timer); + }); + if (lifecycle === timeout) { + return { defer: true, reason: 'cf_container_lifecycle_timeout' }; + } + return { + defer: !isVmAgentContainerLifecycleTerminal(lifecycle.status), + reason: `cf_container_${lifecycle.status ?? 'unknown'}`, + }; +} diff --git a/apps/api/src/durable-objects/project-data/types.ts b/apps/api/src/durable-objects/project-data/types.ts index 4e014f83d..4058840f7 100644 --- a/apps/api/src/durable-objects/project-data/types.ts +++ b/apps/api/src/durable-objects/project-data/types.ts @@ -2,8 +2,12 @@ * Shared types and utilities for ProjectData DO modules. */ +import type { VmAgentContainer } from '../vm-agent-container'; + export type Env = { DATABASE: D1Database; + VM_AGENT_CONTAINER?: DurableObjectNamespace; + TASK_LIVENESS_PROBE_TIMEOUT_MS?: string; BASE_DOMAIN?: string; DO_SUMMARY_SYNC_DEBOUNCE_MS?: string; MAX_SESSIONS_PER_PROJECT?: string; diff --git a/apps/api/src/durable-objects/vm-agent-container-active-work.ts b/apps/api/src/durable-objects/vm-agent-container-active-work.ts new file mode 100644 index 000000000..78f80c766 --- /dev/null +++ b/apps/api/src/durable-objects/vm-agent-container-active-work.ts @@ -0,0 +1,120 @@ +import { log } from '../lib/logger'; + +export type ActiveWorkStatus = 'active' | 'ended' | 'expired'; + +export interface ActiveWorkState { + status: ActiveWorkStatus; + nodeId: string; + workspaceId: string; + agentSessionId: string; + reason: string; + activeSince: number; + lastRenewedAt: number; + deadlineAt: number; + endedAt?: number; + endReason?: string; +} + +export interface ActiveWorkRuntime { + storage: DurableObjectStorage; + activeWorkMaxMs: number; + renewIntervalMs: number; + renewActivityTimeout: () => void; + replaceSchedule: (delayMs: number) => Promise; + clearSchedule: () => Promise; +} + +export const ACTIVE_WORK_KEY = 'activeWork'; + +// Concurrency note: these read-check-write helpers on ACTIVE_WORK_KEY are +// deliberately NOT serialized under the DO's lifecycleChain/wakeChain mutexes. +// Active-work state is advisory keepalive bookkeeping (it only extends the +// container's idle deadline), so a lost interleaving self-heals on the next +// keepalive tick or prompt: the worst case is one extra/short renewal window, +// never an incorrect lifecycle transition. Folding these under the lifecycle +// mutex would needlessly serialize the frequent keepalive schedule against +// stop/wake/recovery critical sections. The authoritative lifecycle writes +// (lifecycleStatus, runtimeRecovery) remain mutex-protected in the DO. + +export async function startActiveWork( + runtime: ActiveWorkRuntime, + nodeId: string, + input: { workspaceId: string; agentSessionId: string; reason: string } +): Promise { + const now = Date.now(); + const activeWork: ActiveWorkState = { + status: 'active', + nodeId, + workspaceId: input.workspaceId, + agentSessionId: input.agentSessionId, + reason: input.reason, + activeSince: now, + lastRenewedAt: now, + deadlineAt: now + runtime.activeWorkMaxMs, + }; + runtime.renewActivityTimeout(); + await runtime.storage.put(ACTIVE_WORK_KEY, activeWork); + await runtime.replaceSchedule(runtime.renewIntervalMs); + log.info('vm_agent_container_active_work_started', { + nodeId, + workspaceId: input.workspaceId, + agentSessionId: input.agentSessionId, + reason: input.reason, + activeSince: new Date(now).toISOString(), + deadlineAt: new Date(activeWork.deadlineAt).toISOString(), + }); +} + +export async function endActiveWork(runtime: ActiveWorkRuntime, reason: string): Promise { + const activeWork = await runtime.storage.get(ACTIVE_WORK_KEY); + if (!activeWork || activeWork.status !== 'active') { + await runtime.clearSchedule(); + return; + } + const now = Date.now(); + await runtime.storage.put(ACTIVE_WORK_KEY, { + ...activeWork, + status: 'ended', + endedAt: now, + endReason: reason, + } satisfies ActiveWorkState); + await runtime.clearSchedule(); + log.info('vm_agent_container_active_work_ended', { + nodeId: activeWork.nodeId, + workspaceId: activeWork.workspaceId, + agentSessionId: activeWork.agentSessionId, + reason, + activeSince: new Date(activeWork.activeSince).toISOString(), + lastRenewedAt: new Date(activeWork.lastRenewedAt).toISOString(), + endedAt: new Date(now).toISOString(), + }); +} + +export async function renewActiveWork(runtime: ActiveWorkRuntime): Promise { + const activeWork = await runtime.storage.get(ACTIVE_WORK_KEY); + if (!activeWork || activeWork.status !== 'active') { + await runtime.clearSchedule(); + return; + } + const now = Date.now(); + if (now >= activeWork.deadlineAt) { + await runtime.storage.put(ACTIVE_WORK_KEY, { + ...activeWork, + status: 'expired', + endedAt: now, + endReason: 'keepalive_deadline_exceeded', + } satisfies ActiveWorkState); + await runtime.clearSchedule(); + log.warn('vm_agent_container_active_work_keepalive_expired', { + nodeId: activeWork.nodeId, + workspaceId: activeWork.workspaceId, + agentSessionId: activeWork.agentSessionId, + deadlineAt: new Date(activeWork.deadlineAt).toISOString(), + }); + return; + } + + runtime.renewActivityTimeout(); + await runtime.storage.put(ACTIVE_WORK_KEY, { ...activeWork, lastRenewedAt: now }); + await runtime.replaceSchedule(runtime.renewIntervalMs); +} diff --git a/apps/api/src/durable-objects/vm-agent-container-lifecycle.ts b/apps/api/src/durable-objects/vm-agent-container-lifecycle.ts new file mode 100644 index 000000000..0e9b3cbf0 --- /dev/null +++ b/apps/api/src/durable-objects/vm-agent-container-lifecycle.ts @@ -0,0 +1,86 @@ +import { + type ActiveWorkState, + type ActiveWorkStatus, +} from './vm-agent-container-active-work'; +import type { + RuntimeRecoveryPhase, + RuntimeRecoveryState, + RuntimeRecoveryTrigger, +} from './vm-agent-container-recovery'; + +export type VmAgentContainerLifecycleStatus = + | 'launching' + | 'running' + | 'stopping' + | 'stopped' + | 'sleeping' + | 'recovering' + | 'waking' + | 'restoring' + | 'degraded' + | 'expired' + | 'error'; + +export interface VmAgentContainerLifecycleInspection { + status: VmAgentContainerLifecycleStatus | null; + recoveryPhase: RuntimeRecoveryPhase | null; + recoveryTrigger: RuntimeRecoveryTrigger | null; + activeWorkStatus: ActiveWorkStatus | null; +} + +export interface VmAgentContainerLivenessClassification { + live: boolean; + conclusive: boolean; + reason: string; +} + +export async function inspectStoredVmAgentContainerLifecycle( + storage: DurableObjectStorage, + recoveryStateKey: string, + activeWorkKey: string +): Promise { + const [status, recovery, activeWork] = await Promise.all([ + storage.get('lifecycleStatus'), + storage.get(recoveryStateKey), + storage.get(activeWorkKey), + ]); + return { + status: status ?? null, + recoveryPhase: recovery?.phase ?? null, + recoveryTrigger: recovery?.trigger ?? null, + activeWorkStatus: activeWork?.status ?? null, + }; +} + +const TERMINAL_LIFECYCLE_STATUSES = new Set([ + 'stopping', + 'stopped', + 'expired', + 'error', +]); + +export function isVmAgentContainerLifecycleTerminal( + status: VmAgentContainerLifecycleStatus | null +): boolean { + return status !== null && TERMINAL_LIFECYCLE_STATUSES.has(status); +} + +export function classifyVmAgentContainerLiveness( + inspection: VmAgentContainerLifecycleInspection +): VmAgentContainerLivenessClassification { + if (isVmAgentContainerLifecycleTerminal(inspection.status)) { + return { + live: false, + conclusive: true, + reason: `cf_container_${inspection.status}`, + }; + } + if (inspection.status === 'running' && inspection.activeWorkStatus === 'active') { + return { live: true, conclusive: true, reason: 'cf_container_active_work' }; + } + return { + live: false, + conclusive: false, + reason: `cf_container_${inspection.status ?? 'unknown'}_resumable`, + }; +} diff --git a/apps/api/src/durable-objects/vm-agent-container-recovery.ts b/apps/api/src/durable-objects/vm-agent-container-recovery.ts new file mode 100644 index 000000000..18df2c7bc --- /dev/null +++ b/apps/api/src/durable-objects/vm-agent-container-recovery.ts @@ -0,0 +1,251 @@ +import { and, desc, eq, inArray } 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 { ulid } from '../lib/ulid'; +import * as projectDataService from '../services/project-data'; + +export const RUNTIME_RECOVERING_MESSAGE = + 'Instant session interrupted; restoring the last safe checkpoint.'; +export const RUNTIME_REQUEST_INTERRUPTED_MESSAGE = + 'Your message is saved, but delivery was interrupted and its execution outcome is unknown. It was not replayed automatically. After restore finishes, check the transcript and partial output before deciding whether to send it again.'; +export const RUNTIME_RECOVERY_DEGRADED_MESSAGE = + 'The Instant session could not restore its last safe checkpoint. Your transcript and partial output are still available.'; +export const RUNTIME_STOPPED_MESSAGE = 'This Instant session was stopped and cannot be resumed.'; + +export type RuntimeRecoveryCode = + | 'RUNTIME_RECOVERING' + | 'RUNTIME_REQUEST_INTERRUPTED' + | 'RUNTIME_RECOVERY_DEGRADED' + | 'RUNTIME_STOPPED'; + +export function getRuntimeRecoveryMessage(code: RuntimeRecoveryCode): string { + if (code === 'RUNTIME_RECOVERING') return RUNTIME_RECOVERING_MESSAGE; + if (code === 'RUNTIME_REQUEST_INTERRUPTED') return RUNTIME_REQUEST_INTERRUPTED_MESSAGE; + if (code === 'RUNTIME_STOPPED') return RUNTIME_STOPPED_MESSAGE; + return RUNTIME_RECOVERY_DEGRADED_MESSAGE; +} + +export type RuntimeRecoveryPhase = 'pending' | 'waking' | 'restoring' | 'degraded' | 'exhausted'; + +export type RuntimeRecoveryTrigger = 'idle' | 'stop' | 'error' | 'request'; + +export type RuntimeRecoveryCause = + | { kind: 'idle_sleep' } + | { kind: 'container_stop'; reason: 'exit' | 'runtime_signal'; exitCode: number } + | { kind: 'container_error'; errorName: string } + | { kind: 'transport_interrupted'; errorName: string } + | { kind: 'missing_session_host'; httpStatus: number }; + +export interface RuntimeRecoveryState { + version: 1; + phase: RuntimeRecoveryPhase; + trigger: RuntimeRecoveryTrigger; + cause: RuntimeRecoveryCause; + attempts: number; + promptDisposition: 'none' | 'manual_retry'; + agentSessionId: string | null; + startedAt: number; + updatedAt: number; + lastFailure?: { + kind: 'launch' | 'restore_http' | 'restore_status' | 'unexpected'; + httpStatus?: number; + }; +} + +export interface RuntimeRecoveryTarget { + nodeId: string; + workspaceId: string; + projectId: string; + chatSessionId: string; + agentSessionId: string; +} + +export interface RuntimeRecoveryContext { + userId: string; + chatSessionId: string; + agentSessionId: string; + agentType: string | null; +} + +export function toRuntimeRecoveryTarget( + config: { nodeId: string; workspaceId: string; projectId: string }, + context: RuntimeRecoveryContext +): RuntimeRecoveryTarget { + return { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + projectId: config.projectId, + chatSessionId: context.chatSessionId, + agentSessionId: context.agentSessionId, + }; +} + +const ACTIVE_TASK_STATUSES = ['in_progress', 'delegated', 'awaiting_followup'] as const; + +export async function loadRuntimeRecoveryContext( + env: Env, + input: { workspaceId: string; preferredAgentSessionId?: string | null } +): Promise { + const db = drizzle(env.DATABASE, { schema }); + const workspace = await db + .select({ + userId: schema.workspaces.userId, + chatSessionId: schema.workspaces.chatSessionId, + }) + .from(schema.workspaces) + .where(eq(schema.workspaces.id, input.workspaceId)) + .get(); + if (!workspace?.chatSessionId) return null; + + const agentSession = await db + .select({ id: schema.agentSessions.id, agentType: schema.agentSessions.agentType }) + .from(schema.agentSessions) + .where( + input.preferredAgentSessionId + ? and( + eq(schema.agentSessions.workspaceId, input.workspaceId), + eq(schema.agentSessions.id, input.preferredAgentSessionId) + ) + : eq(schema.agentSessions.workspaceId, input.workspaceId) + ) + .orderBy(desc(schema.agentSessions.updatedAt)) + .get(); + if (!agentSession) return null; + + return { + userId: workspace.userId, + chatSessionId: workspace.chatSessionId, + agentSessionId: agentSession.id, + agentType: agentSession.agentType, + }; +} + +export async function persistRuntimeRecovering( + env: Env, + target: RuntimeRecoveryTarget +): Promise { + const now = new Date().toISOString(); + await env.DATABASE.batch([ + env.DATABASE.prepare( + `UPDATE nodes + SET status = 'recovery', health_status = 'unhealthy', error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(RUNTIME_RECOVERING_MESSAGE, now, target.nodeId), + env.DATABASE.prepare( + `UPDATE workspaces SET status = 'recovery', error_message = ?, updated_at = ? WHERE id = ?` + ).bind(RUNTIME_RECOVERING_MESSAGE, now, target.workspaceId), + env.DATABASE.prepare( + `UPDATE agent_sessions + SET status = 'recovery', stopped_at = NULL, error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(RUNTIME_RECOVERING_MESSAGE, now, target.agentSessionId), + ]); +} + +export async function persistRuntimeRecovered( + env: Env, + target: RuntimeRecoveryTarget, + promptDisposition: RuntimeRecoveryState['promptDisposition'] +): Promise { + const now = new Date().toISOString(); + const agentMessage = + promptDisposition === 'manual_retry' ? RUNTIME_REQUEST_INTERRUPTED_MESSAGE : null; + await env.DATABASE.batch([ + env.DATABASE.prepare( + `UPDATE nodes + SET status = 'running', health_status = 'healthy', error_message = NULL, updated_at = ? + WHERE id = ?` + ).bind(now, target.nodeId), + env.DATABASE.prepare( + `UPDATE workspaces SET status = 'running', error_message = NULL, updated_at = ? WHERE id = ?` + ).bind(now, target.workspaceId), + env.DATABASE.prepare( + `UPDATE agent_sessions + SET status = 'running', stopped_at = NULL, error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(agentMessage, now, target.agentSessionId), + ]); +} + +export async function persistRuntimeRecoveryFailed( + env: Env, + target: RuntimeRecoveryTarget +): Promise { + const db = drizzle(env.DATABASE, { schema }); + const now = new Date().toISOString(); + const task = await db + .select({ id: schema.tasks.id }) + .from(schema.tasks) + .where( + and( + eq(schema.tasks.workspaceId, target.workspaceId), + inArray(schema.tasks.status, [...ACTIVE_TASK_STATUSES]) + ) + ) + .orderBy(desc(schema.tasks.updatedAt)) + .get(); + + const statements: D1PreparedStatement[] = [ + env.DATABASE.prepare( + `UPDATE nodes + SET status = 'error', health_status = 'unhealthy', error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(RUNTIME_RECOVERY_DEGRADED_MESSAGE, now, target.nodeId), + env.DATABASE.prepare( + `UPDATE workspaces SET status = 'error', error_message = ?, updated_at = ? WHERE id = ?` + ).bind(RUNTIME_RECOVERY_DEGRADED_MESSAGE, now, target.workspaceId), + env.DATABASE.prepare( + `UPDATE agent_sessions + SET status = 'error', stopped_at = ?, error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(now, RUNTIME_RECOVERY_DEGRADED_MESSAGE, now, target.agentSessionId), + ]; + + if (task) { + statements.push( + env.DATABASE.prepare( + `INSERT INTO task_status_events + (id, task_id, from_status, to_status, actor_type, actor_id, reason, created_at) + SELECT ?, id, status, 'failed', 'system', ?, ?, ? + FROM tasks + WHERE id = ? AND status IN ('in_progress', 'delegated', 'awaiting_followup')` + ).bind(ulid(), target.nodeId, 'Instant runtime recovery exhausted', now, task.id), + env.DATABASE.prepare( + `UPDATE tasks + SET status = 'failed', execution_step = NULL, error_message = ?, updated_at = ? + WHERE id = ? AND status IN ('in_progress', 'delegated', 'awaiting_followup')` + ).bind(RUNTIME_RECOVERY_DEGRADED_MESSAGE, now, task.id) + ); + } + + await env.DATABASE.batch(statements); + + await projectDataService + .transitionAcpSession(env, target.projectId, target.agentSessionId, 'failed', { + actorType: 'system', + actorId: target.nodeId, + reason: 'Instant runtime recovery exhausted', + errorMessage: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + workspaceId: target.workspaceId, + nodeId: target.nodeId, + }) + .catch((error) => { + log.warn('vm_agent_container_recovery.acp_reconcile_failed', { + nodeId: target.nodeId, + workspaceId: target.workspaceId, + error, + }); + }); + await projectDataService + .failSession(env, target.projectId, target.chatSessionId, RUNTIME_RECOVERY_DEGRADED_MESSAGE) + .catch((error) => { + log.warn('vm_agent_container_recovery.chat_reconcile_failed', { + nodeId: target.nodeId, + workspaceId: target.workspaceId, + error, + }); + }); +} diff --git a/apps/api/src/durable-objects/vm-agent-container-runtime.ts b/apps/api/src/durable-objects/vm-agent-container-runtime.ts new file mode 100644 index 000000000..22a157d33 --- /dev/null +++ b/apps/api/src/durable-objects/vm-agent-container-runtime.ts @@ -0,0 +1,190 @@ +import type { Env } from '../env'; +import { log } from '../lib/logger'; +import { signNodeManagementToken } from '../services/jwt'; +import { + RUNTIME_RECOVERING_MESSAGE, + RUNTIME_RECOVERY_DEGRADED_MESSAGE, + RUNTIME_REQUEST_INTERRUPTED_MESSAGE, + RUNTIME_STOPPED_MESSAGE, + type RuntimeRecoveryCode, +} from './vm-agent-container-recovery'; + +interface RuntimeRecoveryResponseResult { + code?: RuntimeRecoveryCode; + message?: string; +} + +export interface RuntimeIdentity { + nodeId: string; + workspaceId: string; +} + +export function isMutatingRuntimeRequest(request: Request): boolean { + return !['GET', 'HEAD', 'OPTIONS'].includes(request.method.toUpperCase()); +} + +export function runtimeRecoveryResponse( + code: RuntimeRecoveryCode, + message: string, + status: 409 | 410 | 503 +): Response { + return Response.json({ error: code, message }, { status }); +} + +export function runtimeResultResponse(result: RuntimeRecoveryResponseResult): Response { + if (result.code === 'RUNTIME_STOPPED') { + return runtimeRecoveryResponse(result.code, result.message ?? RUNTIME_STOPPED_MESSAGE, 410); + } + return runtimeRecoveryResponse( + result.code ?? 'RUNTIME_RECOVERY_DEGRADED', + result.message ?? RUNTIME_RECOVERY_DEGRADED_MESSAGE, + result.code === 'RUNTIME_RECOVERING' ? 503 : 409 + ); +} + +export function interruptedRuntimeRequestResponse(request: Request): Response { + if (isMutatingRuntimeRequest(request)) { + return runtimeRecoveryResponse( + 'RUNTIME_REQUEST_INTERRUPTED', + RUNTIME_REQUEST_INTERRUPTED_MESSAGE, + 409 + ); + } + return runtimeRecoveryResponse('RUNTIME_RECOVERING', RUNTIME_RECOVERING_MESSAGE, 503); +} +export async function isMissingSessionHostResponse(response: Response): Promise { + if (response.status !== 404) return false; + const body = await response + .clone() + .text() + .catch(() => ''); + return body.toLowerCase().includes('no active agent session found'); +} + +const LIVE_SESSION_HOST_STATUSES = new Set(['idle', 'starting', 'ready', 'prompting']); + +export async function probeLiveRuntimeSession(input: { + env: Env; + userId: string; + nodeId: string; + workspaceId: string; + agentSessionId: string; + vmAgentPort: number; + containerFetch: (request: Request, port: number) => Promise; +}): Promise { + const { token } = await signNodeManagementToken( + input.userId, + input.nodeId, + input.workspaceId, + input.env + ); + const url = new URL( + `http://localhost:${input.vmAgentPort}/workspaces/${input.workspaceId}/agent-sessions` + ); + const response = await input.containerFetch( + new Request(url.toString(), { + headers: { + Authorization: `Bearer ${token}`, + 'X-SAM-Node-Id': input.nodeId, + 'X-SAM-Workspace-Id': input.workspaceId, + }, + }), + input.vmAgentPort + ); + if (!response.ok) { + log.warn('vm_agent_container_runtime_session_probe_failed', { + nodeId: input.nodeId, + workspaceId: input.workspaceId, + agentSessionId: input.agentSessionId, + status: response.status, + }); + throw new Error(`runtime_session_probe_http_${response.status}`); + } + const body = (await response.json()) as { + sessions?: Array<{ id?: unknown; hostStatus?: unknown }>; + }; + return Boolean( + body.sessions?.some( + (session) => + session.id === input.agentSessionId && + typeof session.hostStatus === 'string' && + LIVE_SESSION_HOST_STATUSES.has(session.hostStatus) + ) + ); +} + +export function parsePositiveRuntimeSetting(raw: string | undefined, fallback: number): number { + const parsed = raw ? Number.parseInt(raw, 10) : fallback; + return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback; +} + +export function resolveRuntimeSettings( + env: Env, + defaults: { + portReadyTimeoutMs: number; + activeWorkMaxMs: number; + keepaliveRenewIntervalMs: number; + recoveryMaxAttempts: number; + } +) { + return { + portReadyTimeoutMs: parsePositiveRuntimeSetting( + env.CF_CONTAINER_PORT_READY_TIMEOUT_MS || env.SANDBOX_EXEC_TIMEOUT_MS, + defaults.portReadyTimeoutMs + ), + activeWorkMaxMs: parsePositiveRuntimeSetting( + env.CF_CONTAINER_ACTIVE_WORK_MAX_MS, + defaults.activeWorkMaxMs + ), + keepaliveRenewIntervalMs: parsePositiveRuntimeSetting( + env.CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS, + defaults.keepaliveRenewIntervalMs + ), + recoveryMaxAttempts: parsePositiveRuntimeSetting( + env.CF_CONTAINER_RECOVERY_MAX_ATTEMPTS, + defaults.recoveryMaxAttempts + ), + }; +} + +export async function persistRuntimeSleeping(env: Env, identity: RuntimeIdentity): Promise { + const now = new Date().toISOString(); + await env.DATABASE.batch([ + env.DATABASE.prepare( + `UPDATE nodes + SET status = 'sleeping', health_status = 'unhealthy', error_message = NULL, updated_at = ? + WHERE id = ?` + ).bind(now, identity.nodeId), + env.DATABASE.prepare( + `UPDATE workspaces SET status = 'sleeping', error_message = NULL, updated_at = ? WHERE id = ?` + ).bind(now, identity.workspaceId), + env.DATABASE.prepare( + `UPDATE agent_sessions SET status = 'sleeping', error_message = NULL, updated_at = ? WHERE workspace_id = ?` + ).bind(now, identity.workspaceId), + ]); +} + +export async function persistRuntimeEnded( + env: Env, + identity: RuntimeIdentity, + status: 'stopped' | 'error', + message: string +): Promise { + const now = new Date().toISOString(); + const errorMessage = status === 'stopped' ? null : message; + await env.DATABASE.batch([ + env.DATABASE.prepare( + `UPDATE nodes + SET status = ?, health_status = 'unhealthy', error_message = ?, updated_at = ? + WHERE id = ?` + ).bind(status, errorMessage, now, identity.nodeId), + env.DATABASE.prepare( + `UPDATE workspaces SET status = ?, error_message = ?, updated_at = ? WHERE id = ?` + ).bind(status, errorMessage, now, identity.workspaceId), + env.DATABASE.prepare( + `UPDATE agent_sessions + SET status = ?, stopped_at = ?, error_message = ?, updated_at = ? + WHERE workspace_id = ?` + ).bind(status, now, errorMessage, now, identity.workspaceId), + ]); +} diff --git a/apps/api/src/durable-objects/vm-agent-container.ts b/apps/api/src/durable-objects/vm-agent-container.ts index 49a06aef7..e2c3258f8 100644 --- a/apps/api/src/durable-objects/vm-agent-container.ts +++ b/apps/api/src/durable-objects/vm-agent-container.ts @@ -1,16 +1,55 @@ +// FILE SIZE EXCEPTION: Durable Object recovery state machine — method groups are already extracted into vm-agent-container-{recovery,runtime,lifecycle,active-work}.ts; the remaining class body is the interlocking mutex-guarded lifecycle critical sections (rule 45), which must stay reviewable as one unit. See .claude/rules/18-file-size-limits.md import { Container, switchPort } from '@cloudflare/containers'; -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'; +import { + ACTIVE_WORK_KEY, + type ActiveWorkRuntime, + type ActiveWorkState, + endActiveWork, + renewActiveWork, + startActiveWork, +} from './vm-agent-container-active-work'; +import { + inspectStoredVmAgentContainerLifecycle, + type VmAgentContainerLifecycleInspection, + type VmAgentContainerLifecycleStatus, +} from './vm-agent-container-lifecycle'; +import { + loadRuntimeRecoveryContext, + persistRuntimeRecovered, + persistRuntimeRecovering, + persistRuntimeRecoveryFailed, + RUNTIME_RECOVERING_MESSAGE, + RUNTIME_RECOVERY_DEGRADED_MESSAGE, + RUNTIME_REQUEST_INTERRUPTED_MESSAGE, + RUNTIME_STOPPED_MESSAGE, + type RuntimeRecoveryCause, + type RuntimeRecoveryCode, + type RuntimeRecoveryState, + type RuntimeRecoveryTarget, + type RuntimeRecoveryTrigger, + toRuntimeRecoveryTarget, +} from './vm-agent-container-recovery'; +import { + interruptedRuntimeRequestResponse as interruptedRequestResponse, + isMissingSessionHostResponse, + isMutatingRuntimeRequest as isMutatingRequest, + persistRuntimeEnded, + persistRuntimeSleeping, + probeLiveRuntimeSession, + resolveRuntimeSettings, + runtimeRecoveryResponse as recoveryResponse, + runtimeResultResponse as resultResponse, +} from './vm-agent-container-runtime'; export const DEFAULT_CF_CONTAINER_SLEEP_AFTER = '1h'; +export const DEFAULT_CF_CONTAINER_PORT_READY_TIMEOUT_MS = 30_000; export const DEFAULT_CF_CONTAINER_ACTIVE_WORK_MAX_MS = 2 * 60 * 60 * 1000; export const DEFAULT_CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS = 5 * 60 * 1000; - +export const DEFAULT_CF_CONTAINER_RECOVERY_MAX_ATTEMPTS = 2; export interface VmAgentContainerLaunchConfig { nodeId: string; workspaceId: string; @@ -22,31 +61,28 @@ export interface VmAgentContainerLaunchConfig { controlPlaneUrl: string; vmAgentPort: number; } - export interface VmAgentContainerLaunchSecrets { nodeCallbackToken: string; } -type LifecycleStatus = 'launching' | 'running' | 'stopping' | 'stopped' | 'sleeping' | 'expired' | 'error'; - -type ActiveWorkStatus = 'active' | 'ended' | 'expired'; - -interface ActiveWorkState { - status: ActiveWorkStatus; - nodeId: string; - workspaceId: string; - agentSessionId: string; - reason: string; - activeSince: number; - lastRenewedAt: number; - deadlineAt: number; - endedAt?: number; - endReason?: string; +export interface VmAgentContainerRecoveryResult { + ok: boolean; + status: 'running' | 'recovering' | 'degraded' | 'stopped'; + code?: RuntimeRecoveryCode; + message?: string; } +type LifecycleStatus = VmAgentContainerLifecycleStatus; -const ACTIVE_WORK_KEY = 'activeWork'; +const RECOVERY_STATE_KEY = 'runtimeRecovery'; const KEEPALIVE_CALLBACK = 'renewActiveWorkKeepalive'; -const WAKE_DEGRADED_RESPONSE = 'Workspace woke with degraded snapshot restore; retry the prompt or fork from transcript history.'; +function stoppedRecoveryResult(): VmAgentContainerRecoveryResult { + return { + ok: false, + status: 'stopped', + code: 'RUNTIME_STOPPED', + message: RUNTIME_STOPPED_MESSAGE, + }; +} export class VmAgentContainer extends Container { defaultPort = 8080; @@ -54,20 +90,21 @@ 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(); + private lifecycleChain: 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); + const configuredPort = Number.parseInt( + env.CF_CONTAINER_VM_AGENT_PORT || env.SANDBOX_VM_AGENT_PORT || '', + 10 + ); if (Number.isFinite(configuredPort) && configuredPort > 0) { this.defaultPort = configuredPort; this.requiredPorts = [configuredPort]; } - this.sleepAfter = env.CF_CONTAINER_SLEEP_AFTER || env.SANDBOX_SLEEP_AFTER || DEFAULT_CF_CONTAINER_SLEEP_AFTER; + this.sleepAfter = + env.CF_CONTAINER_SLEEP_AFTER || env.SANDBOX_SLEEP_AFTER || DEFAULT_CF_CONTAINER_SLEEP_AFTER; } async launch( @@ -77,77 +114,188 @@ export class VmAgentContainer extends Container { await this.ctx.storage.put('launchConfig', config); await this.ctx.storage.put('lifecycleStatus', 'launching' satisfies LifecycleStatus); await this.ctx.storage.delete(ACTIVE_WORK_KEY); + await this.ctx.storage.delete(RECOVERY_STATE_KEY); await this.clearKeepaliveSchedule(); - await this.startAndWaitForPorts({ - ports: config.vmAgentPort, - startOptions: { - envVars: { - NODE_ROLE: 'standalone', - NODE_ID: config.nodeId, - WORKSPACE_ID: config.workspaceId, - PROJECT_ID: config.projectId, - CHAT_SESSION_ID: config.chatSessionId, - CONTROL_PLANE_URL: config.controlPlaneUrl, - CALLBACK_TOKEN: secrets.nodeCallbackToken, - REPOSITORY: config.repository, - BRANCH: config.branch, - WORKSPACE_DIR: config.workspaceDir, - CONTAINER_WORK_DIR: config.workspaceDir, - CONTAINER_MODE: 'false', - PORT_SCAN_ENABLED: 'false', - VM_AGENT_PORT: String(config.vmAgentPort), - VM_AGENT_PROTOCOL: 'http', - COOKIE_SECURE: 'true', - // Operator override for the standalone git partial-clone filter; - // the vm-agent defaults to blob:none when unset. - ...(this.env.CF_CONTAINER_CLONE_FILTER - ? { STANDALONE_CLONE_FILTER: this.env.CF_CONTAINER_CLONE_FILTER } - : {}), - }, - labels: { - nodeId: config.nodeId, - workspaceId: config.workspaceId, - runtime: 'cf-container', - }, - }, - cancellationOptions: { - portReadyTimeoutMS: this.getPortReadyTimeoutMs(), - }, - }); - - await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus); + try { + await this.startRuntime(config, secrets); + await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus); + } catch (error) { + await this.ctx.storage.put('lifecycleStatus', 'error' satisfies LifecycleStatus); + throw error; + } } async proxyHttp(request: Request, port?: number): Promise { - let state = await this.getState(); - const lifecycleStatus = await this.ctx.storage.get('lifecycleStatus'); - if (lifecycleStatus === 'sleeping') { - const wake = await this.ensureAwake(); - if (!wake.ok) { - return new Response(wake.message || WAKE_DEGRADED_RESPONSE, { status: 503 }); + const ready = await this.prepareForRequest(); + if (!ready.ok) return resultResponse(ready); + + const state = await this.getState(); + if (state.status === 'stopped' || state.status === 'stopped_with_code') { + const recovery = await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { kind: 'transport_interrupted', errorName: 'container_stopped' }, + promptDisposition: isMutatingRequest(request) ? 'manual_retry' : 'none', + }); + if (!recovery) { + return recoveryResponse('RUNTIME_STOPPED', RUNTIME_STOPPED_MESSAGE, 410); } - // 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(); + return interruptedRequestResponse(request); } + + try { + const response = await this.containerFetch(request, port ?? this.defaultPort); + if (await isMissingSessionHostResponse(response)) { + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { kind: 'missing_session_host', httpStatus: response.status }, + promptDisposition: isMutatingRequest(request) ? 'manual_retry' : 'none', + }); + return interruptedRequestResponse(request); + } + return response; + } catch (error) { + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { + kind: 'transport_interrupted', + errorName: error instanceof Error ? error.name : 'unknown', + }, + promptDisposition: isMutatingRequest(request) ? 'manual_retry' : 'none', + }); + return interruptedRequestResponse(request); + } + } + + override async fetch(request: Request): Promise { + const ready = await this.prepareForRequest(); + if (!ready.ok) return resultResponse(ready); + + const 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 }); + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { kind: 'transport_interrupted', errorName: 'container_stopped' }, + promptDisposition: 'none', + }); + return recoveryResponse('RUNTIME_RECOVERING', RUNTIME_RECOVERING_MESSAGE, 503); + } + + try { + return await super.fetch(switchPort(request, this.defaultPort)); + } catch (error) { + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { + kind: 'transport_interrupted', + errorName: error instanceof Error ? error.name : 'unknown', + }, + promptDisposition: 'none', + }); + return recoveryResponse('RUNTIME_RECOVERING', RUNTIME_RECOVERING_MESSAGE, 503); + } + } + + async resumeRuntime(agentSessionId?: string): Promise { + const lifecycle = await this.ctx.storage.get('lifecycleStatus'); + if (lifecycle === 'running') { + const config = await this.ctx.storage.get('launchConfig'); + const context = config + ? await loadRuntimeRecoveryContext(this.env, { + workspaceId: config.workspaceId, + preferredAgentSessionId: agentSessionId, + }) + : null; + if (!config || !context) { + return { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }; + } + try { + const live = await probeLiveRuntimeSession({ + env: this.env, + userId: context.userId, + nodeId: config.nodeId, + workspaceId: config.workspaceId, + agentSessionId: context.agentSessionId, + vmAgentPort: config.vmAgentPort, + containerFetch: (request, port) => this.containerFetch(request, port), + }); + if (live) { + const reconciled = await this.withLifecycleLock(async () => { + const current = await this.ctx.storage.get('lifecycleStatus'); + if (current !== 'running') return false; + await persistRuntimeRecovered( + this.env, + toRuntimeRecoveryTarget(config, context), + 'none' + ); + return true; + }); + if (reconciled) return { ok: true, status: 'running' }; + return this.ensureAwake(); + } + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { kind: 'missing_session_host', httpStatus: 404 }, + promptDisposition: 'none', + }); + } catch (error) { + await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { + kind: 'transport_interrupted', + errorName: error instanceof Error ? error.name : 'unknown', + }, + promptDisposition: 'none', + }); + } + } + return this.ensureAwake(); + } + + async markRequestInterrupted(input: { + method: string; + errorName: string; + }): Promise { + const mutating = !['GET', 'HEAD', 'OPTIONS'].includes(input.method.toUpperCase()); + const state = await this.beginUnexpectedRecovery({ + trigger: 'request', + cause: { kind: 'transport_interrupted', errorName: input.errorName }, + promptDisposition: mutating ? 'manual_retry' : 'none', + }); + if (!state) { + const status = await this.ctx.storage.get('lifecycleStatus'); + if (status === 'stopped' || status === 'stopping') { + return stoppedRecoveryResult(); + } } - return this.containerFetch(request, port ?? this.defaultPort); + return { + ok: false, + status: 'recovering', + code: mutating ? 'RUNTIME_REQUEST_INTERRUPTED' : 'RUNTIME_RECOVERING', + message: mutating ? RUNTIME_REQUEST_INTERRUPTED_MESSAGE : RUNTIME_RECOVERING_MESSAGE, + }; } async stopForUser(): Promise { await this.markActiveWorkEnded('user_stop'); - await this.ctx.storage.put('lifecycleStatus', 'stopping' satisfies LifecycleStatus); + await this.withLifecycleLock(async () => { + await this.ctx.storage.put('lifecycleStatus', 'stopping' satisfies LifecycleStatus); + await this.ctx.storage.delete(RECOVERY_STATE_KEY); + }); await this.stop(); } async destroyForUser(): Promise { await this.markActiveWorkEnded('user_destroy'); - await this.ctx.storage.put('lifecycleStatus', 'stopping' satisfies LifecycleStatus); + await this.withLifecycleLock(async () => { + await this.ctx.storage.put('lifecycleStatus', 'stopping' satisfies LifecycleStatus); + await this.ctx.storage.delete(RECOVERY_STATE_KEY); + }); await this.destroy(); } @@ -157,131 +305,61 @@ export class VmAgentContainer extends Container { reason: string; }): Promise { const config = await this.ctx.storage.get('launchConfig'); - const now = Date.now(); - const nodeId = config?.nodeId ?? input.workspaceId; - const activeWork: ActiveWorkState = { - status: 'active', - nodeId, - workspaceId: input.workspaceId, - agentSessionId: input.agentSessionId, - reason: input.reason, - activeSince: now, - lastRenewedAt: now, - deadlineAt: now + this.getActiveWorkMaxMs(), - }; - this.renewActivityTimeout(); - await this.ctx.storage.put(ACTIVE_WORK_KEY, activeWork); - await this.replaceKeepaliveSchedule(this.getKeepaliveRenewIntervalMs()); - log.info('vm_agent_container_active_work_started', { - nodeId, - workspaceId: input.workspaceId, - agentSessionId: input.agentSessionId, - reason: input.reason, - activeSince: new Date(now).toISOString(), - deadlineAt: new Date(activeWork.deadlineAt).toISOString(), - }); + await startActiveWork(this.activeWorkRuntime(), config?.nodeId ?? input.workspaceId, input); } async markActiveWorkEnded(reason: string): Promise { - const activeWork = await this.ctx.storage.get(ACTIVE_WORK_KEY); - if (!activeWork || activeWork.status !== 'active') { - await this.clearKeepaliveSchedule(); - return; - } - const now = Date.now(); - await this.ctx.storage.put(ACTIVE_WORK_KEY, { - ...activeWork, - status: 'ended', - endedAt: now, - endReason: reason, - } satisfies ActiveWorkState); - await this.clearKeepaliveSchedule(); - log.info('vm_agent_container_active_work_ended', { - nodeId: activeWork.nodeId, - workspaceId: activeWork.workspaceId, - agentSessionId: activeWork.agentSessionId, - reason, - activeSince: new Date(activeWork.activeSince).toISOString(), - lastRenewedAt: new Date(activeWork.lastRenewedAt).toISOString(), - endedAt: new Date(now).toISOString(), - }); + await endActiveWork(this.activeWorkRuntime(), reason); } - async renewActiveWorkKeepalive(): Promise { - const activeWork = await this.ctx.storage.get(ACTIVE_WORK_KEY); - if (!activeWork || activeWork.status !== 'active') { - await this.clearKeepaliveSchedule(); - return; - } - const now = Date.now(); - if (now >= activeWork.deadlineAt) { - await this.ctx.storage.put(ACTIVE_WORK_KEY, { - ...activeWork, - status: 'expired', - endedAt: now, - endReason: 'keepalive_deadline_exceeded', - } satisfies ActiveWorkState); - await this.clearKeepaliveSchedule(); - log.warn('vm_agent_container_active_work_keepalive_expired', { - nodeId: activeWork.nodeId, - workspaceId: activeWork.workspaceId, - agentSessionId: activeWork.agentSessionId, - activeSince: new Date(activeWork.activeSince).toISOString(), - lastRenewedAt: new Date(activeWork.lastRenewedAt).toISOString(), - deadlineAt: new Date(activeWork.deadlineAt).toISOString(), - }); - return; - } - - this.renewActivityTimeout(); - await this.ctx.storage.put(ACTIVE_WORK_KEY, { - ...activeWork, - lastRenewedAt: now, - } satisfies ActiveWorkState); - await this.replaceKeepaliveSchedule(this.getKeepaliveRenewIntervalMs()); - log.debug('vm_agent_container_active_work_keepalive_renewed', { - nodeId: activeWork.nodeId, - workspaceId: activeWork.workspaceId, - agentSessionId: activeWork.agentSessionId, - activeSince: new Date(activeWork.activeSince).toISOString(), - lastRenewedAt: new Date(now).toISOString(), - deadlineAt: new Date(activeWork.deadlineAt).toISOString(), - }); + async inspectLifecycle(): Promise { + return inspectStoredVmAgentContainerLifecycle(this.ctx.storage, RECOVERY_STATE_KEY, ACTIVE_WORK_KEY); } - override async fetch(request: Request): Promise { - let state = await this.getState(); - const lifecycleStatus = await this.ctx.storage.get('lifecycleStatus'); - if (lifecycleStatus === 'sleeping') { - 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 }); - } - return super.fetch(switchPort(request, this.defaultPort)); + async renewActiveWorkKeepalive(): Promise { + await renewActiveWork(this.activeWorkRuntime()); } override async onStart(): Promise { - await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus); + log.info('vm_agent_container_runtime_started', {}); } - override async onStop(params: { exitCode: number; reason: 'exit' | 'runtime_signal' }): Promise { + override async onStop(params: { + exitCode: number; + reason: 'exit' | 'runtime_signal'; + }): Promise { const status = await this.ctx.storage.get('lifecycleStatus'); - if (status === 'expired' || status === 'sleeping') { + if ( + status === 'expired' || + status === 'sleeping' || + status === 'recovering' || + status === 'waking' || + status === 'restoring' || + status === 'degraded' || + status === 'stopped' + ) { return; } - const explicitStop = status === 'stopping'; - await this.markRuntimeEnded( - explicitStop ? 'stopped' : 'error', - explicitStop ? 'Container stopped by user request' : `Container stopped: ${params.reason} (${params.exitCode})` - ); - await this.ctx.storage.put('lifecycleStatus', explicitStop ? 'stopped' : 'error'); + if (status === 'stopping') { + await this.markRuntimeEnded('stopped', 'Container stopped by user request'); + await this.ctx.storage.put('lifecycleStatus', 'stopped' satisfies LifecycleStatus); + return; + } + if (status === 'launching') { + await this.markRuntimeEnded('error', 'Instant runtime failed during initial launch.'); + await this.ctx.storage.put('lifecycleStatus', 'error' satisfies LifecycleStatus); + return; + } + + await this.beginUnexpectedRecovery({ + trigger: 'stop', + cause: { + kind: 'container_stop', + reason: params.reason, + exitCode: params.exitCode, + }, + promptDisposition: 'none', + }); } override async onActivityExpired(): Promise { @@ -296,47 +374,111 @@ export class VmAgentContainer extends Container { } override async onError(error: unknown): Promise { - await this.markRuntimeEnded( - 'error', - error instanceof Error ? `Container error: ${error.message}` : `Container error: ${String(error)}` - ); - await this.ctx.storage.put('lifecycleStatus', 'error' satisfies LifecycleStatus); - } - - private getPortReadyTimeoutMs(): number { - const raw = this.env.CF_CONTAINER_PORT_READY_TIMEOUT_MS || this.env.SANDBOX_EXEC_TIMEOUT_MS; - const parsed = raw ? Number.parseInt(raw, 10) : 30_000; - return Number.isFinite(parsed) && parsed > 0 ? parsed : 30_000; - } + const status = await this.ctx.storage.get('lifecycleStatus'); + if (status === 'stopping') { + await this.markRuntimeEnded('stopped', 'Container stopped by user request'); + await this.ctx.storage.put('lifecycleStatus', 'stopped' satisfies LifecycleStatus); + return; + } + if (status === 'launching') { + await this.markRuntimeEnded('error', 'Instant runtime failed during initial launch.'); + await this.ctx.storage.put('lifecycleStatus', 'error' satisfies LifecycleStatus); + return; + } + if (status === 'stopped' || status === 'sleeping' || status === 'expired') return; - private getActiveWorkMaxMs(): number { - const raw = this.env.CF_CONTAINER_ACTIVE_WORK_MAX_MS; - const parsed = raw ? Number.parseInt(raw, 10) : DEFAULT_CF_CONTAINER_ACTIVE_WORK_MAX_MS; - return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_CF_CONTAINER_ACTIVE_WORK_MAX_MS; + await this.beginUnexpectedRecovery({ + trigger: 'error', + cause: { + kind: 'container_error', + errorName: error instanceof Error ? error.name : 'unknown', + }, + promptDisposition: 'none', + }); } - private getKeepaliveRenewIntervalMs(): number { - const raw = this.env.CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS; - const parsed = raw ? Number.parseInt(raw, 10) : DEFAULT_CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS; - return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS; + private async prepareForRequest(): Promise { + const status = await this.ctx.storage.get('lifecycleStatus'); + if (status === 'running' || status === 'launching') { + return { ok: true, status: 'running' }; + } + if (status === 'stopping' || status === 'stopped') { + return stoppedRecoveryResult(); + } + return this.ensureAwake(); } - /** - * 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 }> { + private async ensureAwake(): Promise { const run = this.wakeChain.then(async () => { - const status = await this.ctx.storage.get('lifecycleStatus'); - if (status !== 'sleeping') { - return { ok: true }; + let status = await this.ctx.storage.get('lifecycleStatus'); + if (status === 'running' || status === 'launching') { + return { ok: true, status: 'running' } satisfies VmAgentContainerRecoveryResult; + } + if (status === 'stopping' || status === 'stopped') { + return stoppedRecoveryResult(); + } + + let recovery: RuntimeRecoveryState | null | undefined = + await this.ctx.storage.get(RECOVERY_STATE_KEY); + if (!recovery) { + const trigger: RuntimeRecoveryTrigger = status === 'sleeping' ? 'idle' : 'error'; + const cause: RuntimeRecoveryCause = + status === 'sleeping' + ? { kind: 'idle_sleep' } + : { kind: 'container_error', errorName: 'legacy_terminal_state' }; + recovery = await this.beginUnexpectedRecovery({ + trigger, + cause, + promptDisposition: 'none', + }); + status = await this.ctx.storage.get('lifecycleStatus'); + } + + // A concurrent stopForUser()/destroyForUser() runs on the same + // lifecycleChain as beginUnexpectedRecovery above, so an explicit stop can + // land between beginUnexpectedRecovery returning and this unlocked + // continuation — flipping lifecycleStatus to 'stopping' and deleting the + // recovery record. 'stopping' is terminal exactly like 'stopped'; both must + // short-circuit here so we never re-arm recovery or wake a container the + // user just stopped ("explicit stop is terminal, never recover"). + if (status === 'stopping' || status === 'stopped') { + return stoppedRecoveryResult(); + } + if (!recovery) { + return { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + } satisfies VmAgentContainerRecoveryResult; + } + + // Once recovery is exhausted, exhaustRecovery() has already reconciled D1 + // and set lifecycleStatus 'error'. Return the same sanitized degraded result + // without re-running the full reconciliation batch on every later request. + if (recovery.phase === 'exhausted') { + return { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + } satisfies VmAgentContainerRecoveryResult; + } + + if (recovery.attempts >= this.getRuntimeSettings().recoveryMaxAttempts) { + return this.exhaustRecovery(recovery); } - return this.wakeFromSnapshot(); + + recovery = { + ...recovery, + phase: 'waking', + attempts: recovery.attempts + 1, + updatedAt: Date.now(), + }; + await this.ctx.storage.put(RECOVERY_STATE_KEY, recovery); + await this.ctx.storage.put('lifecycleStatus', 'waking' satisfies LifecycleStatus); + return this.wakeFromSnapshot(recovery); }); - // 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 @@ -344,139 +486,333 @@ export class VmAgentContainer extends Container { return run; } - private async wakeFromSnapshot(): Promise<{ ok: boolean; message?: string }> { + private async beginUnexpectedRecovery(input: { + trigger: RuntimeRecoveryTrigger; + cause: RuntimeRecoveryCause; + promptDisposition: RuntimeRecoveryState['promptDisposition']; + }): Promise { + return this.withLifecycleLock(async () => { + const lifecycle = await this.ctx.storage.get('lifecycleStatus'); + if (lifecycle === 'stopping' || lifecycle === 'stopped' || lifecycle === 'expired') { + return null; + } + + const existing = await this.ctx.storage.get(RECOVERY_STATE_KEY); + if (existing) { + if ( + input.promptDisposition === 'manual_retry' && + existing.promptDisposition !== 'manual_retry' + ) { + const promoted = { + ...existing, + promptDisposition: 'manual_retry' as const, + updatedAt: Date.now(), + }; + await this.ctx.storage.put(RECOVERY_STATE_KEY, promoted); + return promoted; + } + return existing; + } + + const config = await this.ctx.storage.get('launchConfig'); + if (!config) return null; + const activeWork = await this.ctx.storage.get(ACTIVE_WORK_KEY); + const context = await loadRuntimeRecoveryContext(this.env, { + workspaceId: config.workspaceId, + preferredAgentSessionId: activeWork?.agentSessionId, + }); + if (!context) return null; + + const now = Date.now(); + const recovery: RuntimeRecoveryState = { + version: 1, + phase: 'pending', + trigger: input.trigger, + cause: input.cause, + attempts: 0, + promptDisposition: input.promptDisposition, + agentSessionId: context.agentSessionId, + startedAt: now, + updatedAt: now, + }; + await this.ctx.storage.put(RECOVERY_STATE_KEY, recovery); + await this.ctx.storage.put('lifecycleStatus', 'recovering' satisfies LifecycleStatus); + await this.clearKeepaliveSchedule(); + await persistRuntimeRecovering(this.env, toRuntimeRecoveryTarget(config, context)); + log.warn('vm_agent_container_recovery_started', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + agentSessionId: context.agentSessionId, + trigger: input.trigger, + cause: input.cause, + promptDisposition: input.promptDisposition, + }); + return recovery; + }); + } + + private async wakeFromSnapshot( + recovery: RuntimeRecoveryState + ): Promise { 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.' }; + if (!config) return this.degradeRecovery(recovery, 'launch'); + + // loadRuntimeRecoveryContext reads D1; a query failure must degrade through the + // sanitized recovery path, not escape as an uncaught 500. It sits outside the + // restore try below, so guard it explicitly here. + let context: Awaited>; + try { + context = await loadRuntimeRecoveryContext(this.env, { + workspaceId: config.workspaceId, + preferredAgentSessionId: recovery.agentSessionId, + }); + } catch (error) { + log.warn('vm_agent_container_recovery_context_failed', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + errorName: error instanceof Error ? error.name : 'unknown', + }); + return this.degradeRecovery(recovery, 'unexpected'); } - 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.' }; + if (!context) return this.degradeRecovery(recovery, 'unexpected'); + const target = toRuntimeRecoveryTarget(config, context); + + try { + const nodeCallbackToken = await signNodeCallbackToken(config.nodeId, this.env); + await this.startRuntime(config, { nodeCallbackToken }); + } catch (error) { + log.warn('vm_agent_container_recovery_launch_failed', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + errorName: error instanceof Error ? error.name : 'unknown', + }); + return this.degradeRecovery(recovery, 'launch', target); } - log.info('vm_agent_container_wake_started', { - nodeId: config.nodeId, - workspaceId: config.workspaceId, - chatSessionId: workspace.chatSessionId, - agentSessionId: agentSession.id, - }); + const restoring = { ...recovery, phase: 'restoring' as const, updatedAt: Date.now() }; + await this.ctx.storage.put(RECOVERY_STATE_KEY, restoring); + await this.ctx.storage.put('lifecycleStatus', 'restoring' satisfies LifecycleStatus); - 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, + try { + const workspaceCallbackToken = await signCallbackToken(config.workspaceId, this.env); + const { token } = await signNodeManagementToken( + context.userId, + config.nodeId, + config.workspaceId, + this.env + ); + const restoreUrl = new URL( + `http://localhost:${config.vmAgentPort}/workspaces/${config.workspaceId}/agent-sessions/${context.agentSessionId}/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: context.chatSessionId, + runtime: 'cf-container', + agentType: context.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.' }; + config.vmAgentPort + ); + const restoreBody = await restoreResponse.text().catch(() => ''); + if (!restoreResponse.ok) { + return this.degradeRecovery(restoring, 'restore_http', target, restoreResponse.status); + } + let restoreStatus = ''; + try { + const parsed = JSON.parse(restoreBody) as { status?: unknown }; + restoreStatus = typeof parsed.status === 'string' ? parsed.status : ''; + } catch { + // An explicit restored status is required before D1 can become running. + } + if (restoreStatus !== 'restored') { + return this.degradeRecovery(restoring, 'restore_status', target); + } + + const completed = await this.withLifecycleLock(async () => { + const lifecycle = await this.ctx.storage.get('lifecycleStatus'); + if (lifecycle === 'stopping' || lifecycle === 'stopped') return false; + await persistRuntimeRecovered(this.env, target, restoring.promptDisposition); + await this.ctx.storage.delete(RECOVERY_STATE_KEY); + await this.ctx.storage.put('lifecycleStatus', 'running' satisfies LifecycleStatus); + return true; + }); + if (!completed) { + // An explicit stop crossed the restore after startRuntime() already + // launched a fresh container. Tear it down so the just-stopped session + // does not leak compute until sleepAfter. this.stop() is idempotent/safe + // if stopForUser() already issued it. + await this.stop().catch(() => undefined); + return stoppedRecoveryResult(); + } + log.info('vm_agent_container_recovery_completed', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + agentSessionId: context.agentSessionId, + attempts: restoring.attempts, + promptDisposition: restoring.promptDisposition, + }); + return { ok: true, status: 'running' }; + } catch (error) { + log.warn('vm_agent_container_recovery_restore_failed', { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + errorName: error instanceof Error ? error.name : 'unknown', + }); + return this.degradeRecovery(restoring, 'unexpected', target); } - 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. + } + + private async degradeRecovery( + recovery: RuntimeRecoveryState, + kind: NonNullable['kind'], + target?: RuntimeRecoveryTarget, + httpStatus?: number + ): Promise { + const degraded: RuntimeRecoveryState = { + ...recovery, + phase: 'degraded', + updatedAt: Date.now(), + lastFailure: { kind, ...(httpStatus === undefined ? {} : { httpStatus }) }, + }; + const applied = await this.withLifecycleLock(async () => { + const lifecycle = await this.ctx.storage.get('lifecycleStatus'); + if (lifecycle === 'stopping' || lifecycle === 'stopped') return false; + await this.ctx.storage.put(RECOVERY_STATE_KEY, degraded); + await this.ctx.storage.put('lifecycleStatus', 'degraded' satisfies LifecycleStatus); + return true; + }); + if (!applied) { + // A stop crossed the wake after startRuntime() launched a container (the + // restore_http / restore_status / restore-catch degrade calls all run + // post-launch). Tear it down before returning terminal. Idempotent/safe if + // stopForUser() already stopped it. + await this.stop().catch(() => undefined); + return stoppedRecoveryResult(); } - if (restoreStatus !== 'restored') { - const message = restoreBody || 'Session restore did not report restored status.'; - await this.markWakeDegraded(config, message); - return { ok: false, message }; + await this.stop().catch(() => undefined); + + if (degraded.attempts >= this.getRuntimeSettings().recoveryMaxAttempts) { + return this.exhaustRecovery(degraded, target); } + return { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_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, + private async exhaustRecovery( + recovery: RuntimeRecoveryState, + providedTarget?: RuntimeRecoveryTarget + ): Promise { + const applied = await this.withLifecycleLock(async () => { + const lifecycle = await this.ctx.storage.get('lifecycleStatus'); + if (lifecycle === 'stopping' || lifecycle === 'stopped') return false; + const config = await this.ctx.storage.get('launchConfig'); + let target = providedTarget; + if (!target && config) { + const context = await loadRuntimeRecoveryContext(this.env, { + workspaceId: config.workspaceId, + preferredAgentSessionId: recovery.agentSessionId, + }); + if (context) target = toRuntimeRecoveryTarget(config, context); + } + const exhausted = { ...recovery, phase: 'exhausted' as const, updatedAt: Date.now() }; + await this.ctx.storage.put(RECOVERY_STATE_KEY, exhausted); + await this.ctx.storage.put('lifecycleStatus', 'error' satisfies LifecycleStatus); + if (target) await persistRuntimeRecoveryFailed(this.env, target); + return true; }); - return { ok: true }; + if (!applied) { + // A stop crossed exhaustion. If a fresh container was launched earlier in + // this wake, tear it down before returning terminal. Idempotent/safe if + // stopForUser() already stopped it. + await this.stop().catch(() => undefined); + return stoppedRecoveryResult(); + } + return { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }; } - 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 startRuntime( + config: VmAgentContainerLaunchConfig, + secrets: VmAgentContainerLaunchSecrets + ): Promise { + await this.startAndWaitForPorts({ + ports: config.vmAgentPort, + startOptions: { + envVars: { + NODE_ROLE: 'standalone', + NODE_ID: config.nodeId, + WORKSPACE_ID: config.workspaceId, + PROJECT_ID: config.projectId, + CHAT_SESSION_ID: config.chatSessionId, + CONTROL_PLANE_URL: config.controlPlaneUrl, + CALLBACK_TOKEN: secrets.nodeCallbackToken, + REPOSITORY: config.repository, + BRANCH: config.branch, + WORKSPACE_DIR: config.workspaceDir, + CONTAINER_WORK_DIR: config.workspaceDir, + CONTAINER_MODE: 'false', + PORT_SCAN_ENABLED: 'false', + VM_AGENT_PORT: String(config.vmAgentPort), + VM_AGENT_PROTOCOL: 'http', + COOKIE_SECURE: 'true', + ...(this.env.CF_CONTAINER_CLONE_FILTER + ? { STANDALONE_CLONE_FILTER: this.env.CF_CONTAINER_CLONE_FILTER } + : {}), + }, + labels: { + nodeId: config.nodeId, + workspaceId: config.workspaceId, + runtime: 'cf-container', + }, + }, + cancellationOptions: { portReadyTimeoutMS: this.getRuntimeSettings().portReadyTimeoutMs }, + }); + } + + private withLifecycleLock(operation: () => Promise): Promise { + const run = this.lifecycleChain.then(operation); + this.lifecycleChain = run.then( + () => undefined, + () => undefined + ); + return run; + } + + private activeWorkRuntime(): ActiveWorkRuntime { + const settings = this.getRuntimeSettings(); + return { + storage: this.ctx.storage, + activeWorkMaxMs: settings.activeWorkMaxMs, + renewIntervalMs: settings.keepaliveRenewIntervalMs, + renewActivityTimeout: () => this.renewActivityTimeout(), + replaceSchedule: (delayMs) => this.replaceKeepaliveSchedule(delayMs), + clearSchedule: () => this.clearKeepaliveSchedule(), + }; + } + + private getRuntimeSettings() { + return resolveRuntimeSettings(this.env, { + portReadyTimeoutMs: DEFAULT_CF_CONTAINER_PORT_READY_TIMEOUT_MS, + activeWorkMaxMs: DEFAULT_CF_CONTAINER_ACTIVE_WORK_MAX_MS, + keepaliveRenewIntervalMs: DEFAULT_CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS, + recoveryMaxAttempts: DEFAULT_CF_CONTAINER_RECOVERY_MAX_ATTEMPTS, }); } @@ -491,97 +827,24 @@ export class VmAgentContainer extends Container { private async markRuntimeSleeping(message: string): Promise { const config = await this.ctx.storage.get('launchConfig'); - if (!config) { - return; - } - + if (!config) return; await this.markActiveWorkEnded('container_idle_sleeping'); - - const now = new Date().toISOString(); - const db = drizzle(this.env.DATABASE, { schema }); - - await db - .update(schema.nodes) - .set({ - status: 'sleeping', - healthStatus: 'unhealthy', - errorMessage: null, - updatedAt: now, - }) - .where(eq(schema.nodes.id, config.nodeId)); - - await db - .update(schema.workspaces) - .set({ - status: 'sleeping', - errorMessage: null, - updatedAt: now, - }) - .where(eq(schema.workspaces.id, config.workspaceId)); - - await db - .update(schema.agentSessions) - .set({ - status: 'sleeping', - errorMessage: null, - updatedAt: now, - }) - .where(eq(schema.agentSessions.workspaceId, config.workspaceId)); - + await persistRuntimeSleeping(this.env, config); log.info('vm_agent_container_runtime_sleeping', { nodeId: config.nodeId, workspaceId: config.workspaceId, - status: 'sleeping', message, }); } - private async markRuntimeEnded(status: Exclude, message: string): Promise { + private async markRuntimeEnded(status: 'stopped' | 'error', message: string): Promise { const config = await this.ctx.storage.get('launchConfig'); - if (!config) { - return; - } - - const now = new Date().toISOString(); - const db = drizzle(this.env.DATABASE, { schema }); - const workspaceStatus = status === 'stopped' ? 'stopped' : 'error'; - const agentStatus = status === 'stopped' ? 'stopped' : 'error'; - const nodeStatus = status === 'stopped' ? 'stopped' : 'error'; - - await db - .update(schema.nodes) - .set({ - status: nodeStatus, - healthStatus: 'unhealthy', - errorMessage: status === 'stopped' ? null : message, - updatedAt: now, - }) - .where(eq(schema.nodes.id, config.nodeId)); - - await db - .update(schema.workspaces) - .set({ - status: workspaceStatus, - errorMessage: status === 'stopped' ? null : message, - updatedAt: now, - }) - .where(eq(schema.workspaces.id, config.workspaceId)); - - await db - .update(schema.agentSessions) - .set({ - status: agentStatus, - stoppedAt: now, - errorMessage: status === 'stopped' ? null : message, - updatedAt: now, - }) - .where(eq(schema.agentSessions.workspaceId, config.workspaceId)); - + if (!config) return; + await persistRuntimeEnded(this.env, config, status, message); log.warn('vm_agent_container_runtime_ended', { nodeId: config.nodeId, workspaceId: config.workspaceId, status, - message, }); } } diff --git a/apps/api/src/env.ts b/apps/api/src/env.ts index f0b06806a..a0fda3ddb 100644 --- a/apps/api/src/env.ts +++ b/apps/api/src/env.ts @@ -836,6 +836,8 @@ export interface Env extends WebhookTriggerEnv, TaskRecoveryEnv { 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_RECOVERY_MAX_ATTEMPTS?: string; // Max snapshot restore attempts before terminal reconciliation (default: 2) + INSTANT_STALE_CALLBACK_MARGIN_MS?: string; // Freshness margin for rejecting destructive callbacks from superseded Instant containers (default: 60000) CF_CONTAINER_CREATE_WORKSPACE_TIMEOUT_MS?: string; // Max time for the synchronous standalone create-workspace request incl. clone (default: 120000) CF_CONTAINER_CLONE_FILTER?: string; // Git partial-clone filter passed to the container as STANDALONE_CLONE_FILTER (vm-agent default: blob:none; "off" disables) CF_CONTAINER_WORKSPACE_BASE_DIR?: string; // Base checkout dir inside raw container (default: /workspaces) diff --git a/apps/api/src/routes/_stale-callback-guard.ts b/apps/api/src/routes/_stale-callback-guard.ts new file mode 100644 index 000000000..0f5f1be4b --- /dev/null +++ b/apps/api/src/routes/_stale-callback-guard.ts @@ -0,0 +1,89 @@ +import { decodeJwt } from 'jose'; + +import type { Env } from '../env'; + +/** + * Staleness guard for VM-agent → control-plane DESTRUCTIVE callbacks (S2). + * + * Instant (cf-container) sessions recover by replacing a dead container with a + * NEW generation under the SAME nodeId. Callback tokens are stateless RS256 JWTs + * with no jti/revocation, minted fresh per cold wake, so a superseded (dead) + * container's token stays valid. A mid-turn rollout can make the OLD container + * fire a fire-and-forget `error` / `failed` callback that lands AFTER the DO has + * recovered a new container to running. Left unguarded, that late callback + * regresses the freshly recovered, healthy session to error/failed. + * + * We reject a destructive callback only when it is provably OLDER than a + * completed D1 reconciliation of the target row: the row's `updated_at` (written + * by the Durable Object recovery via `persistRuntimeRecovered` / + * `persistRuntimeRecovering`) is newer than the callback token's `iat` by more + * than a configurable freshness margin. + * + * Fail-open by construction: any ambiguity (non-Instant runtime, missing iat, + * unparseable timestamp, gap within the margin) PROCESSES the callback, so a + * genuinely crashed CURRENT container still fails its session (requirement #2). + */ + +/** + * Default freshness margin (ms). Chosen to sit comfortably ABOVE the worst-case + * same-generation gap (a recovered container mints its token just BEFORE the + * snapshot restore that later writes `updated_at`, so its own genuine later + * error has an `updated_at - iat` gap bounded by the restore duration — a few + * seconds) yet BELOW a typical superseded-container lifetime (a real user turn), + * so a stale generation's callback is reliably rejected. + */ +export const DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS = 60_000; + +export function getInstantStaleCallbackMarginMs(env: Env): number { + const raw = env.INSTANT_STALE_CALLBACK_MARGIN_MS; + const parsed = raw ? Number.parseInt(raw, 10) : Number.NaN; + return Number.isFinite(parsed) && parsed >= 0 + ? parsed + : DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS; +} + +/** + * Read the `iat` (issued-at) claim from an ALREADY-VERIFIED callback token and + * return it in milliseconds. The token MUST have been verified by + * `verifyCallbackToken` first — `decodeJwt` does not verify the signature; it is + * used here only to read a claim `verifyCallbackToken` does not surface. + * + * Returns null when the token cannot be decoded or has no numeric `iat`. + */ +export function callbackTokenIssuedAtMs(token: string): number | null { + try { + const claims = decodeJwt(token); + return typeof claims.iat === 'number' ? claims.iat * 1000 : null; + } catch { + return null; + } +} + +export interface SupersededInstantCallbackInput { + /** Runtime of the node backing the target row (`nodes.runtime`). */ + runtime: string | null | undefined; + /** ISO-8601 `updated_at` of the D1 row a completed recovery reconciles. */ + rowUpdatedAt: string | null | undefined; + /** Callback token `iat` in ms (see {@link callbackTokenIssuedAtMs}). */ + tokenIssuedAtMs: number | null; + /** Freshness margin in ms (see {@link getInstantStaleCallbackMarginMs}). */ + marginMs: number; +} + +/** + * True when a destructive callback provably originates from a superseded Instant + * container generation (older than a completed recovery). See file header. + */ +export function isSupersededInstantCallback(input: SupersededInstantCallbackInput): boolean { + // Only Instant (cf-container) has the same-nodeId generation-replacement race. + // VM-runtime nodes never swap generations under a still-valid old token. + if (input.runtime !== 'cf-container') return false; + // Cannot establish the token's generation age → fail open (process the error). + if (input.tokenIssuedAtMs === null) return false; + if (!input.rowUpdatedAt) return false; + const rowMs = Date.parse(input.rowUpdatedAt); + if (!Number.isFinite(rowMs)) return false; + // Reconciliation strictly newer than the token by more than the margin ⇒ a + // recovery completed after this token's generation cold-woke ⇒ superseded. + return rowMs > input.tokenIssuedAtMs + input.marginMs; +} diff --git a/apps/api/src/routes/chat-workspace-resolver.ts b/apps/api/src/routes/chat-workspace-resolver.ts index 928723f52..332f7c1a7 100644 --- a/apps/api/src/routes/chat-workspace-resolver.ts +++ b/apps/api/src/routes/chat-workspace-resolver.ts @@ -1,10 +1,9 @@ /** - * Tenant-scoped resolution of the live workspace that bridges a chat session - * to its running VM. Extracted from routes/chat.ts to keep that file under the - * 800-line limit (.claude/rules/18) and to make the security-critical resolver - * independently testable against real D1. + * Tenant-scoped resolution of the workspace bridging a chat session to its + * runtime. Recovery states are admitted only for cf-container nodes; VM + * sessions preserve their fail-fast dead-node guard. */ -import { and, eq, inArray } from 'drizzle-orm'; +import { and, eq, inArray, or } from 'drizzle-orm'; import type { drizzle } from 'drizzle-orm/d1'; import * as schema from '../db/schema'; @@ -13,32 +12,28 @@ import { errors } from '../middleware/error'; type ChatDb = ReturnType>; +export interface ChatWorkspaceRuntime { + id: string; + nodeId: string | null; + nodeStatus: string | null; + nodeRuntime: string | null; +} + /** - * Resolve the live workspace that bridges a chat session to its running VM, - * scoped to the owning project AND user. - * - * Security: the workspace bridge MUST NOT be resolvable by `chatSessionId` - * alone. Without project + user scoping, any authenticated caller who supplies - * (or guesses/replays) a sessionId belonging to another tenant can drive the - * VM agent for a workspace they do not own — an IDOR in the same class as past - * cross-tenant leaks. We resolve via the narrowest canonical chat-scoped - * identifier (`chatSessionId`) AND enforce ownership in the query WHERE clause - * (see .claude/rules/06 canonical session routing, .claude/rules/11 identity - * validation). A post-query defence-in-depth assertion rejects any row whose - * ownership does not match the caller, guarding against a future WHERE-clause - * regression (typo/refactor/ORM bug) per .claude/rules/28. - * - * Returns null when no matching active workspace exists for this owner. + * Resolve the workspace bridge, scoped to the owning project and user. + * `chatSessionId` alone is never sufficient: retaining project/user predicates + * is the defence against cross-tenant workspace control. */ export async function resolveLiveWorkspaceForSession( db: ChatDb, { projectId, sessionId, userId }: { projectId: string; sessionId: string; userId: string } -): Promise<{ id: string; nodeId: string | null; nodeStatus: string | null } | null> { +): Promise { const [workspace] = await db .select({ id: schema.workspaces.id, nodeId: schema.workspaces.nodeId, nodeStatus: schema.nodes.status, + nodeRuntime: schema.nodes.runtime, userId: schema.workspaces.userId, projectId: schema.workspaces.projectId, }) @@ -49,16 +44,16 @@ export async function resolveLiveWorkspaceForSession( eq(schema.workspaces.chatSessionId, sessionId), eq(schema.workspaces.projectId, projectId), eq(schema.workspaces.userId, userId), - inArray(schema.workspaces.status, ['running', 'recovery', 'sleeping']) + or( + inArray(schema.workspaces.status, ['running', 'recovery', 'sleeping']), + and(eq(schema.nodes.runtime, 'cf-container'), eq(schema.workspaces.status, 'error')) + ) ) ) .limit(1); - if (!workspace) { - return null; - } + if (!workspace) return null; - // Defence-in-depth: reject any row whose ownership doesn't match the caller. if (workspace.userId !== userId || workspace.projectId !== projectId) { log.error('chat: workspace ownership mismatch on session bridge', { sessionId, @@ -72,53 +67,44 @@ export async function resolveLiveWorkspaceForSession( return null; } - return { id: workspace.id, nodeId: workspace.nodeId, nodeStatus: workspace.nodeStatus }; + return { + id: workspace.id, + nodeId: workspace.nodeId, + nodeStatus: workspace.nodeStatus, + nodeRuntime: workspace.nodeRuntime, + }; } /** - * 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 - * handlers: workspace bridge resolution (project + user scoped), node-liveness - * guards, and the defence-in-depth user-scoped agent-session lookup (backed by - * idx_agent_sessions_ws_user_status). Throws the same fail-fast errors both - * handlers relied on so callers can forward straight to the VM agent. - * - * @throws notFound when no active workspace/node or running agent session exists - * @throws conflict when the workspace node is no longer running + * Resolve the workspace and resumable agent session for a chat action. A + * cf-container Durable Object owns replacement classification, so recoverable + * D1 states must reach it. Other runtimes still require a running node. */ export async function resolveLiveAgentSessionForChat( db: ChatDb, { projectId, sessionId, userId }: { projectId: string; sessionId: string; userId: 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 - // task timeout), its DNS record is cleaned up but the workspace may still be - // marked 'running' in D1. Without this check, the request to the VM agent - // would hit the wildcard DNS record and loop back to this Worker (404). +): Promise<{ + workspace: { id: string; nodeId: string; nodeStatus: string; nodeRuntime: string }; + agentSession: { id: string }; +}> { const workspace = await resolveLiveWorkspaceForSession(db, { projectId, sessionId, userId }); - - if (!workspace || !workspace.nodeId) { + if (!workspace?.nodeId || !workspace.nodeStatus || !workspace.nodeRuntime) { throw errors.notFound('No active workspace found for this session'); } - // 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). - // 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') { + const isContainer = workspace.nodeRuntime === 'cf-container'; + const nodeIsReachable = isContainer + ? ['running', 'sleeping', 'recovery', 'error'].includes(workspace.nodeStatus) + : workspace.nodeStatus === 'running'; + if (!nodeIsReachable) { throw errors.conflict( 'The workspace node is no longer running. Start a new chat to create a fresh workspace.' ); } - // Find the running agent session on that workspace, scoped to the user - // for defence-in-depth (uses idx_agent_sessions_ws_user_status composite index). + const agentStatuses = isContainer + ? ['running', 'sleeping', 'recovery', 'error'] + : ['running', 'sleeping']; const [agentSession] = await db .select({ id: schema.agentSessions.id }) .from(schema.agentSessions) @@ -126,14 +112,20 @@ export async function resolveLiveAgentSessionForChat( and( eq(schema.agentSessions.workspaceId, workspace.id), eq(schema.agentSessions.userId, userId), - inArray(schema.agentSessions.status, ['running', 'sleeping']) + inArray(schema.agentSessions.status, agentStatuses) ) ) .limit(1); - if (!agentSession) { - throw errors.notFound('No running agent session found'); - } + if (!agentSession) throw errors.notFound('No running agent session found'); - return { workspace: { id: workspace.id, nodeId: workspace.nodeId, nodeStatus: workspace.nodeStatus }, agentSession }; + return { + workspace: { + id: workspace.id, + nodeId: workspace.nodeId, + nodeStatus: workspace.nodeStatus, + nodeRuntime: workspace.nodeRuntime, + }, + agentSession, + }; } diff --git a/apps/api/src/routes/chat.ts b/apps/api/src/routes/chat.ts index e88682a4a..839fcf6d1 100644 --- a/apps/api/src/routes/chat.ts +++ b/apps/api/src/routes/chat.ts @@ -541,7 +541,7 @@ chatRoutes.post('/:sessionId/prompt', async (c) => { c.env, userId, undefined, - workspace.nodeStatus === 'sleeping' + workspace.nodeRuntime === 'cf-container' && workspace.nodeStatus !== 'running' ? { requestTimeoutMs: getCfContainerWakeTimeoutMs(c.env) } : undefined ); diff --git a/apps/api/src/routes/projects/agent-activity-callback.ts b/apps/api/src/routes/projects/agent-activity-callback.ts index 40be0be9e..d72e49c0c 100644 --- a/apps/api/src/routes/projects/agent-activity-callback.ts +++ b/apps/api/src/routes/projects/agent-activity-callback.ts @@ -17,6 +17,11 @@ import { } from '../../services/node-callback-auth'; import * as projectDataService from '../../services/project-data'; import { markVmAgentContainerActiveWorkEndedBestEffort } from '../../services/vm-agent-container'; +import { + callbackTokenIssuedAtMs, + getInstantStaleCallbackMarginMs, + isSupersededInstantCallback, +} from '../_stale-callback-guard'; /** * Agent activity callback route — mounted BEFORE projectsRoutes in index.ts @@ -114,6 +119,53 @@ agentActivityCallbackRoute.post( throw errors.forbidden('Node identity verification failed'); } + // Staleness guard (S2): reject a DESTRUCTIVE `error` callback that provably + // originates from a superseded Instant (cf-container) generation. A mid-turn + // rollout can make a dead container fire `reportActivity("error")` (retried + // with backoff) with a still-valid callback token AFTER the DO has recovered + // a NEW container generation to running under the same nodeId. Processing it + // would regress the freshly recovered, healthy session to error/failed. + // Short-circuit BEFORE the activity mirror + destructive transition + active- + // work-ended, so nothing touches the live recovered generation. + if (body.activity === 'error') { + const guardDb = drizzle(c.env.DATABASE, { schema }); + const guardRow = await guardDb + .select({ + updatedAt: schema.agentSessions.updatedAt, + runtime: schema.nodes.runtime, + }) + .from(schema.agentSessions) + .leftJoin(schema.workspaces, eq(schema.workspaces.id, schema.agentSessions.workspaceId)) + .leftJoin(schema.nodes, eq(schema.nodes.id, schema.workspaces.nodeId)) + .where(eq(schema.agentSessions.id, sessionId)) + .get(); + const tokenIssuedAtMs = callbackTokenIssuedAtMs(token); + const marginMs = getInstantStaleCallbackMarginMs(c.env); + if ( + isSupersededInstantCallback({ + runtime: guardRow?.runtime, + rowUpdatedAt: guardRow?.updatedAt, + tokenIssuedAtMs, + marginMs, + }) + ) { + log.warn('acp_activity.rejected_stale_callback', { + projectId, + sessionId, + nodeId: body.nodeId, + runtime: guardRow?.runtime ?? null, + rowUpdatedAt: guardRow?.updatedAt ?? null, + tokenIssuedAtMs, + marginMs, + reportedActivity: body.activity, + action: 'rejected_stale_callback', + }); + // 2xx: the vm-agent activity client retries only on 5xx, so 204 stops + // retries without error-log noise. The superseded generation moves on. + return c.body(null, 204); + } + } + await projectDataService.reportAcpSessionActivity(c.env, projectId, sessionId, body.activity, { promptStartedAt: body.promptStartedAt, agentType: body.agentType, @@ -201,6 +253,7 @@ agentActivityCallbackRoute.post( { chatSessionId: workspace.chatSessionId, runtime: 'cf-container', + agentType: body.agentType ?? existing.agentType ?? undefined, } ).catch((err) => { log.warn('acp_activity.session_snapshot_failed', { diff --git a/apps/api/src/routes/tasks/callback.ts b/apps/api/src/routes/tasks/callback.ts index 4fc2bb098..08a5ebc0e 100644 --- a/apps/api/src/routes/tasks/callback.ts +++ b/apps/api/src/routes/tasks/callback.ts @@ -1,5 +1,5 @@ import { isTaskExecutionStep } from '@simple-agent-manager/shared'; -import { and, eq } from 'drizzle-orm'; +import { and, desc, eq } from 'drizzle-orm'; import { drizzle } from 'drizzle-orm/d1'; import { Hono } from 'hono'; @@ -20,6 +20,11 @@ import { isTaskStatus, } from '../../services/task-status'; import { cleanupTerminalTaskResourcesOrThrow } from '../../services/task-terminal-cleanup'; +import { + callbackTokenIssuedAtMs, + getInstantStaleCallbackMarginMs, + isSupersededInstantCallback, +} from '../_stale-callback-guard'; import { computeBlockedForTask, setTaskStatus, @@ -218,6 +223,55 @@ taskCallbackRoute.post('/:projectId/tasks/:taskId/status/callback', jsonValidato ); } + // Staleness guard (S2): reject a DESTRUCTIVE `failed` callback that provably + // originates from a superseded Instant (cf-container) generation. A dead + // container killed mid-rollout can POST `toStatus:'failed'` with a still-valid + // callback token AFTER the DO has recovered a NEW generation to running; that + // late callback would otherwise fail an already-recovered, healthy task. Use + // the workspace's most-recently-reconciled agent session `updated_at` (written + // by the DO recovery) as the completed-recovery marker. + if (body.toStatus === 'failed') { + const [guardRow] = await db + .select({ + updatedAt: schema.agentSessions.updatedAt, + runtime: schema.nodes.runtime, + }) + .from(schema.agentSessions) + .leftJoin(schema.workspaces, eq(schema.workspaces.id, schema.agentSessions.workspaceId)) + .leftJoin(schema.nodes, eq(schema.nodes.id, schema.workspaces.nodeId)) + .where(eq(schema.agentSessions.workspaceId, task.workspaceId)) + .orderBy(desc(schema.agentSessions.updatedAt)) + .limit(1); + const tokenIssuedAtMs = callbackTokenIssuedAtMs(token); + const marginMs = getInstantStaleCallbackMarginMs(c.env); + if ( + isSupersededInstantCallback({ + runtime: guardRow?.runtime, + rowUpdatedAt: guardRow?.updatedAt, + tokenIssuedAtMs, + marginMs, + }) + ) { + log.warn('task.rejected_stale_callback', { + projectId, + taskId, + workspaceId: task.workspaceId, + callerWorkspace: payload.workspace, + runtime: guardRow?.runtime ?? null, + rowUpdatedAt: guardRow?.updatedAt ?? null, + tokenIssuedAtMs, + marginMs, + fromStatus: task.status, + toStatus: body.toStatus, + action: 'rejected_stale_callback', + }); + // 2xx: postTaskCallback is fire-and-forget (no retry) and logs non-2xx as + // an error, so return the unchanged task with 200 — no regression, no noise. + const blocked = await computeBlockedForTask(db, task.id); + return c.json(toTaskResponse(task, blocked)); + } + } + const updatedTask = await setTaskStatus(db, task, body.toStatus, 'workspace_callback', payload.workspace, { reason: body.reason, outputSummary: body.outputSummary, diff --git a/apps/api/src/routes/workspaces/agent-sessions.ts b/apps/api/src/routes/workspaces/agent-sessions.ts index 69c14246a..2a6715194 100644 --- a/apps/api/src/routes/workspaces/agent-sessions.ts +++ b/apps/api/src/routes/workspaces/agent-sessions.ts @@ -1,6 +1,4 @@ -import type { - AgentSession, -} from '@simple-agent-manager/shared'; +import type { AgentSession } from '@simple-agent-manager/shared'; import { and, desc, eq } from 'drizzle-orm'; import { drizzle } from 'drizzle-orm/d1'; import { Hono } from 'hono'; @@ -11,8 +9,8 @@ import { log } from '../../lib/logger'; import { toAgentSessionResponse } from '../../lib/mappers'; import { parsePositiveInt } from '../../lib/route-helpers'; import { ulid } from '../../lib/ulid'; -import { getUserId, requireApproved,requireAuth } from '../../middleware/auth'; -import { errors } from '../../middleware/error'; +import { getUserId, requireApproved, requireAuth } from '../../middleware/auth'; +import { AppError, errors } from '../../middleware/error'; import { CreateAgentSessionSchema, jsonValidator, UpdateAgentSessionSchema } from '../../schemas'; import { getRuntimeLimits } from '../../services/limits'; import { generateMcpToken, revokeMcpToken, storeMcpToken } from '../../services/mcp-token'; @@ -22,8 +20,9 @@ import { stopAgentSessionOnNode, suspendAgentSessionOnNode, } from '../../services/node-agent'; +import { resumeVmAgentContainer } from '../../services/vm-agent-container'; import { requireRepositoryOwnerAccess } from '../projects/_helpers'; -import { assertNodeOperational,getOwnedNode, getOwnedWorkspace } from './_helpers'; +import { assertNodeOperational, getOwnedNode, getOwnedWorkspace } from './_helpers'; const agentSessionRoutes = new Hono<{ Bindings: Env }>(); @@ -73,362 +72,469 @@ agentSessionRoutes.get('/:id/agent-sessions', requireAuth(), requireApproved(), return c.json(sessions.map(toAgentSessionResponse)); }); -agentSessionRoutes.post('/:id/agent-sessions', requireAuth(), requireApproved(), jsonValidator(CreateAgentSessionSchema), async (c) => { - const userId = getUserId(c); - const workspaceId = c.req.param('id'); - const db = drizzle(c.env.DATABASE, { schema }); - const body = c.req.valid('json'); - const limits = getRuntimeLimits(c.env); - - const workspace = await getOwnedWorkspace(db, workspaceId, userId); - if (!workspace.nodeId) { - throw errors.badRequest('Workspace is not attached to a node'); - } - - const node = await getOwnedNode(db, workspace.nodeId, userId); - assertNodeOperational(node, 'create agent session'); - await requireWorkspaceAgentGitHubAccess(c.env, db, workspace, userId); - - const existingRunning = await db - .select({ id: schema.agentSessions.id }) - .from(schema.agentSessions) - .where( - and( - eq(schema.agentSessions.workspaceId, workspace.id), - eq(schema.agentSessions.userId, userId), - eq(schema.agentSessions.status, 'running') - ) - ); +agentSessionRoutes.post( + '/:id/agent-sessions', + requireAuth(), + requireApproved(), + jsonValidator(CreateAgentSessionSchema), + async (c) => { + const userId = getUserId(c); + const workspaceId = c.req.param('id'); + const db = drizzle(c.env.DATABASE, { schema }); + const body = c.req.valid('json'); + const limits = getRuntimeLimits(c.env); + + const workspace = await getOwnedWorkspace(db, workspaceId, userId); + if (!workspace.nodeId) { + throw errors.badRequest('Workspace is not attached to a node'); + } - if (existingRunning.length >= limits.maxAgentSessionsPerWorkspace) { - throw errors.badRequest( - `Maximum ${limits.maxAgentSessionsPerWorkspace} agent sessions per workspace` - ); - } + const node = await getOwnedNode(db, workspace.nodeId, userId); + assertNodeOperational(node, 'create agent session'); + await requireWorkspaceAgentGitHubAccess(c.env, db, workspace, userId); + + const existingRunning = await db + .select({ id: schema.agentSessions.id }) + .from(schema.agentSessions) + .where( + and( + eq(schema.agentSessions.workspaceId, workspace.id), + eq(schema.agentSessions.userId, userId), + eq(schema.agentSessions.status, 'running') + ) + ); - const sessionId = ulid(); - const now = new Date().toISOString(); - - await db.insert(schema.agentSessions).values({ - id: sessionId, - workspaceId: workspace.id, - userId, - status: 'running', - label: body.label?.trim() || null, - agentType: body.agentType?.trim() || null, - worktreePath: body.worktreePath?.trim() || null, - createdAt: now, - updatedAt: now, - }); - - let mcpToken: string | null = null; - try { - if (workspace.projectId) { - mcpToken = generateMcpToken(); - await storeMcpToken( - c.env.KV, - mcpToken, - { - // Empty taskId for direct project-chat sessions — only task-runner dispatched - // sessions have a real task row. Setting sessionId as taskId was wrong because - // MCP tools query tasks by this ID and would get "Task not found". Empty string - // is falsy so tools guarding on !tokenData.taskId correctly reject early. - taskId: '', - contextType: workspace.chatSessionId ? 'conversation' : 'direct-workspace', - taskMode: workspace.chatSessionId ? 'conversation' : undefined, - projectId: workspace.projectId, - userId, - workspaceId: workspace.id, - chatSessionId: workspace.chatSessionId ?? undefined, - agentSessionId: sessionId, - createdAt: new Date().toISOString(), - }, - c.env, + if (existingRunning.length >= limits.maxAgentSessionsPerWorkspace) { + throw errors.badRequest( + `Maximum ${limits.maxAgentSessionsPerWorkspace} agent sessions per workspace` ); } - await createAgentSessionOnNode( - workspace.nodeId, - workspace.id, - sessionId, - body.label?.trim() || null, - c.env, + const sessionId = ulid(); + const now = new Date().toISOString(); + + await db.insert(schema.agentSessions).values({ + id: sessionId, + workspaceId: workspace.id, userId, - workspace.chatSessionId, - workspace.projectId, - mcpToken - ? { - url: `https://api.${c.env.BASE_DOMAIN}/mcp`, - token: mcpToken, - } - : undefined, - ); - } catch (err) { - if (mcpToken) { - await revokeMcpToken(c.env.KV, mcpToken).catch((revokeErr) => { - log.warn('agent_session.mcp_token_revoke_failed', { - sessionId, - workspaceId: workspace.id, - error: revokeErr instanceof Error ? revokeErr.message : String(revokeErr), - }); - }); - } + status: 'running', + label: body.label?.trim() || null, + agentType: body.agentType?.trim() || null, + worktreePath: body.worktreePath?.trim() || null, + createdAt: now, + updatedAt: now, + }); - await db - .update(schema.agentSessions) - .set({ - status: 'error', - errorMessage: err instanceof Error ? err.message : 'Failed to create agent session', - updatedAt: new Date().toISOString(), - }) - .where(eq(schema.agentSessions.id, sessionId)); + let mcpToken: string | null = null; + try { + if (workspace.projectId) { + mcpToken = generateMcpToken(); + await storeMcpToken( + c.env.KV, + mcpToken, + { + // Empty taskId for direct project-chat sessions — only task-runner dispatched + // sessions have a real task row. Setting sessionId as taskId was wrong because + // MCP tools query tasks by this ID and would get "Task not found". Empty string + // is falsy so tools guarding on !tokenData.taskId correctly reject early. + taskId: '', + contextType: workspace.chatSessionId ? 'conversation' : 'direct-workspace', + taskMode: workspace.chatSessionId ? 'conversation' : undefined, + projectId: workspace.projectId, + userId, + workspaceId: workspace.id, + chatSessionId: workspace.chatSessionId ?? undefined, + agentSessionId: sessionId, + createdAt: new Date().toISOString(), + }, + c.env + ); + } - throw errors.internal('Failed to create agent session on node'); - } + await createAgentSessionOnNode( + workspace.nodeId, + workspace.id, + sessionId, + body.label?.trim() || null, + c.env, + userId, + workspace.chatSessionId, + workspace.projectId, + mcpToken + ? { + url: `https://api.${c.env.BASE_DOMAIN}/mcp`, + token: mcpToken, + } + : undefined + ); + } catch (err) { + if (mcpToken) { + await revokeMcpToken(c.env.KV, mcpToken).catch((revokeErr) => { + log.warn('agent_session.mcp_token_revoke_failed', { + sessionId, + workspaceId: workspace.id, + error: revokeErr instanceof Error ? revokeErr.message : String(revokeErr), + }); + }); + } - const rows = await db - .select() - .from(schema.agentSessions) - .where(eq(schema.agentSessions.id, sessionId)) - .limit(1); + await db + .update(schema.agentSessions) + .set({ + status: 'error', + errorMessage: err instanceof Error ? err.message : 'Failed to create agent session', + updatedAt: new Date().toISOString(), + }) + .where(eq(schema.agentSessions.id, sessionId)); - return c.json(toAgentSessionResponse(rows[0]!), 201); -}); + throw errors.internal('Failed to create agent session on node'); + } -agentSessionRoutes.patch('/:id/agent-sessions/:sessionId', requireAuth(), requireApproved(), jsonValidator(UpdateAgentSessionSchema), async (c) => { - const userId = getUserId(c); - const workspaceId = c.req.param('id'); - const sessionId = c.req.param('sessionId'); - const db = drizzle(c.env.DATABASE, { schema }); + const rows = await db + .select() + .from(schema.agentSessions) + .where(eq(schema.agentSessions.id, sessionId)) + .limit(1); - const workspace = await getOwnedWorkspace(db, workspaceId, userId); - if (!workspace) { - throw errors.notFound('Workspace'); + return c.json(toAgentSessionResponse(rows[0]!), 201); } +); + +agentSessionRoutes.patch( + '/:id/agent-sessions/:sessionId', + requireAuth(), + requireApproved(), + jsonValidator(UpdateAgentSessionSchema), + async (c) => { + const userId = getUserId(c); + const workspaceId = c.req.param('id'); + const sessionId = c.req.param('sessionId'); + const db = drizzle(c.env.DATABASE, { schema }); + + const workspace = await getOwnedWorkspace(db, workspaceId, userId); + if (!workspace) { + throw errors.notFound('Workspace'); + } - const body = c.req.valid('json'); - const maxLabelLength = parsePositiveInt(c.env.MAX_AGENT_SESSION_LABEL_LENGTH, 50); - const label = body.label?.trim()?.slice(0, maxLabelLength); - if (!label) { - throw errors.badRequest('Label is required and must be non-empty'); - } + const body = c.req.valid('json'); + const maxLabelLength = parsePositiveInt(c.env.MAX_AGENT_SESSION_LABEL_LENGTH, 50); + const label = body.label?.trim()?.slice(0, maxLabelLength); + if (!label) { + throw errors.badRequest('Label is required and must be non-empty'); + } - const rows = await db - .select() - .from(schema.agentSessions) - .where( - and( - eq(schema.agentSessions.id, sessionId), - eq(schema.agentSessions.workspaceId, workspace.id), - eq(schema.agentSessions.userId, userId) + const rows = await db + .select() + .from(schema.agentSessions) + .where( + and( + eq(schema.agentSessions.id, sessionId), + eq(schema.agentSessions.workspaceId, workspace.id), + eq(schema.agentSessions.userId, userId) + ) ) - ) - .limit(1); - - const session = rows[0]; - if (!session) { - throw errors.notFound('Agent session'); - } + .limit(1); - if (session.status !== 'running') { - throw errors.badRequest('Cannot rename a session that is not running'); - } - - await db - .update(schema.agentSessions) - .set({ - label, - updatedAt: new Date().toISOString(), - }) - .where(eq(schema.agentSessions.id, session.id)); + const session = rows[0]; + if (!session) { + throw errors.notFound('Agent session'); + } - return c.json(toAgentSessionResponse({ ...session, label, updatedAt: new Date().toISOString() })); -}); + if (session.status !== 'running') { + throw errors.badRequest('Cannot rename a session that is not running'); + } -agentSessionRoutes.post('/:id/agent-sessions/:sessionId/stop', requireAuth(), requireApproved(), async (c) => { - const userId = getUserId(c); - const workspaceId = c.req.param('id'); - const sessionId = c.req.param('sessionId'); - const db = drizzle(c.env.DATABASE, { schema }); + await db + .update(schema.agentSessions) + .set({ + label, + updatedAt: new Date().toISOString(), + }) + .where(eq(schema.agentSessions.id, session.id)); - const workspace = await getOwnedWorkspace(db, workspaceId, userId); - if (!workspace.nodeId) { - throw errors.badRequest('Workspace is not attached to a node'); + return c.json( + toAgentSessionResponse({ ...session, label, updatedAt: new Date().toISOString() }) + ); } +); + +agentSessionRoutes.post( + '/:id/agent-sessions/:sessionId/stop', + requireAuth(), + requireApproved(), + async (c) => { + const userId = getUserId(c); + const workspaceId = c.req.param('id'); + const sessionId = c.req.param('sessionId'); + const db = drizzle(c.env.DATABASE, { schema }); + + const workspace = await getOwnedWorkspace(db, workspaceId, userId); + if (!workspace.nodeId) { + throw errors.badRequest('Workspace is not attached to a node'); + } - const rows = await db - .select() - .from(schema.agentSessions) - .where( - and( - eq(schema.agentSessions.id, sessionId), - eq(schema.agentSessions.workspaceId, workspace.id), - eq(schema.agentSessions.userId, userId) + const rows = await db + .select() + .from(schema.agentSessions) + .where( + and( + eq(schema.agentSessions.id, sessionId), + eq(schema.agentSessions.workspaceId, workspace.id), + eq(schema.agentSessions.userId, userId) + ) ) - ) - .limit(1); + .limit(1); - const session = rows[0]; - if (!session) { - throw errors.notFound('Agent session'); - } + const session = rows[0]; + if (!session) { + throw errors.notFound('Agent session'); + } - if (session.status !== 'running') { - // Still attempt VM stop for orphaned sessions whose process may be alive - if (workspace.nodeId) { - try { - await stopAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); - } catch (e) { - log.error('agent_session.orphaned_stop_failed', { sessionId: session.id, workspaceId: workspace.id, nodeId: workspace.nodeId, error: String(e) }); + if (session.status !== 'running') { + // Still attempt VM stop for orphaned sessions whose process may be alive + if (workspace.nodeId) { + try { + await stopAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); + } catch (e) { + log.error('agent_session.orphaned_stop_failed', { + sessionId: session.id, + workspaceId: workspace.id, + nodeId: workspace.nodeId, + error: String(e), + }); + } } + return c.json({ status: session.status }); } - return c.json({ status: session.status }); - } - - try { - await stopAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); - } catch (e) { - log.error('agent_session.stop_on_node_failed', { sessionId: session.id, workspaceId: workspace.id, nodeId: workspace.nodeId, error: String(e) }); - } - await db - .update(schema.agentSessions) - .set({ - status: 'stopped', - stoppedAt: new Date().toISOString(), - errorMessage: null, - updatedAt: new Date().toISOString(), - }) - .where(eq(schema.agentSessions.id, session.id)); - - return c.json({ status: 'stopped' }); -}); + try { + await stopAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); + } catch (e) { + log.error('agent_session.stop_on_node_failed', { + sessionId: session.id, + workspaceId: workspace.id, + nodeId: workspace.nodeId, + error: String(e), + }); + } -agentSessionRoutes.post('/:id/agent-sessions/:sessionId/suspend', requireAuth(), requireApproved(), async (c) => { - const userId = getUserId(c); - const workspaceId = c.req.param('id'); - const sessionId = c.req.param('sessionId'); - const db = drizzle(c.env.DATABASE, { schema }); + await db + .update(schema.agentSessions) + .set({ + status: 'stopped', + stoppedAt: new Date().toISOString(), + errorMessage: null, + updatedAt: new Date().toISOString(), + }) + .where(eq(schema.agentSessions.id, session.id)); - const workspace = await getOwnedWorkspace(db, workspaceId, userId); - if (!workspace.nodeId) { - throw errors.badRequest('Workspace is not attached to a node'); + return c.json({ status: 'stopped' }); } +); + +agentSessionRoutes.post( + '/:id/agent-sessions/:sessionId/suspend', + requireAuth(), + requireApproved(), + async (c) => { + const userId = getUserId(c); + const workspaceId = c.req.param('id'); + const sessionId = c.req.param('sessionId'); + const db = drizzle(c.env.DATABASE, { schema }); + + const workspace = await getOwnedWorkspace(db, workspaceId, userId); + if (!workspace.nodeId) { + throw errors.badRequest('Workspace is not attached to a node'); + } - const rows = await db - .select() - .from(schema.agentSessions) - .where( - and( - eq(schema.agentSessions.id, sessionId), - eq(schema.agentSessions.workspaceId, workspace.id), - eq(schema.agentSessions.userId, userId) + const rows = await db + .select() + .from(schema.agentSessions) + .where( + and( + eq(schema.agentSessions.id, sessionId), + eq(schema.agentSessions.workspaceId, workspace.id), + eq(schema.agentSessions.userId, userId) + ) ) - ) - .limit(1); - - const session = rows[0]; - if (!session) { - throw errors.notFound('Agent session'); - } - - if (session.status !== 'running' && session.status !== 'error') { - throw errors.badRequest(`Session cannot be suspended from status: ${session.status}`); - } + .limit(1); - try { - await suspendAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); - } catch (e) { - log.warn('agent_session.suspend_on_node_failed', { sessionId: session.id, workspaceId: workspace.id, nodeId: workspace.nodeId, error: String(e) }); - } + const session = rows[0]; + if (!session) { + throw errors.notFound('Agent session'); + } - const now = new Date().toISOString(); - await db - .update(schema.agentSessions) - .set({ - status: 'suspended', - suspendedAt: now, - errorMessage: null, - updatedAt: now, - }) - .where(eq(schema.agentSessions.id, session.id)); - - return c.json( - toAgentSessionResponse({ - ...session, - status: 'suspended', - suspendedAt: now, - errorMessage: null, - updatedAt: now, - }) - ); -}); + if (session.status !== 'running' && session.status !== 'error') { + throw errors.badRequest(`Session cannot be suspended from status: ${session.status}`); + } -agentSessionRoutes.post('/:id/agent-sessions/:sessionId/resume', requireAuth(), requireApproved(), async (c) => { - const userId = getUserId(c); - const workspaceId = c.req.param('id'); - const sessionId = c.req.param('sessionId'); - const db = drizzle(c.env.DATABASE, { schema }); + try { + await suspendAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); + } catch (e) { + log.warn('agent_session.suspend_on_node_failed', { + sessionId: session.id, + workspaceId: workspace.id, + nodeId: workspace.nodeId, + error: String(e), + }); + } - const workspace = await getOwnedWorkspace(db, workspaceId, userId); - if (!workspace.nodeId) { - throw errors.badRequest('Workspace is not attached to a node'); + const now = new Date().toISOString(); + await db + .update(schema.agentSessions) + .set({ + status: 'suspended', + suspendedAt: now, + errorMessage: null, + updatedAt: now, + }) + .where(eq(schema.agentSessions.id, session.id)); + + return c.json( + toAgentSessionResponse({ + ...session, + status: 'suspended', + suspendedAt: now, + errorMessage: null, + updatedAt: now, + }) + ); } +); + +agentSessionRoutes.post( + '/:id/agent-sessions/:sessionId/resume', + requireAuth(), + requireApproved(), + async (c) => { + const userId = getUserId(c); + const workspaceId = c.req.param('id'); + const sessionId = c.req.param('sessionId'); + const db = drizzle(c.env.DATABASE, { schema }); + + const workspace = await getOwnedWorkspace(db, workspaceId, userId); + if (!workspace.nodeId) { + throw errors.badRequest('Workspace is not attached to a node'); + } - const rows = await db - .select() - .from(schema.agentSessions) - .where( - and( - eq(schema.agentSessions.id, sessionId), - eq(schema.agentSessions.workspaceId, workspace.id), - eq(schema.agentSessions.userId, userId) + const rows = await db + .select() + .from(schema.agentSessions) + .where( + and( + eq(schema.agentSessions.id, sessionId), + eq(schema.agentSessions.workspaceId, workspace.id), + eq(schema.agentSessions.userId, userId) + ) ) - ) - .limit(1); + .limit(1); - const session = rows[0]; - if (!session) { - throw errors.notFound('Agent session'); - } + const session = rows[0]; + if (!session) { + throw errors.notFound('Agent session'); + } - // Already running -- idempotent - if (session.status === 'running') { - return c.json(toAgentSessionResponse(session)); - } + const node = await getOwnedNode(db, workspace.nodeId, userId); + // Intentional WIDE gate: recovery must fire whenever EITHER the node or the + // session looks not-running in D1. These D1 prechecks can be stale — the + // container may have already died and been marked 'recovery'/'error' — so a + // narrower gate would let a dead-D1 precheck block the very recovery that + // heals it. The Durable Object owns the live container generation and + // reconciles D1 itself; over-triggering here is safe because the DO + // fast-paths an already-running container. Do NOT narrow this gate without + // re-checking the recovery contract. + if ( + node.runtime === 'cf-container' && + (node.status !== 'running' || session.status !== 'running') + ) { + const recovery = await resumeVmAgentContainer(c.env, node.id, session.id); + if (!recovery.ok) { + if (!recovery.code || !recovery.message) { + throw errors.internal('Instant session recovery failed.'); + } + const statusCode = + recovery.code === 'RUNTIME_STOPPED' + ? 410 + : recovery.code === 'RUNTIME_RECOVERING' + ? 503 + : 409; + throw new AppError(statusCode, recovery.code, recovery.message); + } - // Resume is allowed from suspended, stopped, or error states. - // For suspended sessions, also tell the VM agent to resume. - if (session.status === 'suspended') { - try { - await resumeAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); - } catch (e) { - log.warn('agent_session.resume_on_node_failed', { sessionId: session.id, workspaceId: workspace.id, nodeId: workspace.nodeId, error: String(e) }); + // Recovery succeeded. The Durable Object (persistRuntimeRecovered) has + // ALREADY reconciled the agent_sessions row in D1 for THIS request — + // including error_message: when the interrupted prompt needs manual retry + // it persists RUNTIME_REQUEST_INTERRUPTED_MESSAGE so reloads / other + // devices see "your message needs manual retry". The local `session` + // snapshot read before recovery is now stale. Re-fetch and return the + // DO-reconciled row; do NOT fall through to the status rewrite below + // (S4+CF3), which would clobber the DO-written status/error_message back + // to running / null. + const [recovered] = await db + .select() + .from(schema.agentSessions) + .where(eq(schema.agentSessions.id, session.id)) + .limit(1); + return c.json( + toAgentSessionResponse( + recovered ?? { + ...session, + status: 'running', + stoppedAt: null, + suspendedAt: null, + errorMessage: null, + updatedAt: new Date().toISOString(), + } + ) + ); } - } - const now = new Date().toISOString(); - await db - .update(schema.agentSessions) - .set({ - status: 'running', - stoppedAt: null, - suspendedAt: null, - errorMessage: null, - updatedAt: now, - }) - .where(eq(schema.agentSessions.id, session.id)); + // Non cf-container, or an already-running cf-container session: fall through + // to the legacy/VM lifecycle rewrite below. - return c.json( - toAgentSessionResponse({ - ...session, - status: 'running', - stoppedAt: null, - suspendedAt: null, - errorMessage: null, - updatedAt: now, - }) - ); -}); + // Already running -- idempotent + if (session.status === 'running') { + return c.json(toAgentSessionResponse(session)); + } + + // Resume is allowed from suspended, stopped, or error states. + // For suspended sessions, also tell the VM agent to resume. + if (session.status === 'suspended') { + try { + await resumeAgentSessionOnNode(workspace.nodeId, workspace.id, session.id, c.env, userId); + } catch (e) { + log.warn('agent_session.resume_on_node_failed', { + sessionId: session.id, + workspaceId: workspace.id, + nodeId: workspace.nodeId, + error: String(e), + }); + } + } + + const now = new Date().toISOString(); + await db + .update(schema.agentSessions) + .set({ + status: 'running', + stoppedAt: null, + suspendedAt: null, + errorMessage: null, + updatedAt: now, + }) + .where(eq(schema.agentSessions.id, session.id)); + + return c.json( + toAgentSessionResponse({ + ...session, + status: 'running', + stoppedAt: null, + suspendedAt: null, + errorMessage: null, + updatedAt: now, + }) + ); + } +); export { agentSessionRoutes }; diff --git a/apps/api/src/routes/workspaces/session-snapshots.ts b/apps/api/src/routes/workspaces/session-snapshots.ts index ac4b72eae..f0ca3c3b5 100644 --- a/apps/api/src/routes/workspaces/session-snapshots.ts +++ b/apps/api/src/routes/workspaces/session-snapshots.ts @@ -54,6 +54,11 @@ function stringField(body: Record, key: string, required = true return value.trim(); } +function optionalStringField(body: Record, key: string): string | null { + if (!(key in body) || body[key] === null || body[key] === undefined) return null; + return requiredStringField(body, key); +} + function requiredStringField(body: Record, key: string): string { const value = stringField(body, key, true); if (value === null) throw errors.badRequest(`${key} is required`); @@ -169,6 +174,16 @@ sessionSnapshotRoutes.post('/:id/session-snapshot/complete', async (c) => { ) { throw errors.badRequest('Snapshot manifest identity does not match request'); } + const agentSessionId = optionalStringField(body, 'agentSessionId'); + const manifestAgentSessionId = optionalStringField(manifestRecord, 'agentSessionId'); + if (manifestAgentSessionId !== agentSessionId) { + throw errors.badRequest('Snapshot agent session does not match request'); + } + const acpSessionId = optionalStringField(manifestRecord, 'acpSessionId'); + const agentType = optionalStringField(manifestRecord, 'agentType'); + if ((acpSessionId === null) !== (agentType === null)) { + throw errors.badRequest('Snapshot agent context identity is incomplete'); + } const manifestArtifacts = expectJsonRecord( manifestRecord.artifacts ?? {}, 'snapshot manifest artifacts' @@ -201,7 +216,7 @@ sessionSnapshotRoutes.post('/:id/session-snapshot/complete', async (c) => { await completeSessionSnapshot(db, c.env, { workspaceId, chatSessionId, - agentSessionId: stringField(body, 'agentSessionId', false), + agentSessionId, runtime: stringField(body, 'runtime', false) || 'runtime-neutral', baseCommit: stringField(body, 'baseCommit', false), status, diff --git a/apps/api/src/scheduled/stuck-tasks.ts b/apps/api/src/scheduled/stuck-tasks.ts index cb805fe9e..5ed918634 100644 --- a/apps/api/src/scheduled/stuck-tasks.ts +++ b/apps/api/src/scheduled/stuck-tasks.ts @@ -39,6 +39,7 @@ import { drizzle } from 'drizzle-orm/d1'; import * as schema from '../db/schema'; import type { TaskRunner } from '../durable-objects/task-runner'; +import { classifyVmAgentContainerLiveness } from '../durable-objects/vm-agent-container-lifecycle'; import type { Env } from '../env'; import { log } from '../lib/logger'; import { ulid } from '../lib/ulid'; @@ -46,6 +47,7 @@ import { persistError } from '../services/observability'; import * as projectDataService from '../services/project-data'; import { cleanupTaskRun } from '../services/task-runner'; import { syncTriggerExecutionStatus } from '../services/trigger-execution-sync'; +import { inspectVmAgentContainerLifecycle } from '../services/vm-agent-container'; import { type CompactionLoopRecovery, detectTaskCompactionLoop, @@ -196,7 +198,18 @@ export interface TaskReconciliationDiagnostics { }; } +// 'running' is always a live-candidate. 'recovery' is a live-candidate ONLY for +// VM runtimes: the VM /ready callback path transiently marks a workspace +// 'recovery' while the agent reconnects, and such workspaces must still fall +// through to the node-heartbeat staleness + ACP-session liveness checks below +// (pre-PR behavior). For cf-container (Instant) runtimes, 'recovery'/'sleeping' +// are instead handled by the resumable short-circuit, which defers to the +// authoritative DO lifecycle probe. const LIVE_WORKSPACE_STATUSES = new Set(['running', 'recovery']); +// Resumable statuses short-circuit to an inconclusive result ONLY for +// cf-container runtimes. Gating on node_runtime keeps VM 'recovery'/'sleeping' +// workspaces on a dead node conclusively reconcilable instead of stranding them. +const RESUMABLE_WORKSPACE_STATUSES = new Set(['sleeping', 'recovery']); const ACTIVE_ACP_STATUSES = new Set(['assigned', 'running']); function stuckTaskScanCursorKey(env: Env): string { @@ -424,7 +437,8 @@ export async function getTaskRuntimeLiveness( const row = await env.DATABASE.prepare( `SELECT w.status AS workspace_status, w.chat_session_id, w.node_id, - n.status AS node_status, n.health_status, n.last_heartbeat_at + n.status AS node_status, n.health_status, n.last_heartbeat_at, + n.runtime AS node_runtime FROM workspaces w LEFT JOIN nodes n ON n.id = w.node_id WHERE w.id = ? @@ -438,9 +452,20 @@ export async function getTaskRuntimeLiveness( node_status: string | null; health_status: string | null; last_heartbeat_at: string | null; + node_runtime: string | null; }>(); if (!row) return dead('workspace_missing', null, null); + if (row.node_runtime === 'cf-container' && RESUMABLE_WORKSPACE_STATUSES.has(row.workspace_status)) { + return { + live: false, + conclusive: false, + reason: `workspace_${row.workspace_status}_resumable`, + workspaceStatus: row.workspace_status, + nodeId: row.node_id, + activeAcpSessionId: null, + }; + } if (!LIVE_WORKSPACE_STATUSES.has(row.workspace_status)) { return dead(`workspace_${row.workspace_status}`, row.workspace_status, row.node_id); } @@ -455,6 +480,56 @@ export async function getTaskRuntimeLiveness( }; } + const probeTimeoutMs = parseMs( + env.TASK_LIVENESS_PROBE_TIMEOUT_MS, + DEFAULT_TASK_LIVENESS_PROBE_TIMEOUT_MS + ); + if (row.node_runtime === 'cf-container') { + const TIMEOUT = Symbol('container_lifecycle_probe_timeout'); + let timer: ReturnType | undefined; + try { + const probe = await Promise.race([ + inspectVmAgentContainerLifecycle(env, row.node_id), + new Promise((resolve) => { + timer = setTimeout(() => resolve(TIMEOUT), probeTimeoutMs); + }), + ]); + if (probe === TIMEOUT) { + return { + live: false, + conclusive: false, + reason: 'cf_container_lifecycle_timeout', + workspaceStatus: row.workspace_status, + nodeId: row.node_id, + activeAcpSessionId: null, + }; + } + const classification = classifyVmAgentContainerLiveness(probe); + return { + ...classification, + workspaceStatus: row.workspace_status, + nodeId: row.node_id, + activeAcpSessionId: null, + }; + } catch (err) { + log.warn('stuck_task.container_lifecycle_probe_failed', { + workspaceId: task.workspace_id, + nodeId: row.node_id, + error: err instanceof Error ? err.message : String(err), + }); + return { + live: false, + conclusive: false, + reason: 'cf_container_lifecycle_unknown', + workspaceStatus: row.workspace_status, + nodeId: row.node_id, + activeAcpSessionId: null, + }; + } finally { + if (timer) clearTimeout(timer); + } + } + const staleSeconds = parseInt(env.NODE_HEARTBEAT_STALE_SECONDS || '', 10) || DEFAULT_NODE_HEARTBEAT_STALE_SECONDS; const staleMs = staleSeconds * 1000; @@ -477,10 +552,6 @@ export async function getTaskRuntimeLiveness( ); // Bound the ProjectData DO probe so a slow/unresponsive DO cannot stall the // control-loop sweep (rule 47). A timeout is inconclusive, never fatal. - const probeTimeoutMs = parseMs( - env.TASK_LIVENESS_PROBE_TIMEOUT_MS, - DEFAULT_TASK_LIVENESS_PROBE_TIMEOUT_MS - ); const TIMEOUT = Symbol('liveness_probe_timeout'); let timer: ReturnType | undefined; const probe = await Promise.race([ diff --git a/apps/api/src/services/node-agent-readiness.ts b/apps/api/src/services/node-agent-readiness.ts new file mode 100644 index 000000000..2a6d888af --- /dev/null +++ b/apps/api/src/services/node-agent-readiness.ts @@ -0,0 +1,107 @@ +import type { Env } from '../env'; + +const DEFAULT_NODE_AGENT_READY_TIMEOUT_MS = 900_000; +const DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS = 5000; + +type NodeAgentFetch = ( + nodeId: string, + env: Env, + url: string, + options: RequestInit, + requestTimeoutMs: number +) => Promise; + +export function getNodeBackendBaseUrl(nodeId: string, env: Env): string { + const protocol = env.VM_AGENT_PROTOCOL || 'https'; + const port = env.VM_AGENT_PORT || '8443'; + // Two-level subdomain ({nodeId}.vm.{domain}) bypasses Cloudflare same-zone routing. + // The wildcard Worker route *.{domain}/* matches exactly one subdomain level, + // so {nodeId}.vm.{domain} is NOT intercepted — requests reach the VM directly. + return `${protocol}://${nodeId.toLowerCase()}.vm.${env.BASE_DOMAIN}:${port}`; +} + +export function getNodeAgentReadyTimeoutMs(env: { + NODE_AGENT_READY_TIMEOUT_MS?: string; +}): number { + const parsed = env.NODE_AGENT_READY_TIMEOUT_MS + ? Number.parseInt(env.NODE_AGENT_READY_TIMEOUT_MS, 10) + : DEFAULT_NODE_AGENT_READY_TIMEOUT_MS; + if (!Number.isFinite(parsed) || parsed <= 0) { + return DEFAULT_NODE_AGENT_READY_TIMEOUT_MS; + } + return parsed; +} + +export function getNodeAgentReadyPollIntervalMs(env: { + NODE_AGENT_READY_POLL_INTERVAL_MS?: string; +}): number { + const parsed = env.NODE_AGENT_READY_POLL_INTERVAL_MS + ? Number.parseInt(env.NODE_AGENT_READY_POLL_INTERVAL_MS, 10) + : DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS; + if (!Number.isFinite(parsed) || parsed <= 0) { + return DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS; + } + return parsed; +} + +export async function waitForNodeAgentReadyWith( + fetchNodeAgent: NodeAgentFetch, + nodeId: string, + env: Env +): Promise { + const timeoutMs = getNodeAgentReadyTimeoutMs(env); + const pollIntervalMs = getNodeAgentReadyPollIntervalMs(env); + const healthUrl = `${getNodeBackendBaseUrl(nodeId, env)}/health`; + const deadline = Date.now() + timeoutMs; + + let lastError = ''; + + while (Date.now() < deadline) { + const remainingMs = deadline - Date.now(); + const requestTimeoutMs = Math.max(1, Math.min(pollIntervalMs, remainingMs)); + const controller = new AbortController(); + let timeoutHandle: ReturnType | null = null; + + try { + const requestTimeoutError = `request timeout after ${requestTimeoutMs}ms`; + const response = await Promise.race([ + fetchNodeAgent( + nodeId, + env, + healthUrl, + { method: 'GET', signal: controller.signal }, + requestTimeoutMs + ), + new Promise((_resolve, reject) => { + timeoutHandle = setTimeout(() => { + controller.abort(); + reject(new Error(requestTimeoutError)); + }, requestTimeoutMs); + }), + ]); + if (response.ok) return; + + const responseBody = await response.text().catch(() => ''); + lastError = `HTTP ${response.status}${responseBody ? ` ${responseBody}` : ''}`.trim(); + } catch (err) { + if (err instanceof Error && err.message.startsWith('request timeout after ')) { + lastError = err.message; + } else if (err instanceof Error && err.message.startsWith('Request timed out after ')) { + lastError = err.message.replace('Request timed out after ', 'request timeout after '); + } else if (err instanceof Error && err.name === 'AbortError') { + lastError = `request timeout after ${requestTimeoutMs}ms`; + } else { + lastError = err instanceof Error ? err.message : String(err); + } + } finally { + if (timeoutHandle) clearTimeout(timeoutHandle); + } + + const nextRemainingMs = deadline - Date.now(); + if (nextRemainingMs <= 0) break; + await new Promise((resolve) => setTimeout(resolve, Math.min(pollIntervalMs, nextRemainingMs))); + } + + const details = lastError ? ` Last error: ${lastError}` : ''; + throw new Error(`Node Agent not reachable at ${healthUrl} within ${timeoutMs}ms.${details}`); +} diff --git a/apps/api/src/services/node-agent-session-snapshots.ts b/apps/api/src/services/node-agent-session-snapshots.ts index c94944a7e..287a1af07 100644 --- a/apps/api/src/services/node-agent-session-snapshots.ts +++ b/apps/api/src/services/node-agent-session-snapshots.ts @@ -4,6 +4,7 @@ import { getNodeAgentRequestTimeoutMs, nodeAgentRequest } from './node-agent'; interface SessionSnapshotRequest { chatSessionId: string; runtime: string; + agentType?: string; } function requestSessionSnapshot( @@ -15,13 +16,18 @@ function requestSessionSnapshot( 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), - }); + return nodeAgentRequest( + nodeId, + env, + `/workspaces/${workspaceId}/agent-sessions/${sessionId}/${action}`, + { + method: 'POST', + userId, + workspaceId, + requestTimeoutMs: getNodeAgentRequestTimeoutMs(env), + body: JSON.stringify(input), + } + ); } export function hibernateAgentSessionOnNode( diff --git a/apps/api/src/services/node-agent.ts b/apps/api/src/services/node-agent.ts index 452046ae4..8e730cfcc 100644 --- a/apps/api/src/services/node-agent.ts +++ b/apps/api/src/services/node-agent.ts @@ -1,17 +1,30 @@ +// FILE SIZE EXCEPTION: Cross-boundary node-agent request layer marginally over the limit; interactive/background timeout tiers and cf-container interruption classification are one cohesive concern. Split candidate if it grows further. See .claude/rules/18-file-size-limits.md import { eq } from 'drizzle-orm'; import { drizzle } from 'drizzle-orm/d1'; import * as schema from '../db/schema'; +import { + getRuntimeRecoveryMessage, + type RuntimeRecoveryCode, +} from '../durable-objects/vm-agent-container-recovery'; import type { Env } from '../env'; import { expectJsonRecord } from '../lib/runtime-validation'; +import { AppError } from '../middleware/error'; import { fetchWithTimeout, getTimeoutMs } from './fetch-timeout'; import { signNodeManagementToken, signTerminalToken } from './jwt'; +import { + getNodeAgentReadyPollIntervalMs, + getNodeAgentReadyTimeoutMs, + getNodeBackendBaseUrl, + waitForNodeAgentReadyWith, +} from './node-agent-readiness'; import { recordNodeRoutingMetric } from './telemetry'; import { fetchVmAgentContainer, getVmAgentContainerConfig, markVmAgentContainerActiveWorkEndedBestEffort, markVmAgentContainerActiveWorkStarted, + markVmAgentContainerRequestInterrupted, } from './vm-agent-container'; const DEFAULT_NODE_AGENT_REQUEST_TIMEOUT_MS = 30_000; @@ -22,57 +35,53 @@ const DEFAULT_CF_CONTAINER_WAKE_TIMEOUT_MS = 120_000; // the wake/restore budget, which re-runs the same clone. See rule 43. const DEFAULT_CF_CONTAINER_CREATE_WORKSPACE_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; - -function getNodeBackendBaseUrl(nodeId: string, env: Env): string { - const protocol = env.VM_AGENT_PROTOCOL || 'https'; - const port = env.VM_AGENT_PORT || '8443'; - // Two-level subdomain ({nodeId}.vm.{domain}) bypasses Cloudflare same-zone routing. - // The wildcard Worker route *.{domain}/* matches exactly one subdomain level, - // so {nodeId}.vm.{domain} is NOT intercepted — requests reach the VM directly. - return `${protocol}://${nodeId.toLowerCase()}.vm.${env.BASE_DOMAIN}:${port}`; -} - interface NodeAgentRequestOptions extends RequestInit { userId: string; workspaceId?: string | null; requestTimeoutMs?: number; } +const RUNTIME_RECOVERY_CODES: ReadonlySet = new Set([ + 'RUNTIME_RECOVERING', + 'RUNTIME_REQUEST_INTERRUPTED', + 'RUNTIME_RECOVERY_DEGRADED', + 'RUNTIME_STOPPED', +]); + +function isRuntimeRecoveryCode(value: unknown): value is RuntimeRecoveryCode { + return typeof value === 'string' && RUNTIME_RECOVERY_CODES.has(value); +} + +function runtimeRecoveryStatus(code: RuntimeRecoveryCode): number { + if (code === 'RUNTIME_STOPPED') return 410; + if (code === 'RUNTIME_RECOVERING') return 503; + return 409; +} + +export class NodeAgentRequestError extends AppError { + constructor(statusCode: number, code: RuntimeRecoveryCode, message: string) { + super(statusCode, code, message); + this.name = 'NodeAgentRequestError'; + } +} + function requestInitWithoutSignal(options: RequestInit): RequestInit { const serializableOptions = { ...options }; delete serializableOptions.signal; return serializableOptions; } -export function getNodeAgentReadyTimeoutMs(env: { NODE_AGENT_READY_TIMEOUT_MS?: string }): number { - const parsed = env.NODE_AGENT_READY_TIMEOUT_MS - ? Number.parseInt(env.NODE_AGENT_READY_TIMEOUT_MS, 10) - : DEFAULT_NODE_AGENT_READY_TIMEOUT_MS; - if (!Number.isFinite(parsed) || parsed <= 0) { - return DEFAULT_NODE_AGENT_READY_TIMEOUT_MS; - } - return parsed; -} +export { getNodeAgentReadyPollIntervalMs, getNodeAgentReadyTimeoutMs }; -export function getNodeAgentReadyPollIntervalMs(env: { - NODE_AGENT_READY_POLL_INTERVAL_MS?: string; +export function getNodeAgentRequestTimeoutMs(env: { + NODE_AGENT_REQUEST_TIMEOUT_MS?: string; }): number { - const parsed = env.NODE_AGENT_READY_POLL_INTERVAL_MS - ? Number.parseInt(env.NODE_AGENT_READY_POLL_INTERVAL_MS, 10) - : DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS; - if (!Number.isFinite(parsed) || parsed <= 0) { - return DEFAULT_NODE_AGENT_READY_POLL_INTERVAL_MS; - } - return parsed; -} - -export function getNodeAgentRequestTimeoutMs(env: { NODE_AGENT_REQUEST_TIMEOUT_MS?: string }): number { 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 { +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); } @@ -86,62 +95,7 @@ export function getCfContainerCreateWorkspaceTimeoutMs(env: { } export async function waitForNodeAgentReady(nodeId: string, env: Env): Promise { - const timeoutMs = getNodeAgentReadyTimeoutMs(env); - const pollIntervalMs = getNodeAgentReadyPollIntervalMs(env); - const baseUrl = getNodeBackendBaseUrl(nodeId, env); - const healthUrl = `${baseUrl}/health`; - const deadline = Date.now() + timeoutMs; - - let lastError = ''; - - while (Date.now() < deadline) { - const remainingMs = deadline - Date.now(); - const requestTimeoutMs = Math.max(1, Math.min(pollIntervalMs, remainingMs)); - const controller = new AbortController(); - let timeoutHandle: ReturnType | null = null; - - try { - const requestTimeoutError = `request timeout after ${requestTimeoutMs}ms`; - const response = await Promise.race([ - fetchNodeAgent(nodeId, env, healthUrl, { method: 'GET', signal: controller.signal }, requestTimeoutMs), - new Promise((_resolve, reject) => { - timeoutHandle = setTimeout(() => { - controller.abort(); - reject(new Error(requestTimeoutError)); - }, requestTimeoutMs); - }), - ]); - if (response.ok) { - return; - } - - const responseBody = await response.text().catch(() => ''); - lastError = `HTTP ${response.status}${responseBody ? ` ${responseBody}` : ''}`.trim(); - } catch (err) { - if (err instanceof Error && err.message.startsWith('request timeout after ')) { - lastError = err.message; - } else if (err instanceof Error && err.message.startsWith('Request timed out after ')) { - lastError = err.message.replace('Request timed out after ', 'request timeout after '); - } else if (err instanceof Error && err.name === 'AbortError') { - lastError = `request timeout after ${requestTimeoutMs}ms`; - } else { - lastError = err instanceof Error ? err.message : String(err); - } - } finally { - if (timeoutHandle) { - clearTimeout(timeoutHandle); - } - } - - const nextRemainingMs = deadline - Date.now(); - if (nextRemainingMs <= 0) { - break; - } - await new Promise((resolve) => setTimeout(resolve, Math.min(pollIntervalMs, nextRemainingMs))); - } - - const details = lastError ? ` Last error: ${lastError}` : ''; - throw new Error(`Node Agent not reachable at ${healthUrl} within ${timeoutMs}ms.${details}`); + return waitForNodeAgentReadyWith(fetchNodeAgent, nodeId, env); } export async function nodeAgentRequest( @@ -180,7 +134,13 @@ export async function nodeAgentRequest( ); const requestTimeoutMs = options.requestTimeoutMs ?? getNodeAgentRequestTimeoutMs(env); - const response = await fetchNodeAgent(nodeId, env, url, { ...options, headers }, requestTimeoutMs); + const response = await fetchNodeAgent( + nodeId, + env, + url, + { ...options, headers }, + requestTimeoutMs + ); recordNodeRoutingMetric( { @@ -196,6 +156,20 @@ export async function nodeAgentRequest( if (!response.ok) { const body = await response.text().catch(() => ''); + let recoveryPayload: { error?: unknown; message?: unknown } | null = null; + try { + recoveryPayload = JSON.parse(body) as { error?: unknown; message?: unknown }; + } catch { + // Non-recovery Node Agent responses retain the existing generic handling. + } + if (recoveryPayload && isRuntimeRecoveryCode(recoveryPayload.error)) { + throw new NodeAgentRequestError( + runtimeRecoveryStatus(recoveryPayload.error), + recoveryPayload.error, + getRuntimeRecoveryMessage(recoveryPayload.error) + ); + } + // Detect Worker loop-back: when the vm-{nodeId} DNS record is missing, // the wildcard DNS record routes the request back to this API Worker, // which returns its own 404 format. Provide a clear error instead. @@ -260,6 +234,18 @@ export async function fetchNodeAgent( containerUrl.hostname = 'localhost'; containerUrl.port = String(vmAgentPort); + // Known limitation (tracked: tasks/backlog/2026-07-21-instant-container-request-timeout-cancellation.md): + // the loser of this race is NOT cancelled. The container fetch is a + // VmAgentContainer DO RPC (fetchVmAgentContainer -> stub.proxyHttp), and an + // AbortSignal cannot be wired to it — passing a signal into the container node + // fetch is a proven regression (SPIKE PR #1544 "avoid abort signal in container + // node fetch"), which is exactly why requestInitWithoutSignal strips it. So on + // timeout the DO RPC keeps running and may still succeed, yet we mark the + // runtime interrupted below. That is over-eager for a slow-but-healthy request. + // Genuine cancellation requires the per-request timeout to live INSIDE + // proxyHttp, wrapping this.containerFetch (a real fetch that honors signals) + // with an AbortController — which needs the budget plumbed through + // fetchVmAgentContainer. Deferred to the tracked backlog task. let timeoutHandle: ReturnType | null = null; try { const response = await Promise.race([ @@ -277,6 +263,22 @@ export async function fetchNodeAgent( }), ]); return response; + } catch (error) { + const timedOut = error instanceof Error && error.message.startsWith('Request timed out after '); + if (!timedOut) throw error; + + const recovery = await markVmAgentContainerRequestInterrupted(env, nodeId, { + method: options.method ?? 'GET', + errorName: 'request_timeout', + }).catch(() => null); + if (recovery?.code && recovery.message) { + throw new NodeAgentRequestError( + runtimeRecoveryStatus(recovery.code), + recovery.code, + getRuntimeRecoveryMessage(recovery.code) + ); + } + throw error; } finally { if (timeoutHandle) { clearTimeout(timeoutHandle); @@ -553,7 +555,10 @@ export async function sendPromptToAgentOnNode( } } -export { hibernateAgentSessionOnNode, restoreAgentSessionOnNode } from './node-agent-session-snapshots'; +export { + hibernateAgentSessionOnNode, + restoreAgentSessionOnNode, +} from './node-agent-session-snapshots'; /** * Cancel a running prompt on an agent session. @@ -588,7 +593,11 @@ export async function cancelAgentSessionOnNode( const statusMatch = msg.match(/failed:\s*(\d{3})/); const status = statusMatch?.[1] ? parseInt(statusMatch[1], 10) : 500; if (status === 409) { - await markVmAgentContainerActiveWorkEndedBestEffort(env, nodeId, 'cancel_agent_session_no_prompt'); + await markVmAgentContainerActiveWorkEndedBestEffort( + env, + nodeId, + 'cancel_agent_session_no_prompt' + ); } return { success: false, status }; } diff --git a/apps/api/src/services/session-snapshots.ts b/apps/api/src/services/session-snapshots.ts index 393209b7d..266624bee 100644 --- a/apps/api/src/services/session-snapshots.ts +++ b/apps/api/src/services/session-snapshots.ts @@ -25,6 +25,8 @@ export interface SessionSnapshotManifest { chatSessionId: string; workspaceId: string; agentSessionId?: string; + acpSessionId?: string; + agentType?: string; baseCommit?: string; status: SessionSnapshotStatus; degradation: SessionSnapshotDegradation; @@ -93,7 +95,9 @@ export function getSessionSnapshotConfig(env: Env): SessionSnapshotConfig { 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), + r2Prefix: sanitizeSnapshotPrefix( + env.SESSION_SNAPSHOT_R2_PREFIX || DEFAULT_SESSION_SNAPSHOT_R2_PREFIX + ), }; } @@ -116,7 +120,8 @@ export function buildSessionSnapshotR2Key( ): string { const { r2Prefix } = getSessionSnapshotConfig(env); const sessionKey = sanitizeKeySegment(chatSessionId); - const suffix = artifact === 'home' ? 'home.tar' : artifact === 'wip' ? 'wip.bundle' : 'manifest.json'; + const suffix = + artifact === 'home' ? 'home.tar' : artifact === 'wip' ? 'wip.bundle' : 'manifest.json'; return `${r2Prefix}/${sessionKey}/${suffix}`; } @@ -174,7 +179,10 @@ export async function prepareSessionSnapshot( } satisfies schema.NewSessionSnapshot; if (existing[0]) { - await db.update(schema.sessionSnapshots).set(row).where(eq(schema.sessionSnapshots.id, snapshotId)); + await db + .update(schema.sessionSnapshots) + .set(row) + .where(eq(schema.sessionSnapshots.id, snapshotId)); } else { await db.insert(schema.sessionSnapshots).values({ ...row, createdAt: now.toISOString() }); } @@ -214,11 +222,13 @@ export async function completeSessionSnapshot( 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), + ( + 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)); @@ -264,7 +274,10 @@ export async function recordSessionSnapshotRestoreResult( .where(eq(schema.sessionSnapshots.chatSessionId, input.chatSessionId)); } -export async function deleteSessionSnapshotArtifacts(env: Env, chatSessionId: string): Promise { +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) => { diff --git a/apps/api/src/services/vm-agent-container.ts b/apps/api/src/services/vm-agent-container.ts index 4b6cc0c8f..1ca0163a5 100644 --- a/apps/api/src/services/vm-agent-container.ts +++ b/apps/api/src/services/vm-agent-container.ts @@ -7,7 +7,9 @@ import { type VmAgentContainer, type VmAgentContainerLaunchConfig, type VmAgentContainerLaunchSecrets, + type VmAgentContainerRecoveryResult, } from '../durable-objects/vm-agent-container'; +import type { VmAgentContainerLifecycleInspection } from '../durable-objects/vm-agent-container-lifecycle'; import type { Env } from '../env'; import { log } from '../lib/logger'; import { errors } from '../middleware/error'; @@ -22,8 +24,12 @@ export interface VmAgentContainerConfig { export function getVmAgentContainerConfig(env: Env): VmAgentContainerConfig { return { enabled: (env.CF_CONTAINER_ENABLED ?? env.SANDBOX_ENABLED) === 'true', - vmAgentPort: Number.parseInt(env.CF_CONTAINER_VM_AGENT_PORT || env.SANDBOX_VM_AGENT_PORT || '8080', 10), - sleepAfter: env.CF_CONTAINER_SLEEP_AFTER || env.SANDBOX_SLEEP_AFTER || DEFAULT_CF_CONTAINER_SLEEP_AFTER, + vmAgentPort: Number.parseInt( + env.CF_CONTAINER_VM_AGENT_PORT || env.SANDBOX_VM_AGENT_PORT || '8080', + 10 + ), + sleepAfter: + env.CF_CONTAINER_SLEEP_AFTER || env.SANDBOX_SLEEP_AFTER || DEFAULT_CF_CONTAINER_SLEEP_AFTER, }; } @@ -73,6 +79,31 @@ export async function fetchVmAgentContainer( return container.proxyHttp(request, port); } +export async function resumeVmAgentContainer( + env: Env, + nodeId: string, + agentSessionId: string +): Promise { + const container = getVmAgentContainer(env, nodeId); + return container.resumeRuntime(agentSessionId); +} + +export async function inspectVmAgentContainerLifecycle( + env: Env, + nodeId: string +): Promise { + return getVmAgentContainer(env, nodeId).inspectLifecycle(); +} + +export async function markVmAgentContainerRequestInterrupted( + env: Env, + nodeId: string, + input: { method: string; errorName: string } +): Promise { + const container = getVmAgentContainer(env, nodeId); + return container.markRequestInterrupted(input); +} + export async function destroyVmAgentContainer(env: Env, nodeId: string): Promise { const container = getVmAgentContainer(env, nodeId); await container.destroyForUser(); diff --git a/apps/api/tests/unit/acp-session-heartbeat-timeout.test.ts b/apps/api/tests/unit/acp-session-heartbeat-timeout.test.ts new file mode 100644 index 000000000..0bf1c3a77 --- /dev/null +++ b/apps/api/tests/unit/acp-session-heartbeat-timeout.test.ts @@ -0,0 +1,108 @@ +import Database from 'better-sqlite3'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { runMigrations } from '../../src/durable-objects/migrations'; +import { + checkHeartbeatTimeouts, + computeHeartbeatAlarmTime, +} from '../../src/durable-objects/project-data/acp-sessions'; +import type { Env } from '../../src/durable-objects/project-data/types'; +import { createSqlStorage } from './durable-objects/sql-storage-test-utils'; + +describe('ACP heartbeat timeout policy', () => { + let db: Database.Database; + let sql: SqlStorage; + const now = Date.UTC(2026, 6, 21, 16, 35, 0); + + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(now); + db = new Database(':memory:'); + sql = createSqlStorage(db); + runMigrations(sql); + db.prepare( + `INSERT INTO chat_sessions + (id, topic, status, message_count, started_at, created_at, updated_at) + VALUES ('chat-1', 'Idle recovery', 'active', 0, ?, ?, ?)` + ).run(now - 600_000, now - 600_000, now - 600_000); + db.prepare( + `INSERT INTO acp_sessions + (id, chat_session_id, workspace_id, node_id, status, agent_type, + last_heartbeat_at, created_at, updated_at) + VALUES ('acp-1', 'chat-1', 'ws-1', 'node-1', 'running', 'claude-code', ?, ?, ?)` + ).run(now - 600_000, now - 600_000, now - 600_000); + }); + + afterEach(() => { + db.close(); + vi.useRealTimers(); + }); + + it('does not terminalize a stale ACP session while its runtime is resumable', async () => { + const transition = vi.fn().mockResolvedValue(undefined); + const shouldDeferTimeout = vi.fn().mockResolvedValue({ + defer: true, + reason: 'cf_container_sleeping', + }); + + const timedOut = await checkHeartbeatTimeouts( + sql, + { ACP_SESSION_DETECTION_WINDOW_MS: '300000' } as Env, + transition, + { shouldDeferTimeout } + ); + + expect(shouldDeferTimeout).toHaveBeenCalledWith( + expect.objectContaining({ id: 'acp-1', workspaceId: 'ws-1', nodeId: 'node-1' }) + ); + expect(transition).not.toHaveBeenCalled(); + expect(timedOut).toEqual([]); + }); + + it('still terminalizes a stale ACP session after the runtime is conclusively terminal', async () => { + const transition = vi.fn().mockResolvedValue(undefined); + + const timedOut = await checkHeartbeatTimeouts( + sql, + { ACP_SESSION_DETECTION_WINDOW_MS: '300000' } as Env, + transition, + { + shouldDeferTimeout: vi.fn().mockResolvedValue({ + defer: false, + reason: 'cf_container_error', + }), + } + ); + + expect(transition).toHaveBeenCalledWith( + 'acp-1', + 'interrupted', + expect.objectContaining({ actorType: 'alarm' }) + ); + expect(timedOut).toEqual([{ sessionId: 'acp-1', workspaceId: 'ws-1' }]); + }); + + it('defers conservatively when runtime inspection fails', async () => { + const transition = vi.fn().mockResolvedValue(undefined); + + const timedOut = await checkHeartbeatTimeouts( + sql, + { ACP_SESSION_DETECTION_WINDOW_MS: '300000' } as Env, + transition, + { + shouldDeferTimeout: vi.fn().mockRejectedValue(new Error('DO unavailable')), + } + ); + + expect(transition).not.toHaveBeenCalled(); + expect(timedOut).toEqual([]); + }); + + it('backs off the next alarm when a stale session was deliberately deferred', () => { + const nextAlarm = computeHeartbeatAlarmTime(sql, { + ACP_SESSION_DETECTION_WINDOW_MS: '300000', + } as Env); + + expect(nextAlarm).toBeGreaterThanOrEqual(now + 300_000); + }); +}); 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 2f0024ddf..ec5db3ab5 100644 --- a/apps/api/tests/unit/cf-container-runtime-contract.test.ts +++ b/apps/api/tests/unit/cf-container-runtime-contract.test.ts @@ -18,7 +18,11 @@ function readPackage(relPath: string): string { } function getPinnedAgentPackage(agentType: string): string { - const manifest = agentInstallManifest as Array<{ agentType: string; package: string; version: string }>; + const manifest = agentInstallManifest as Array<{ + agentType: string; + package: string; + version: string; + }>; const entry = manifest.find((candidate) => candidate.agentType === agentType); if (!entry) { throw new Error(`Missing agent install manifest entry for ${agentType}`); @@ -118,6 +122,7 @@ describe('cf-container runtime spike contracts', () => { it('keeps active raw container prompt work alive with a bounded renewActivityTimeout loop', () => { const containerDo = read('durable-objects/vm-agent-container.ts'); + const activeWork = read('durable-objects/vm-agent-container-active-work.ts'); const containerService = read('services/vm-agent-container.ts'); const nodeAgent = read('services/node-agent.ts'); const activityCallback = read('routes/projects/agent-activity-callback.ts'); @@ -134,15 +139,15 @@ describe('cf-container runtime spike contracts', () => { expect(containerDo).toContain( 'await this.schedule(Math.max(1, Math.ceil(delayMs / 1000)), KEEPALIVE_CALLBACK)' ); - expect(containerDo).toContain("endReason: 'keepalive_deadline_exceeded'"); + expect(activeWork).toContain("endReason: 'keepalive_deadline_exceeded'"); expect(nodeAgent).toContain('markVmAgentContainerActiveWorkStarted(env, nodeId'); expect(nodeAgent).toContain("reason: 'start_agent_session'"); expect(nodeAgent).toContain("reason: 'send_prompt'"); expect(nodeAgent).toContain( "markVmAgentContainerActiveWorkEndedBestEffort(env, nodeId, 'cancel_agent_session')" ); - expect(nodeAgent).toContain( - "markVmAgentContainerActiveWorkEndedBestEffort(env, nodeId, 'cancel_agent_session_no_prompt')" + expect(nodeAgent).toMatch( + /markVmAgentContainerActiveWorkEndedBestEffort\(\s*env,\s*nodeId,\s*'cancel_agent_session_no_prompt'\s*\)/ ); expect(nodeAgent).toContain( "markVmAgentContainerActiveWorkEndedBestEffort(env, nodeId, 'stop_agent_session')" @@ -153,10 +158,17 @@ describe('cf-container runtime spike contracts', () => { it('classifies raw container idle expiration as sleeping instead of crash/error', () => { const containerDo = read('durable-objects/vm-agent-container.ts'); + const containerLifecycle = read('durable-objects/vm-agent-container-lifecycle.ts'); const chatResolver = read('routes/chat-workspace-resolver.ts'); + // Recovery-state-machine behavior (onStop ignoring 'recovering', the resume/ + // ensureAwake wiring, and the RUNTIME_RECOVERY_DEGRADED_MESSAGE / + // RUNTIME_STOPPED_MESSAGE codes) is now exercised behaviorally in + // durable-objects/vm-agent-container-recovery.test.ts (wake concurrency, + // persistence, and lifecycle-status parity suites); those source-contract + // string assertions were removed here to avoid duplicate, non-behavioral + // coverage (rule 02). Structural-only checks (type unions, storage keys) stay. expect(containerDo).toContain('override async onStop'); - expect(containerDo).toContain("if (status === 'expired' || status === 'sleeping')"); expect(containerDo).toContain('override async onError'); expect(containerDo).toContain('override async onActivityExpired'); expect(containerDo).toContain( @@ -165,10 +177,8 @@ describe('cf-container runtime spike contracts', () => { expect(containerDo).toContain( "await this.ctx.storage.put('lifecycleStatus', 'sleeping' satisfies LifecycleStatus)" ); - expect(containerDo).toContain("status: 'sleeping'"); - expect(containerDo).toContain("if (lifecycleStatus === 'sleeping')"); - expect(containerDo).toContain('const wake = await this.ensureAwake()'); - expect(containerDo).toContain('WAKE_DEGRADED_RESPONSE'); + expect(containerLifecycle).toContain("| 'sleeping'"); + expect(containerDo).toContain("status === 'sleeping' ? 'idle' : 'error'"); expect(containerDo).toContain( "await this.ctx.storage.put('lifecycleStatus', 'launching' satisfies LifecycleStatus)" ); @@ -179,17 +189,14 @@ describe('cf-container runtime spike contracts', () => { "inArray(schema.workspaces.status, ['running', 'recovery', 'sleeping'])" ); expect(chatResolver).toContain( - "workspace.nodeStatus !== 'running' && workspace.nodeStatus !== 'sleeping'" + "['running', 'sleeping', 'recovery', 'error'].includes(workspace.nodeStatus)" ); - expect(chatResolver).toContain("inArray(schema.agentSessions.status, ['running', 'sleeping'])"); + expect(chatResolver).toContain('inArray(schema.agentSessions.status, agentStatuses)'); 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'); expect(containerDo).toContain('CALLBACK_TOKEN: secrets.nodeCallbackToken'); - expect(containerDo).toContain("status === 'stopped' ? 'stopped' : 'error'"); + expect(containerDo).toContain('persistRuntimeEnded(this.env, config, status, message)'); expect(containerDo).not.toContain( 'Container idle timeout expired; start a new instant session.' ); diff --git a/apps/api/tests/unit/durable-objects/vm-agent-container-launch-env.test.ts b/apps/api/tests/unit/durable-objects/vm-agent-container-launch-env.test.ts index e222d96b3..b5fc3cb0d 100644 --- a/apps/api/tests/unit/durable-objects/vm-agent-container-launch-env.test.ts +++ b/apps/api/tests/unit/durable-objects/vm-agent-container-launch-env.test.ts @@ -39,8 +39,9 @@ function makeFake(env: Record) { const fake = { env, startAndWaitForPorts, + startRuntime: (VmAgentContainer.prototype as unknown as { startRuntime: unknown }).startRuntime, clearKeepaliveSchedule: vi.fn().mockResolvedValue(undefined), - getPortReadyTimeoutMs: () => 30_000, + getRuntimeSettings: () => ({ portReadyTimeoutMs: 30_000 }), ctx: { storage: { put: vi.fn().mockResolvedValue(undefined), @@ -93,3 +94,33 @@ describe('VmAgentContainer.launch env passthrough', () => { expect(envVars.NODE_ROLE).toBe('standalone'); }); }); + +function storagePutMock(fake: unknown): ReturnType { + return (fake as { ctx: { storage: { put: ReturnType } } }).ctx.storage.put; +} + +describe('VmAgentContainer.launch lifecycle status', () => { + it('T3: marks lifecycleStatus running after startRuntime succeeds', async () => { + const { fake, startAndWaitForPorts } = makeFake({}); + + await callLaunch(fake); + + expect(startAndWaitForPorts).toHaveBeenCalledTimes(1); + const put = storagePutMock(fake); + expect(put).toHaveBeenCalledWith('lifecycleStatus', 'launching'); + expect(put).toHaveBeenCalledWith('lifecycleStatus', 'running'); + expect(put).not.toHaveBeenCalledWith('lifecycleStatus', 'error'); + }); + + it('T3: marks lifecycleStatus error and rethrows when startRuntime fails', async () => { + const { fake, startAndWaitForPorts } = makeFake({}); + const boom = Object.assign(new Error('port ready timeout'), { name: 'PortReadyError' }); + startAndWaitForPorts.mockRejectedValueOnce(boom); + + await expect(callLaunch(fake)).rejects.toBe(boom); + + const put = storagePutMock(fake); + expect(put).toHaveBeenCalledWith('lifecycleStatus', 'error'); + expect(put).not.toHaveBeenCalledWith('lifecycleStatus', 'running'); + }); +}); diff --git a/apps/api/tests/unit/durable-objects/vm-agent-container-recovery.test.ts b/apps/api/tests/unit/durable-objects/vm-agent-container-recovery.test.ts new file mode 100644 index 000000000..d627e6400 --- /dev/null +++ b/apps/api/tests/unit/durable-objects/vm-agent-container-recovery.test.ts @@ -0,0 +1,671 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const recoveryMocks = vi.hoisted(() => ({ + loadContext: vi.fn(), + persistRecovering: vi.fn(), + persistRecovered: vi.fn(), + persistFailed: vi.fn(), + signNodeCallbackToken: vi.fn(), + signCallbackToken: vi.fn(), + signNodeManagementToken: vi.fn(), +})); + +vi.mock('../../../src/durable-objects/vm-agent-container-recovery', async (importOriginal) => { + const actual = + await importOriginal< + typeof import('../../../src/durable-objects/vm-agent-container-recovery') + >(); + return { + ...actual, + loadRuntimeRecoveryContext: recoveryMocks.loadContext, + persistRuntimeRecovering: recoveryMocks.persistRecovering, + persistRuntimeRecovered: recoveryMocks.persistRecovered, + persistRuntimeRecoveryFailed: recoveryMocks.persistFailed, + }; +}); + +vi.mock('../../../src/services/jwt', () => ({ + signNodeCallbackToken: recoveryMocks.signNodeCallbackToken, + signCallbackToken: recoveryMocks.signCallbackToken, + signNodeManagementToken: recoveryMocks.signNodeManagementToken, +})); + +import { VmAgentContainer } from '../../../src/durable-objects/vm-agent-container'; +import { + RUNTIME_RECOVERY_DEGRADED_MESSAGE, + type RuntimeRecoveryState, +} from '../../../src/durable-objects/vm-agent-container-recovery'; + +const launchConfig = { + nodeId: 'node-1', + workspaceId: 'workspace-1', + projectId: 'project-1', + chatSessionId: 'chat-1', + repository: 'owner/repo', + branch: 'main', + workspaceDir: '/workspaces/repo', + controlPlaneUrl: 'https://api.example.test', + vmAgentPort: 8080, +}; + +const runtimeContext = { + userId: 'user-1', + chatSessionId: 'chat-1', + agentSessionId: 'agent-session-1', + agentType: 'codex', +}; + +type PrivateContainer = { + ensureAwake: (this: unknown) => Promise; + beginUnexpectedRecovery: (this: unknown, input: unknown) => Promise; + wakeFromSnapshot: (this: unknown, recovery: RuntimeRecoveryState) => Promise; + degradeRecovery: (this: unknown, ...args: unknown[]) => Promise; + exhaustRecovery: (this: unknown, ...args: unknown[]) => Promise; + withLifecycleLock: (this: unknown, operation: () => Promise) => Promise; + prepareForRequest: (this: unknown) => Promise; +}; + +const privateContainer = VmAgentContainer.prototype as unknown as PrivateContainer; + +function makeStorage(initialLifecycle: string) { + const values = new Map([ + ['lifecycleStatus', initialLifecycle], + ['launchConfig', launchConfig], + ]); + return { + values, + storage: { + get: vi.fn(async (key: string) => values.get(key)), + put: vi.fn(async (key: string, value: unknown) => { + values.set(key, value); + }), + delete: vi.fn(async (key: string) => { + values.delete(key); + }), + }, + }; +} + +function makeRecoveryFake(input?: { + lifecycle?: string; + restoreResponse?: Response; + maxAttempts?: number; +}) { + const { values, storage } = makeStorage(input?.lifecycle ?? 'sleeping'); + const fake = { + env: { CF_CONTAINER_RECOVERY_MAX_ATTEMPTS: String(input?.maxAttempts ?? 2) }, + ctx: { storage }, + wakeChain: Promise.resolve(), + lifecycleChain: Promise.resolve(), + defaultPort: 8080, + ensureAwake: privateContainer.ensureAwake, + beginUnexpectedRecovery: privateContainer.beginUnexpectedRecovery, + wakeFromSnapshot: privateContainer.wakeFromSnapshot, + degradeRecovery: privateContainer.degradeRecovery, + exhaustRecovery: privateContainer.exhaustRecovery, + withLifecycleLock: privateContainer.withLifecycleLock, + prepareForRequest: privateContainer.prepareForRequest, + getRuntimeSettings: () => ({ + portReadyTimeoutMs: 30_000, + activeWorkMaxMs: 2 * 60 * 60 * 1000, + keepaliveRenewIntervalMs: 5 * 60 * 1000, + recoveryMaxAttempts: input?.maxAttempts ?? 2, + }), + clearKeepaliveSchedule: vi.fn().mockResolvedValue(undefined), + markActiveWorkEnded: vi.fn().mockResolvedValue(undefined), + startRuntime: vi.fn().mockResolvedValue(undefined), + stop: vi.fn().mockResolvedValue(undefined), + containerFetch: vi + .fn() + .mockResolvedValue( + input?.restoreResponse ?? + Response.json({ status: 'restored', degradation: 'none', skipped: [] }) + ), + getState: vi.fn().mockResolvedValue({ status: 'running' }), + }; + return { fake, values, storage }; +} + +async function callEnsureAwake(fake: unknown) { + return privateContainer.ensureAwake.call(fake) as Promise<{ + ok: boolean; + status: string; + code?: string; + }>; +} + +async function callResumeRuntime(fake: unknown) { + return ( + VmAgentContainer.prototype as unknown as { + resumeRuntime: ( + this: unknown, + agentSessionId?: string + ) => Promise<{ ok: boolean; status: string; code?: string }>; + } + ).resumeRuntime.call(fake, 'agent-session-1'); +} + +function callProxyHttp(fake: unknown, request: Request): Promise { + return ( + VmAgentContainer.prototype as unknown as { + proxyHttp: (this: unknown, request: Request) => Promise; + } + ).proxyHttp.call(fake, request); +} + +beforeEach(() => { + vi.clearAllMocks(); + recoveryMocks.loadContext.mockResolvedValue(runtimeContext); + recoveryMocks.persistRecovering.mockResolvedValue(undefined); + recoveryMocks.persistRecovered.mockResolvedValue(undefined); + recoveryMocks.persistFailed.mockResolvedValue(undefined); + recoveryMocks.signNodeCallbackToken.mockResolvedValue('fresh-node-token'); + recoveryMocks.signCallbackToken.mockResolvedValue('fresh-workspace-token'); + recoveryMocks.signNodeManagementToken.mockResolvedValue({ token: 'management-token' }); +}); + +describe('VmAgentContainer snapshot recovery state machine', () => { + it('reconciles stale D1 state only after proving the target SessionHost is live', async () => { + const { fake } = makeRecoveryFake({ lifecycle: 'running' }); + fake.containerFetch.mockResolvedValueOnce( + Response.json({ + sessions: [{ id: 'agent-session-1', status: 'running', hostStatus: 'ready' }], + }) + ); + + const result = await callResumeRuntime(fake); + + expect(result).toEqual({ ok: true, status: 'running' }); + expect(recoveryMocks.loadContext).toHaveBeenCalledWith(fake.env, { + workspaceId: 'workspace-1', + preferredAgentSessionId: 'agent-session-1', + }); + expect(recoveryMocks.signNodeManagementToken).toHaveBeenCalledWith( + 'user-1', + 'node-1', + 'workspace-1', + fake.env + ); + const probeRequest = fake.containerFetch.mock.calls[0]?.[0] as Request; + expect(probeRequest.headers.get('Authorization')).toBe('Bearer management-token'); + expect(probeRequest.headers.get('X-SAM-Workspace-Id')).toBe('workspace-1'); + expect(fake.startRuntime).not.toHaveBeenCalled(); + expect(recoveryMocks.persistRecovered).toHaveBeenCalledWith( + fake.env, + expect.objectContaining({ agentSessionId: 'agent-session-1' }), + 'none' + ); + }); + + it('restores instead of falsely resuming when D1 says recovery but SessionHost is missing', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + fake.containerFetch + .mockResolvedValueOnce( + Response.json({ sessions: [{ id: 'agent-session-1', status: 'running' }] }) + ) + .mockResolvedValueOnce( + Response.json({ status: 'restored', degradation: 'none', skipped: [] }) + ); + + const result = await callResumeRuntime(fake); + + expect(result).toEqual({ ok: true, status: 'running' }); + expect(fake.startRuntime).toHaveBeenCalledTimes(1); + expect(recoveryMocks.persistRecovering).toHaveBeenCalledTimes(1); + expect(values.has('runtimeRecovery')).toBe(false); + }); + + it('does not reconcile running when an explicit stop crosses the SessionHost probe', async () => { + let finishProbe!: (response: Response) => void; + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + fake.containerFetch.mockImplementationOnce( + () => + new Promise((resolve) => { + finishProbe = resolve; + }) + ); + const stopForUser = ( + VmAgentContainer.prototype as unknown as { + stopForUser: (this: unknown) => Promise; + } + ).stopForUser; + + const resume = callResumeRuntime(fake); + await vi.waitFor(() => expect(fake.containerFetch).toHaveBeenCalledTimes(1)); + await stopForUser.call(fake); + finishProbe( + Response.json({ + sessions: [{ id: 'agent-session-1', status: 'running', hostStatus: 'ready' }], + }) + ); + + await expect(resume).resolves.toMatchObject({ + ok: false, + status: 'stopped', + code: 'RUNTIME_STOPPED', + }); + expect(recoveryMocks.persistRecovered).not.toHaveBeenCalled(); + expect(values.get('lifecycleStatus')).toBe('stopping'); + }); + + it('cold-wakes, reinjects fresh callback tokens, restores, then reconciles running', async () => { + const { fake, values } = makeRecoveryFake(); + + const result = await callEnsureAwake(fake); + + expect(result).toEqual({ ok: true, status: 'running' }); + expect(fake.startRuntime).toHaveBeenCalledWith(launchConfig, { + nodeCallbackToken: 'fresh-node-token', + }); + const restoreRequest = fake.containerFetch.mock.calls[0]?.[0] as Request; + await expect(restoreRequest.json()).resolves.toMatchObject({ + chatSessionId: 'chat-1', + runtime: 'cf-container', + agentType: 'codex', + workspaceCallbackToken: 'fresh-workspace-token', + }); + expect(recoveryMocks.persistRecovering).toHaveBeenCalledWith( + fake.env, + expect.objectContaining({ + nodeId: 'node-1', + workspaceId: 'workspace-1', + agentSessionId: 'agent-session-1', + }) + ); + expect(recoveryMocks.persistRecovered).toHaveBeenCalledWith( + fake.env, + expect.objectContaining({ chatSessionId: 'chat-1' }), + 'none' + ); + expect(values.get('lifecycleStatus')).toBe('running'); + expect(values.has('runtimeRecovery')).toBe(false); + }); + + it('keeps missing snapshots degraded until the bounded attempt is exhausted', async () => { + const { fake, values } = makeRecoveryFake({ + maxAttempts: 1, + restoreResponse: Response.json({ error: 'SNAPSHOT_NOT_FOUND' }, { status: 404 }), + }); + + const result = await callEnsureAwake(fake); + + expect(result).toMatchObject({ + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + }); + expect(recoveryMocks.persistRecovered).not.toHaveBeenCalled(); + expect(recoveryMocks.persistFailed).toHaveBeenCalledTimes(1); + expect(values.get('lifecycleStatus')).toBe('error'); + expect(values.get('runtimeRecovery')).toMatchObject({ + phase: 'exhausted', + lastFailure: { kind: 'restore_http', httpStatus: 404 }, + }); + }); + + it('treats a corrupt successful restore body as degraded, never as true resume', async () => { + const { fake, values } = makeRecoveryFake({ + restoreResponse: new Response('{not-json', { status: 200 }), + }); + + const result = await callEnsureAwake(fake); + + expect(result).toMatchObject({ ok: false, status: 'degraded' }); + expect(recoveryMocks.persistRecovered).not.toHaveBeenCalled(); + expect(recoveryMocks.persistFailed).not.toHaveBeenCalled(); + expect(values.get('runtimeRecovery')).toMatchObject({ + phase: 'degraded', + lastFailure: { kind: 'restore_status' }, + }); + }); + + it('keeps an explicit stop terminal when it crosses an active restore', async () => { + let finishRestore!: (response: Response) => void; + const { fake, values } = makeRecoveryFake(); + fake.containerFetch.mockImplementationOnce( + () => + new Promise((resolve) => { + finishRestore = resolve; + }) + ); + const stopForUser = ( + VmAgentContainer.prototype as unknown as { + stopForUser: (this: unknown) => Promise; + } + ).stopForUser; + + const wake = callEnsureAwake(fake); + await vi.waitFor(() => expect(fake.containerFetch).toHaveBeenCalledTimes(1)); + + await stopForUser.call(fake); + finishRestore( + Response.json({ + status: 'restored', + degradation: 'none', + skipped: [], + }) + ); + + await expect(wake).resolves.toMatchObject({ + ok: false, + status: 'stopped', + code: 'RUNTIME_STOPPED', + }); + expect(recoveryMocks.persistRecovered).not.toHaveBeenCalled(); + expect(values.get('lifecycleStatus')).toBe('stopping'); + expect(values.has('runtimeRecovery')).toBe(false); + // CF2: startRuntime() already launched a fresh container before the restore. + // The completed-block re-check must tear it down (2nd stop), not just return + // terminal — otherwise the just-stopped session leaks compute until sleepAfter. + // Pre-fix this was 1 (only stopForUser's stop). + expect(fake.stop).toHaveBeenCalledTimes(2); + }); +}); + +describe('VmAgentContainer wake concurrency and persistence', () => { + const stopForUser = ( + VmAgentContainer.prototype as unknown as { + stopForUser: (this: unknown) => Promise; + } + ).stopForUser; + + it('CF1: does not re-arm recovery when an explicit stop lands right after beginUnexpectedRecovery', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'sleeping' }); + const realBegin = privateContainer.beginUnexpectedRecovery; + // Land the stop in the EARLIER race window: after beginUnexpectedRecovery + // returns (recovery created, lifecycle 'recovering') but BEFORE ensureAwake's + // unlocked continuation reads the terminal guard and bumps the phase to + // 'waking'. beginUnexpectedRecovery and stopForUser share the lifecycleChain, + // so the stop serializes after begin releases the lock. + fake.beginUnexpectedRecovery = async function (this: unknown, input: unknown) { + const result = await realBegin.call(this, input); + await stopForUser.call(this); + return result; + }; + + const result = await callEnsureAwake(fake); + + expect(result).toMatchObject({ ok: false, status: 'stopped', code: 'RUNTIME_STOPPED' }); + // Pre-fix (guard only checked 'stopped') the continuation flipped to 'waking' + // and woke the just-stopped container. These assertions go red pre-fix. + expect(fake.startRuntime).not.toHaveBeenCalled(); + expect(recoveryMocks.persistRecovered).not.toHaveBeenCalled(); + expect(values.get('lifecycleStatus')).toBe('stopping'); + expect(values.has('runtimeRecovery')).toBe(false); + }); + + it('T10: carries recovery.attempts 1->2 across persisted storage before exhausting', async () => { + const { fake, values } = makeRecoveryFake({ + maxAttempts: 2, + restoreResponse: Response.json({ error: 'SNAPSHOT_NOT_FOUND' }, { status: 404 }), + }); + + // First follow-up during the outage: one bounded wake attempt, degraded but + // NOT yet exhausted. + const first = await callEnsureAwake(fake); + expect(first).toMatchObject({ ok: false, status: 'degraded' }); + expect(values.get('runtimeRecovery')).toMatchObject({ phase: 'degraded', attempts: 1 }); + expect(recoveryMocks.persistFailed).not.toHaveBeenCalled(); + + // Second follow-up reuses the persisted record: attempts 1 -> 2, and only now + // exhausts. + const second = await callEnsureAwake(fake); + expect(second).toMatchObject({ ok: false, status: 'degraded' }); + expect(values.get('runtimeRecovery')).toMatchObject({ phase: 'exhausted', attempts: 2 }); + expect(values.get('lifecycleStatus')).toBe('error'); + expect(recoveryMocks.persistFailed).toHaveBeenCalledTimes(1); + }); + + it('T11: honors a pre-populated recovery record left by onStop without resetting it', async () => { + const { fake, values } = makeRecoveryFake({ + lifecycle: 'recovering', + restoreResponse: Response.json({ error: 'SNAPSHOT_NOT_FOUND' }, { status: 404 }), + maxAttempts: 2, + }); + // Shape onStop leaves behind: phase 'pending', trigger 'stop', a manual_retry + // disposition, attempts 0. ensureAwake must reuse it, not start fresh. + values.set('runtimeRecovery', { + version: 1, + phase: 'pending', + trigger: 'stop', + cause: { kind: 'container_stop', reason: 'runtime_signal', exitCode: 137 }, + attempts: 0, + promptDisposition: 'manual_retry', + agentSessionId: 'agent-session-1', + startedAt: 1000, + updatedAt: 1000, + }); + + await callEnsureAwake(fake); + + // beginUnexpectedRecovery (which would call persistRecovering) must NOT run — + // the existing trigger/cause/disposition/startedAt survive; only the wake + // bump advances attempts 0 -> 1. + expect(recoveryMocks.persistRecovering).not.toHaveBeenCalled(); + expect(values.get('runtimeRecovery')).toMatchObject({ + trigger: 'stop', + cause: { kind: 'container_stop', reason: 'runtime_signal', exitCode: 137 }, + promptDisposition: 'manual_retry', + agentSessionId: 'agent-session-1', + startedAt: 1000, + attempts: 1, + }); + }); + + it('S6: re-invoking an exhausted recovery returns degraded without re-running the D1 batch', async () => { + const { fake, values } = makeRecoveryFake({ + maxAttempts: 1, + restoreResponse: Response.json({ error: 'SNAPSHOT_NOT_FOUND' }, { status: 404 }), + }); + + const first = await callEnsureAwake(fake); + expect(first).toMatchObject({ ok: false, status: 'degraded' }); + expect(values.get('runtimeRecovery')).toMatchObject({ phase: 'exhausted' }); + expect(recoveryMocks.persistFailed).toHaveBeenCalledTimes(1); + + // Every later request hits an exhausted record. Pre-fix this re-ran + // exhaustRecovery() (a full D1 batch + persistRuntimeRecoveryFailed) each time. + const second = await callEnsureAwake(fake); + expect(second).toMatchObject({ ok: false, status: 'degraded', code: 'RUNTIME_RECOVERY_DEGRADED' }); + expect(recoveryMocks.persistFailed).toHaveBeenCalledTimes(1); + }); + + it('T12: degrades (not 500s) when the recovery context loader throws mid-wake', async () => { + const { fake, values } = makeRecoveryFake({ + restoreResponse: Response.json({ error: 'SNAPSHOT_NOT_FOUND' }, { status: 404 }), + }); + // beginUnexpectedRecovery resolves the context; the wakeFromSnapshot reload + // then throws a D1 error. Pre-fix that escaped as an uncaught rejection (500). + recoveryMocks.loadContext + .mockReset() + .mockResolvedValueOnce(runtimeContext) + .mockRejectedValueOnce(new Error('D1 unavailable')); + + const result = await callEnsureAwake(fake); + + expect(result).toMatchObject({ + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }); + expect(values.get('runtimeRecovery')).toMatchObject({ + phase: 'degraded', + lastFailure: { kind: 'unexpected' }, + }); + }); +}); + +describe('VmAgentContainer wake-path lifecycle status parity', () => { + // Parity contract: any lifecycleStatus the wake path can set as an intermediate + // state must never let a later unexpected onStop/onError START A SECOND recovery. + // onStop enforces this via its ignore-list; onError relies on + // beginUnexpectedRecovery's existing-record idempotency. vm-agent-container- + // lifecycle.ts is read-only, so these are literals mirrored from the wake path: + // ensureAwake -> 'waking' + // wakeFromSnapshot -> 'restoring' + // beginUnexpectedRecovery-> 'recovering' + // degradeRecovery -> 'degraded' + // A NEW wake-path status added without preserving this guard fails these tests. + const WAKE_PATH_INTERMEDIATE_STATUSES = [ + 'recovering', + 'waking', + 'restoring', + 'degraded', + ] as const; + + const onStop = ( + VmAgentContainer.prototype as unknown as { + onStop: ( + this: unknown, + input: { exitCode: number; reason: 'exit' | 'runtime_signal' } + ) => Promise; + } + ).onStop; + const onError = ( + VmAgentContainer.prototype as unknown as { + onError: (this: unknown, error: unknown) => Promise; + } + ).onError; + + it.each(WAKE_PATH_INTERMEDIATE_STATUSES)( + 'onStop ignores mid-wake status %s and starts no fresh recovery', + async (status) => { + const { fake, values } = makeRecoveryFake({ lifecycle: status }); + // No pre-existing recovery: if onStop did NOT ignore this status it would + // call beginUnexpectedRecovery -> persistRecovering, which this asserts against. + await onStop.call(fake, { exitCode: 0, reason: 'runtime_signal' }); + expect(recoveryMocks.persistRecovering).not.toHaveBeenCalled(); + expect(values.get('lifecycleStatus')).toBe(status); + } + ); + + it.each(WAKE_PATH_INTERMEDIATE_STATUSES)( + 'onError does not duplicate an in-flight recovery at mid-wake status %s', + async (status) => { + const { fake, values } = makeRecoveryFake({ lifecycle: status }); + // Mid-wake there is always a recovery record; onError must reuse it, not + // start a second one (beginUnexpectedRecovery returns the existing record). + values.set('runtimeRecovery', { + version: 1, + phase: 'waking', + trigger: 'error', + cause: { kind: 'container_error', errorName: 'RuntimeError' }, + attempts: 1, + promptDisposition: 'none', + agentSessionId: 'agent-session-1', + startedAt: 1, + updatedAt: 1, + }); + await onError.call(fake, Object.assign(new Error('boom'), { name: 'RuntimeError' })); + expect(recoveryMocks.persistRecovering).not.toHaveBeenCalled(); + expect(values.get('runtimeRecovery')).toMatchObject({ trigger: 'error', attempts: 1 }); + } + ); +}); + +describe('VmAgentContainer replacement classification', () => { + it('classifies duplicate runtime_signal callbacks once without calling them rollout', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + const onStop = ( + VmAgentContainer.prototype as unknown as { + onStop: ( + this: unknown, + input: { exitCode: number; reason: 'runtime_signal' } + ) => Promise; + } + ).onStop; + + await onStop.call(fake, { exitCode: 0, reason: 'runtime_signal' }); + await onStop.call(fake, { exitCode: 0, reason: 'runtime_signal' }); + + expect(recoveryMocks.persistRecovering).toHaveBeenCalledTimes(1); + expect(values.get('runtimeRecovery')).toMatchObject({ + trigger: 'stop', + cause: { kind: 'container_stop', reason: 'runtime_signal', exitCode: 0 }, + }); + }); + + it('classifies a true container error generically and preserves recovery eligibility', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + const onError = ( + VmAgentContainer.prototype as unknown as { + onError: (this: unknown, error: unknown) => Promise; + } + ).onError; + + await onError.call( + fake, + Object.assign(new Error('sensitive detail'), { name: 'RuntimeError' }) + ); + + expect(values.get('runtimeRecovery')).toMatchObject({ + trigger: 'error', + cause: { kind: 'container_error', errorName: 'RuntimeError' }, + }); + }); + + it('keeps an explicit stop terminal and does not create recovery state', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'stopping' }); + const markRuntimeEnded = vi.fn().mockResolvedValue(undefined); + Object.assign(fake, { markRuntimeEnded }); + const onStop = ( + VmAgentContainer.prototype as unknown as { + onStop: (this: unknown, input: { exitCode: number; reason: 'exit' }) => Promise; + } + ).onStop; + + await onStop.call(fake, { exitCode: 0, reason: 'exit' }); + + expect(markRuntimeEnded).toHaveBeenCalledWith('stopped', 'Container stopped by user request'); + expect(recoveryMocks.persistRecovering).not.toHaveBeenCalled(); + expect(values.get('lifecycleStatus')).toBe('stopped'); + }); + + it('does not replay a prompt whose request crosses replacement', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + fake.containerFetch.mockRejectedValueOnce( + Object.assign(new Error('socket reset'), { name: 'TypeError' }) + ); + const request = new Request( + 'http://container/workspaces/workspace-1/agent-sessions/agent-session-1/prompt', + { + method: 'POST', + body: JSON.stringify({ prompt: 'continue' }), + } + ); + + const response = await callProxyHttp(fake, request); + + expect(response.status).toBe(409); + await expect(response.json()).resolves.toMatchObject({ + error: 'RUNTIME_REQUEST_INTERRUPTED', + }); + expect(fake.containerFetch).toHaveBeenCalledTimes(1); + expect(values.get('runtimeRecovery')).toMatchObject({ + trigger: 'request', + promptDisposition: 'manual_retry', + cause: { kind: 'transport_interrupted', errorName: 'TypeError' }, + }); + }); + + it('classifies a missing SessionHost before another prompt can be sent', async () => { + const { fake, values } = makeRecoveryFake({ lifecycle: 'running' }); + fake.containerFetch.mockResolvedValueOnce( + Response.json({ error: 'no active agent session found' }, { status: 404 }) + ); + + const response = await callProxyHttp( + fake, + new Request('http://container/workspaces/workspace-1/agent-sessions/agent-session-1/prompt', { + method: 'POST', + }) + ); + + expect(response.status).toBe(409); + expect(values.get('runtimeRecovery')).toMatchObject({ + promptDisposition: 'manual_retry', + cause: { kind: 'missing_session_host', httpStatus: 404 }, + }); + }); +}); diff --git a/apps/api/tests/unit/durable-objects/vm-agent-container-runtime-settings.test.ts b/apps/api/tests/unit/durable-objects/vm-agent-container-runtime-settings.test.ts new file mode 100644 index 000000000..45f6bd6cc --- /dev/null +++ b/apps/api/tests/unit/durable-objects/vm-agent-container-runtime-settings.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it } from 'vitest'; + +import { + parsePositiveRuntimeSetting, + resolveRuntimeSettings, +} from '../../../src/durable-objects/vm-agent-container-runtime'; +import type { Env } from '../../../src/env'; + +const DEFAULTS = { + portReadyTimeoutMs: 30_000, + activeWorkMaxMs: 2 * 60 * 60 * 1000, + keepaliveRenewIntervalMs: 5 * 60 * 1000, + recoveryMaxAttempts: 2, +}; + +function envWith(overrides: Record = {}): Env { + return overrides as unknown as Env; +} + +describe('parsePositiveRuntimeSetting', () => { + it('returns the fallback when unset', () => { + expect(parsePositiveRuntimeSetting(undefined, 42)).toBe(42); + }); + + it('honors a valid positive override', () => { + expect(parsePositiveRuntimeSetting('7', 42)).toBe(7); + }); + + it.each([ + ['zero', '0'], + ['negative', '-3'], + ['non-numeric', 'banana'], + ['empty string', ''], + ])('falls back on invalid input (%s)', (_label, raw) => { + expect(parsePositiveRuntimeSetting(raw, 42)).toBe(42); + }); +}); + +describe('resolveRuntimeSettings', () => { + it('uses provided defaults when no env overrides are set', () => { + expect(resolveRuntimeSettings(envWith(), DEFAULTS)).toEqual(DEFAULTS); + }); + + it('honors each env override independently', () => { + const settings = resolveRuntimeSettings( + envWith({ + CF_CONTAINER_PORT_READY_TIMEOUT_MS: '1000', + CF_CONTAINER_ACTIVE_WORK_MAX_MS: '2000', + CF_CONTAINER_KEEPALIVE_RENEW_INTERVAL_MS: '3000', + CF_CONTAINER_RECOVERY_MAX_ATTEMPTS: '4', + }), + DEFAULTS + ); + expect(settings).toEqual({ + portReadyTimeoutMs: 1000, + activeWorkMaxMs: 2000, + keepaliveRenewIntervalMs: 3000, + recoveryMaxAttempts: 4, + }); + }); + + it('falls back per-setting when an override is invalid', () => { + const settings = resolveRuntimeSettings( + envWith({ + CF_CONTAINER_RECOVERY_MAX_ATTEMPTS: '0', + CF_CONTAINER_ACTIVE_WORK_MAX_MS: 'not-a-number', + }), + DEFAULTS + ); + expect(settings.recoveryMaxAttempts).toBe(DEFAULTS.recoveryMaxAttempts); + expect(settings.activeWorkMaxMs).toBe(DEFAULTS.activeWorkMaxMs); + }); +}); 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 index 2c3028a55..b44b7a8be 100644 --- 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 @@ -1,141 +1,165 @@ import { describe, expect, it, vi } from 'vitest'; import { VmAgentContainer } from '../../../src/durable-objects/vm-agent-container'; +import type { RuntimeRecoveryState } from '../../../src/durable-objects/vm-agent-container-recovery'; -// 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. +type PrivateContainer = { + prepareForRequest: (this: unknown) => Promise; + ensureAwake: (this: unknown) => Promise; + resultResponse: (this: unknown, result: unknown) => Response; + interruptedRequestResponse: (this: unknown, request: Request) => Response; +}; -interface FakeState { - status: string; +const privateContainer = VmAgentContainer.prototype as unknown as PrivateContainer; + +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); } -function makeFake(opts: { - statuses: string[]; // sequence returned by getState() - lifecycleStatus: string; - wakeOk: boolean; +function makeProxyFake(input: { + ready: { ok: boolean; status: string; code?: string; message?: string }; + state: string; + response?: Response; }) { - 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, + const containerFetch = vi.fn().mockResolvedValue(input.response ?? new Response('proxied')); + const beginUnexpectedRecovery = vi.fn().mockResolvedValue(null); + return { + fake: { + defaultPort: 8080, + prepareForRequest: vi.fn().mockResolvedValue(input.ready), + resultResponse: privateContainer.resultResponse, + interruptedRequestResponse: privateContainer.interruptedRequestResponse, + getState: vi.fn().mockResolvedValue({ status: input.state }), + containerFetch, + beginUnexpectedRecovery, + }, containerFetch, - wakeFromSnapshot, - defaultPort: 8080, - wakeChain: Promise.resolve(), - ensureAwake: (VmAgentContainer.prototype as unknown as { ensureAwake: unknown }) - .ensureAwake, - ctx: { storage: { get: vi.fn().mockResolvedValue(opts.lifecycleStatus) } }, + beginUnexpectedRecovery, }; - 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, +describe('VmAgentContainer proxy recovery boundaries', () => { + it('forwards a prompt only after wake/restore reports running', async () => { + const { fake, containerFetch } = makeProxyFake({ + ready: { ok: true, status: 'running' }, + state: 'running', }); - const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + const response = 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); + expect(response.status).toBe(200); }); - it('returns 503 (not 410/proxy) when wake fails', async () => { - const { fake, containerFetch } = makeFake({ - statuses: ['stopped', 'stopped'], - lifecycleStatus: 'sleeping', - wakeOk: false, + it('returns a sanitized degraded response without forwarding', async () => { + const { fake, containerFetch } = makeProxyFake({ + ready: { + ok: false, + status: 'degraded', + code: 'RUNTIME_RECOVERY_DEGRADED', + message: 'The Instant session could not restore its last safe checkpoint.', + }, + state: 'stopped', }); - const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + const response = await callProxyHttp( + fake, + new Request('http://container/prompt', { method: 'POST' }) + ); - expect(res.status).toBe(503); + expect(response.status).toBe(409); + await expect(response.json()).resolves.toMatchObject({ + error: 'RUNTIME_RECOVERY_DEGRADED', + }); 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, + it('keeps an explicit stop terminal', async () => { + const { fake, containerFetch } = makeProxyFake({ + ready: { + ok: false, + status: 'stopped', + code: 'RUNTIME_STOPPED', + message: 'This Instant session was stopped and cannot be resumed.', + }, + state: 'stopped', }); - const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' })); + const response = await callProxyHttp( + fake, + new Request('http://container/prompt', { method: 'POST' }) + ); - expect(wakeFromSnapshot).not.toHaveBeenCalled(); + expect(response.status).toBe(410); 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 })); +describe('VmAgentContainer cold-wake serialization', () => { + it('launches and restores exactly once for two concurrent follow-ups', async () => { + const shared = { + lifecycle: 'sleeping', + recovery: { + version: 1, + phase: 'pending', + trigger: 'idle', + cause: { kind: 'idle_sleep' }, + attempts: 0, + promptDisposition: 'none', + agentSessionId: 'session-1', + startedAt: Date.now(), + updatedAt: Date.now(), + } satisfies RuntimeRecoveryState, + }; + const storage = { + get: vi.fn(async (key: string) => { + if (key === 'lifecycleStatus') return shared.lifecycle; + if (key === 'runtimeRecovery') return shared.recovery; + return undefined; + }), + put: vi.fn(async (key: string, value: unknown) => { + if (key === 'lifecycleStatus') shared.lifecycle = String(value); + if (key === 'runtimeRecovery') shared.recovery = value as RuntimeRecoveryState; + }), + }; 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)); + await new Promise((resolve) => setTimeout(resolve, 10)); shared.lifecycle = 'running'; - return { ok: true }; + return { ok: true, status: 'running' }; }); - + const containerFetch = vi.fn().mockResolvedValue(new Response('proxied')); 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) } }, + ctx: { storage }, + prepareForRequest: privateContainer.prepareForRequest, + ensureAwake: privateContainer.ensureAwake, + resultResponse: privateContainer.resultResponse, + interruptedRequestResponse: privateContainer.interruptedRequestResponse, + getRuntimeSettings: () => ({ recoveryMaxAttempts: 2 }), + wakeFromSnapshot, + beginUnexpectedRecovery: vi.fn(), + getState: vi.fn(async () => ({ + status: shared.lifecycle === 'running' ? 'running' : 'stopped', + })), + containerFetch, }; - const [a, b] = await Promise.all([ + const [first, second] = 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); + expect(first.status).toBe(200); + expect(second.status).toBe(200); }); }); diff --git a/apps/api/tests/unit/routes/agent-activity-callback.test.ts b/apps/api/tests/unit/routes/agent-activity-callback.test.ts index 225478283..cbd2ee0b4 100644 --- a/apps/api/tests/unit/routes/agent-activity-callback.test.ts +++ b/apps/api/tests/unit/routes/agent-activity-callback.test.ts @@ -5,6 +5,11 @@ import { AppError } from '../../../src/middleware/error'; const mocks = vi.hoisted(() => ({ updateSets: [] as Array>, + workspace: null as Record | null, + // Row returned to the S2 staleness guard's combined agent_sessions⋈nodes read + // ({ updatedAt, runtime }). Default is non-Instant so existing error tests + // still process (the guard only engages for cf-container runtimes). + guardRow: { updatedAt: null as string | null, runtime: 'vm' as string | null }, jwt: { verifyCallbackToken: vi.fn(), }, @@ -40,13 +45,19 @@ vi.mock('drizzle-orm/d1', () => ({ return { where: vi.fn().mockResolvedValue(undefined) }; }, }), - select: () => ({ - from: () => ({ - leftJoin: () => ({ - where: () => ({ get: vi.fn().mockResolvedValue(null) }), - }), - }), - }), + // Supports both the idle-branch workspace⋈nodes read (selection includes + // `id`) and the S2 guard's agent_sessions⋈workspaces⋈nodes read (selection + // includes `updatedAt`). Any number of leftJoins chain into the same `where`. + select: (selection?: Record) => { + const rowFor = () => + selection && 'updatedAt' in selection ? mocks.guardRow : mocks.workspace; + const terminal = { get: () => Promise.resolve(rowFor()) }; + const joinable: { leftJoin: () => typeof joinable; where: () => typeof terminal } = { + leftJoin: () => joinable, + where: () => terminal, + }; + return { from: () => joinable }; + }, }), })); @@ -97,13 +108,27 @@ function assignedAcpSession(overrides: Record = {}) { }; } -async function postActivity(app: Hono, body: Record): Promise { +/** + * Build an (unsigned) JWT whose payload carries `iat`. `verifyCallbackToken` is + * mocked, so the signature is irrelevant — the S2 guard only needs `decodeJwt` + * to read `iat` from an already-verified token. + */ +function tokenWithIatSeconds(iatSeconds: number): string { + const seg = (obj: unknown) => Buffer.from(JSON.stringify(obj)).toString('base64url'); + return `${seg({ alg: 'RS256', typ: 'JWT' })}.${seg({ iat: iatSeconds, workspace: 'workspace-1' })}.sig`; +} + +async function postActivity( + app: Hono, + body: Record, + token = 'callback-token' +): Promise { return app.request( '/api/projects/project-1/acp-sessions/agent-session-1/activity', { method: 'POST', headers: { - Authorization: 'Bearer callback-token', + Authorization: `Bearer ${token}`, 'Content-Type': 'application/json', }, body: JSON.stringify(body), @@ -116,6 +141,8 @@ describe('agent activity callback', () => { beforeEach(() => { vi.clearAllMocks(); mocks.updateSets.length = 0; + mocks.workspace = null; + mocks.guardRow = { updatedAt: null, runtime: 'vm' }; mocks.jwt.verifyCallbackToken.mockResolvedValue({ workspace: 'workspace-1', type: 'callback', @@ -128,6 +155,45 @@ describe('agent activity callback', () => { mocks.container.markVmAgentContainerActiveWorkEndedBestEffort.mockResolvedValue(undefined); }); + it('snapshots an idle Instant session with the exact agent type before handback', async () => { + mocks.projectData.getAcpSession.mockResolvedValueOnce( + assignedAcpSession({ status: 'active', acpSdkSessionId: 'sdk-session-1' }) + ); + mocks.workspace = { + id: 'workspace-1', + userId: 'user-1', + chatSessionId: 'chat-session-1', + runtime: 'cf-container', + }; + mocks.nodeAgent.hibernateAgentSessionOnNode.mockResolvedValueOnce({ status: 'available' }); + const app = await createTestApp(); + + const response = await postActivity(app, { + activity: 'idle', + nodeId: 'node-1', + agentType: 'openai-codex', + }); + + expect(response.status).toBe(204); + expect(mocks.nodeAgent.hibernateAgentSessionOnNode).toHaveBeenCalledWith( + 'node-1', + 'workspace-1', + 'sdk-session-1', + env, + 'user-1', + { + chatSessionId: 'chat-session-1', + runtime: 'cf-container', + agentType: 'openai-codex', + } + ); + expect(mocks.container.markVmAgentContainerActiveWorkEndedBestEffort).toHaveBeenCalledWith( + env, + 'node-1', + 'agent_activity_idle' + ); + }); + it('turns VM-agent error activity into durable failed control-plane state', async () => { const app = await createTestApp(); @@ -258,4 +324,144 @@ describe('agent activity callback', () => { 'agent_activity_error' ); }); + + // --- S2: stale superseded-generation callback guard (Instant recovery race) --- + describe('stale Instant callback guard', () => { + const TOKEN_IAT_SECONDS = 1_700_000_000; + const TOKEN_IAT_MS = TOKEN_IAT_SECONDS * 1000; + const OLD_TOKEN = tokenWithIatSeconds(TOKEN_IAT_SECONDS); + + beforeEach(() => { + // The session the DO recovered is healthy + running (canTransition→failed), + // so WITHOUT the guard a late error WOULD regress it. + mocks.projectData.getAcpSession.mockResolvedValue( + assignedAcpSession({ status: 'running', acpSdkSessionId: 'sdk-1' }) + ); + }); + + it('(a) rejects a stale error callback after recovery completed — session NOT regressed', async () => { + // Recovery completed: agent_sessions.updated_at reconciled to running well + // after the OLD container's token was issued (gap 180s ≫ 60s margin). + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 180_000).toISOString(), + }; + const app = await createTestApp(); + + const response = await postActivity( + app, + { + activity: 'error', + nodeId: 'node-1', + agentType: 'openai-codex', + statusError: 'peer disconnected before response', + }, + OLD_TOKEN + ); + + expect(response.status).toBe(204); + // Fully short-circuited: no mirror flip, no destructive transition, no + // active-work-ended on the live recovered generation. + expect(mocks.projectData.reportAcpSessionActivity).not.toHaveBeenCalled(); + expect(mocks.projectData.transitionAcpSession).not.toHaveBeenCalled(); + expect(mocks.projectData.failSession).not.toHaveBeenCalled(); + expect(mocks.updateSets).toHaveLength(0); + expect( + mocks.container.markVmAgentContainerActiveWorkEndedBestEffort + ).not.toHaveBeenCalled(); + expect(mocks.log.warn).toHaveBeenCalledWith( + 'acp_activity.rejected_stale_callback', + expect.objectContaining({ + projectId: 'project-1', + sessionId: 'agent-session-1', + nodeId: 'node-1', + runtime: 'cf-container', + action: 'rejected_stale_callback', + }) + ); + }); + + it('(b) still fails the session for a genuine crash of the CURRENT container (no recovery)', async () => { + // Same generation: the row was last reconciled ~at token issuance (gap + // 0.5s ≪ margin), so the error is legitimate and MUST fail the session. + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 500).toISOString(), + }; + const app = await createTestApp(); + + const response = await postActivity( + app, + { + activity: 'error', + nodeId: 'node-1', + agentType: 'openai-codex', + statusError: 'ACP NewSession failed: context deadline exceeded', + }, + OLD_TOKEN + ); + + expect(response.status).toBe(204); + expect(mocks.projectData.reportAcpSessionActivity).toHaveBeenCalled(); + expect(mocks.projectData.transitionAcpSession).toHaveBeenCalledWith( + env, + 'project-1', + 'agent-session-1', + 'failed', + expect.objectContaining({ actorType: 'vm-agent', actorId: 'node-1' }) + ); + expect(mocks.projectData.failSession).toHaveBeenCalled(); + expect(mocks.updateSets).toContainEqual(expect.objectContaining({ status: 'error' })); + expect(mocks.log.warn).not.toHaveBeenCalledWith( + 'acp_activity.rejected_stale_callback', + expect.anything() + ); + }); + + it('(c) rejects a stale error arriving DURING recovery (row reconciled to recovery, not yet running)', async () => { + // persistRuntimeRecovering has stamped the row (updated_at fresh, 300s + // after the old token) but recovery has not completed. The mid-recovery + // session must NOT be regressed by the superseded generation's callback; + // the current generation will report its own state after restore. + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 300_000).toISOString(), + }; + const app = await createTestApp(); + + const response = await postActivity( + app, + { activity: 'error', nodeId: 'node-1', statusError: 'container_stopped' }, + OLD_TOKEN + ); + + expect(response.status).toBe(204); + expect(mocks.projectData.transitionAcpSession).not.toHaveBeenCalled(); + expect(mocks.projectData.failSession).not.toHaveBeenCalled(); + expect(mocks.updateSets).toHaveLength(0); + expect(mocks.log.warn).toHaveBeenCalledWith( + 'acp_activity.rejected_stale_callback', + expect.objectContaining({ action: 'rejected_stale_callback' }) + ); + }); + + it('does not engage for VM-runtime nodes even when the row is newer than the token', async () => { + // Non-Instant runtime never has the generation-replacement race → process. + mocks.guardRow = { + runtime: 'vm', + updatedAt: new Date(TOKEN_IAT_MS + 999_000).toISOString(), + }; + const app = await createTestApp(); + + const response = await postActivity( + app, + { activity: 'error', nodeId: 'node-1', statusError: 'agent crashed' }, + OLD_TOKEN + ); + + expect(response.status).toBe(204); + expect(mocks.projectData.transitionAcpSession).toHaveBeenCalled(); + expect(mocks.projectData.failSession).toHaveBeenCalled(); + }); + }); }); diff --git a/apps/api/tests/unit/routes/agent-session-resume-recovery.test.ts b/apps/api/tests/unit/routes/agent-session-resume-recovery.test.ts new file mode 100644 index 000000000..5fd056dfa --- /dev/null +++ b/apps/api/tests/unit/routes/agent-session-resume-recovery.test.ts @@ -0,0 +1,269 @@ +import { Hono } from 'hono'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { Env } from '../../../src/env'; +import { AppError } from '../../../src/middleware/error'; + +// Must match RUNTIME_REQUEST_INTERRUPTED_MESSAGE in +// src/durable-objects/vm-agent-container-recovery.ts (inlined to keep this unit +// test free of the DO module's drizzle/project-data import side effects). +const RUNTIME_REQUEST_INTERRUPTED_MESSAGE = + 'Your message is saved, but delivery was interrupted and its execution outcome is unknown. It was not replayed automatically. After restore finishes, check the transcript and partial output before deciding whether to send it again.'; + +const mocks = vi.hoisted(() => ({ + getOwnedWorkspace: vi.fn(), + getOwnedNode: vi.fn(), + resumeVmAgentContainer: vi.fn(), + resumeAgentSessionOnNode: vi.fn(), + updateSet: vi.fn(), + session: { + id: 'agent-session-1', + workspaceId: 'workspace-1', + userId: 'user-1', + status: 'error', + label: 'Instant', + agentType: 'codex', + agentProfileId: null, + skillId: null, + worktreePath: null, + stoppedAt: null, + suspendedAt: null, + errorMessage: 'Runtime unavailable', + lastPrompt: null, + createdAt: '2026-07-21T00:00:00.000Z', + updatedAt: '2026-07-21T00:00:00.000Z', + }, +})); + +vi.mock('../../../src/auth', () => ({ + createAuth: () => ({ + api: { + getSession: vi.fn().mockResolvedValue({ + user: { + id: 'user-1', + email: 'user@example.com', + name: 'Test User', + role: 'user', + status: 'active', + }, + session: { id: 'auth-session-1', expiresAt: new Date('2030-01-01T00:00:00Z') }, + }), + }, + }), +})); + +vi.mock('../../../src/routes/workspaces/_helpers', () => ({ + assertNodeOperational: vi.fn(), + getOwnedWorkspace: mocks.getOwnedWorkspace, + getOwnedNode: mocks.getOwnedNode, +})); + +vi.mock('../../../src/services/node-agent', () => ({ + createAgentSessionOnNode: vi.fn(), + resumeAgentSessionOnNode: mocks.resumeAgentSessionOnNode, + stopAgentSessionOnNode: vi.fn(), + suspendAgentSessionOnNode: vi.fn(), +})); + +vi.mock('../../../src/services/vm-agent-container', () => ({ + resumeVmAgentContainer: mocks.resumeVmAgentContainer, +})); + +// The resume route reads the session via select().from().where().limit(1) both +// BEFORE recovery (initial snapshot) and AFTER a successful cf-container recovery +// (re-fetch). Returning a fresh COPY of mocks.session each call models D1: the +// initial read captures the pre-recovery row, and the re-fetch reflects whatever +// the DO's persistRuntimeRecovered wrote during resumeVmAgentContainer. +vi.mock('drizzle-orm/d1', () => ({ + drizzle: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + // Session reads use .limit(1); unrelated middleware lookups use .get(). + limit: () => Promise.resolve([{ ...mocks.session }]), + get: () => Promise.resolve(undefined), + }), + }), + }), + update: () => ({ + set: (values: unknown) => { + mocks.updateSet(values); + return { where: () => Promise.resolve() }; + }, + }), + }), +})); + +async function createTestApp() { + const { agentSessionRoutes } = await import('../../../src/routes/workspaces/agent-sessions'); + const app = new Hono<{ Bindings: Env }>(); + app.route('/api/workspaces', agentSessionRoutes); + app.onError((error, c) => { + if (error instanceof AppError) { + return c.json(error.toJSON(), error.statusCode as 409); + } + return c.json({ error: 'INTERNAL_ERROR', message: 'Internal server error' }, 500); + }); + return app; +} + +async function postResume(app: Hono<{ Bindings: Env }>): Promise { + return app.request( + '/api/workspaces/workspace-1/agent-sessions/agent-session-1/resume', + { method: 'POST' }, + { DATABASE: {} } as Env + ); +} + +describe('agent session Instant recovery route', () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.session.status = 'error'; + mocks.session.errorMessage = 'Runtime unavailable'; + mocks.session.suspendedAt = null; + mocks.session.stoppedAt = null; + mocks.getOwnedWorkspace.mockResolvedValue({ + id: 'workspace-1', + userId: 'user-1', + nodeId: 'node-1', + }); + mocks.getOwnedNode.mockResolvedValue({ + id: 'node-1', + userId: 'user-1', + runtime: 'cf-container', + status: 'error', + }); + mocks.resumeVmAgentContainer.mockResolvedValue({ + ok: true, + status: 'running', + degraded: false, + }); + }); + + it('marks the session running only after the container reports recovery', async () => { + let resolveRecovery!: (value: unknown) => void; + mocks.resumeVmAgentContainer.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveRecovery = resolve; + }) + ); + const app = await createTestApp(); + + const responsePromise = postResume(app); + await vi.waitFor(() => + expect(mocks.resumeVmAgentContainer).toHaveBeenCalledWith( + expect.objectContaining({ DATABASE: {} }), + 'node-1', + 'agent-session-1' + ) + ); + // Response is still pending until the DO reports recovery. + // The route never rewrites the row itself on the recovery path — the DO owns + // the reconciliation, so updateSet must not be called at all. + expect(mocks.updateSet).not.toHaveBeenCalled(); + + // Simulate the DO's persistRuntimeRecovered reconciling D1 to running. + mocks.session.status = 'running'; + mocks.session.errorMessage = null; + mocks.session.updatedAt = '2026-07-21T00:05:00.000Z'; + resolveRecovery({ ok: true, status: 'running', degraded: false }); + const response = await responsePromise; + + expect(response.status).toBe(200); + // No route-side rewrite — the DO reconciled D1; the route re-fetches + returns. + expect(mocks.updateSet).not.toHaveBeenCalled(); + expect(await response.json()).toEqual(expect.objectContaining({ status: 'running' })); + }); + + it('returns the DO-reconciled manual_retry state without clobbering error_message', async () => { + // A mutating prompt was interrupted; the DO recovers and persists the + // manual-retry notice into agent_sessions.error_message in the SAME request. + mocks.resumeVmAgentContainer.mockImplementationOnce(async () => { + mocks.session.status = 'running'; + mocks.session.errorMessage = RUNTIME_REQUEST_INTERRUPTED_MESSAGE; + mocks.session.updatedAt = '2026-07-21T00:05:00.000Z'; + return { ok: true, status: 'running', degraded: false }; + }); + const app = await createTestApp(); + + const response = await postResume(app); + + expect(response.status).toBe(200); + // The route must NOT run its own status rewrite (which nulls error_message). + expect(mocks.updateSet).not.toHaveBeenCalled(); + const body = (await response.json()) as { status: string; errorMessage: string | null }; + expect(body.status).toBe('running'); + expect(body.errorMessage).toBe(RUNTIME_REQUEST_INTERRUPTED_MESSAGE); + }); + + it('returns a stable degraded error without falsely marking the session running', async () => { + mocks.resumeVmAgentContainer.mockResolvedValueOnce({ + ok: false, + status: 'degraded', + degraded: true, + code: 'RUNTIME_RECOVERY_DEGRADED', + message: + 'The Instant session could not restore its last safe checkpoint. Your transcript and partial output are still available.', + }); + const app = await createTestApp(); + + const response = await postResume(app); + + const body = await response.json(); + expect(body).toEqual({ + error: 'RUNTIME_RECOVERY_DEGRADED', + message: + 'The Instant session could not restore its last safe checkpoint. Your transcript and partial output are still available.', + }); + expect(mocks.updateSet).not.toHaveBeenCalled(); + }); + + it('preserves an explicit stop as terminal', async () => { + mocks.resumeVmAgentContainer.mockResolvedValueOnce({ + ok: false, + status: 'stopped', + degraded: false, + code: 'RUNTIME_STOPPED', + message: 'This Instant session was stopped and cannot be resumed.', + }); + const app = await createTestApp(); + + const response = await postResume(app); + + expect(response.status).toBe(410); + expect(mocks.updateSet).not.toHaveBeenCalled(); + }); + + it('never triggers container recovery for a VM-runtime node (T13)', async () => { + // VM-runtime suspended session: recovery is a cf-container-only concern; the + // route must use the suspended VM resume path, not resumeVmAgentContainer. + mocks.getOwnedNode.mockResolvedValue({ + id: 'node-1', + userId: 'user-1', + runtime: 'vm', + status: 'running', + }); + mocks.session.status = 'suspended'; + mocks.session.suspendedAt = '2026-07-21T00:01:00.000Z'; + mocks.session.errorMessage = null; + const app = await createTestApp(); + + const response = await postResume(app); + + expect(response.status).toBe(200); + expect(mocks.resumeVmAgentContainer).not.toHaveBeenCalled(); + expect(mocks.resumeAgentSessionOnNode).toHaveBeenCalledWith( + 'node-1', + 'workspace-1', + 'agent-session-1', + expect.anything(), + 'user-1' + ); + // VM path DOES rewrite the row to running (legacy lifecycle). + expect(mocks.updateSet).toHaveBeenCalledWith( + expect.objectContaining({ status: 'running', errorMessage: null }) + ); + expect(await response.json()).toEqual(expect.objectContaining({ status: 'running' })); + }); +}); 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 b8e88bece..18227f513 100644 --- a/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts +++ b/apps/api/tests/unit/routes/chat-prompt-cancel.test.ts @@ -2,6 +2,7 @@ import { Hono } from 'hono'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { Env } from '../../../src/env'; +import { AppError } from '../../../src/middleware/error'; import { chatRoutes } from '../../../src/routes/chat'; import * as projectDataService from '../../../src/services/project-data'; @@ -94,6 +95,7 @@ function setupDrizzle(opts: { id: string; nodeId: string | null; nodeStatus: string | null; + nodeRuntime?: string | null; userId?: string; projectId?: string | null; } | null; @@ -102,7 +104,7 @@ function setupDrizzle(opts: { // Default the workspace owner to the authenticated caller so the // defence-in-depth ownership assertion passes for positive cases. const workspaceRow = opts.workspace - ? { userId: 'user-1', projectId: 'proj-1', ...opts.workspace } + ? { userId: 'user-1', projectId: 'proj-1', nodeRuntime: 'vm', ...opts.workspace } : null; let callCount = 0; const selectMock = vi.fn().mockImplementation(() => { @@ -112,15 +114,17 @@ function setupDrizzle(opts: { from: vi.fn().mockReturnValue({ leftJoin: vi.fn().mockReturnValue({ where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue( - currentCall === 1 ? (workspaceRow ? [workspaceRow] : []) : [] - ), + limit: vi + .fn() + .mockResolvedValue(currentCall === 1 ? (workspaceRow ? [workspaceRow] : []) : []), }), }), where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue( - currentCall === 2 ? (opts.agentSession ? [opts.agentSession] : []) : [] - ), + limit: vi + .fn() + .mockResolvedValue( + currentCall === 2 ? (opts.agentSession ? [opts.agentSession] : []) : [] + ), }), }), }; @@ -158,7 +162,13 @@ describe('GET /sessions', () => { select: vi.fn().mockReturnValue({ from: vi.fn().mockReturnValue({ where: vi.fn().mockResolvedValue([ - { id: 'user-1', name: 'Alice', email: 'alice@example.com', image: null, avatarUrl: null }, + { + id: 'user-1', + name: 'Alice', + email: 'alice@example.com', + image: null, + avatarUrl: null, + }, { id: 'user-2', name: 'Bob', email: 'bob@example.com', image: null, avatarUrl: null }, ]), }), @@ -176,15 +186,19 @@ describe('GET /sessions', () => { total: 2, }); - const response = await app.request( - '/api/projects/proj-1/sessions', - { method: 'GET' }, - { DATABASE: {} as D1Database } as Env, - ); + const response = await app.request('/api/projects/proj-1/sessions', { method: 'GET' }, { + DATABASE: {} as D1Database, + } as Env); expect(response.status).toBe(200); expect(projectDataService.listSessions).toHaveBeenCalledWith( - expect.anything(), 'proj-1', null, 20, 0, null, null, + expect.anything(), + 'proj-1', + null, + 20, + 0, + null, + null ); const body = await response.json(); expect(body.sessions).toMatchObject([ @@ -203,12 +217,18 @@ describe('GET /sessions', () => { const response = await app.request( '/api/projects/proj-1/sessions?scope=my', { method: 'GET' }, - { DATABASE: {} as D1Database } as Env, + { DATABASE: {} as D1Database } as Env ); expect(response.status).toBe(200); expect(projectDataService.listSessions).toHaveBeenCalledWith( - expect.anything(), 'proj-1', null, 20, 0, null, 'user-1', + expect.anything(), + 'proj-1', + null, + 20, + 0, + null, + 'user-1' ); }); }); @@ -222,7 +242,7 @@ describe('POST /sessions/:sessionId/prompt', () => { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ content: 'hello agent' }), }, - { DATABASE: {} as D1Database } as Env, + { DATABASE: {} as D1Database } as Env ); } @@ -237,14 +257,25 @@ 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', undefined, undefined, + 'node-1', + 'ws-1', + 'agent-sess-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' }, + workspace: { + id: 'ws-1', + nodeId: 'node-1', + nodeStatus: 'sleeping', + nodeRuntime: 'cf-container', + }, agentSession: { id: 'agent-sess-1' }, }); mocks.sendPromptToAgentOnNode.mockResolvedValue({ ok: true }); @@ -253,12 +284,74 @@ 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', undefined, - { requestTimeoutMs: 120_000 }, + 'node-1', + 'ws-1', + 'agent-sess-1', + 'hello agent', + expect.anything(), + 'user-1', + undefined, + { requestTimeoutMs: 120_000 } ); }); + it.each(['recovery', 'error'])( + 'routes cf-container %s state to bounded recovery', + async (nodeStatus) => { + setupDrizzle({ + workspace: { + id: 'ws-1', + nodeId: 'node-1', + nodeStatus, + nodeRuntime: 'cf-container', + }, + 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 } + ); + } + ); + + it('returns the stable sanitized interruption error without replaying', async () => { + setupDrizzle({ + workspace: { + id: 'ws-1', + nodeId: 'node-1', + nodeStatus: 'running', + nodeRuntime: 'cf-container', + }, + agentSession: { id: 'agent-sess-1' }, + }); + mocks.sendPromptToAgentOnNode.mockRejectedValue( + new AppError( + 409, + 'RUNTIME_REQUEST_INTERRUPTED', + 'Your message is saved, but it was not replayed automatically.' + ) + ); + + const response = await postPrompt(); + + expect(response.status).toBe(409); + await expect(response.json()).resolves.toEqual({ + error: 'RUNTIME_REQUEST_INTERRUPTED', + message: 'Your message is saved, but it was not replayed automatically.', + }); + }); + it('returns 404 when no active workspace is found', async () => { setupDrizzle({ workspace: null, agentSession: null }); @@ -385,11 +478,9 @@ describe('POST /sessions/:sessionId/prompt', () => { describe('POST /sessions/:sessionId/cancel', () => { function postCancel() { - return app.request( - '/api/projects/proj-1/sessions/chat-1/cancel', - { method: 'POST' }, - { DATABASE: {} as D1Database } as Env, - ); + return app.request('/api/projects/proj-1/sessions/chat-1/cancel', { method: 'POST' }, { + DATABASE: {} as D1Database, + } as Env); } it('cancels a running prompt successfully', async () => { @@ -406,8 +497,11 @@ describe('POST /sessions/:sessionId/cancel', () => { expect(body.status).toBe('cancelled'); expect(body.message).toBe('Prompt cancel signal sent'); expect(mocks.cancelAgentSessionOnNode).toHaveBeenCalledWith( - 'node-1', 'ws-1', 'agent-sess-1', - expect.anything(), 'user-1', + 'node-1', + 'ws-1', + 'agent-sess-1', + expect.anything(), + 'user-1' ); }); diff --git a/apps/api/tests/unit/routes/deployment-release-compose-submission.test.ts b/apps/api/tests/unit/routes/deployment-release-compose-submission.test.ts index 524e8d439..04f9eeaa8 100644 --- a/apps/api/tests/unit/routes/deployment-release-compose-submission.test.ts +++ b/apps/api/tests/unit/routes/deployment-release-compose-submission.test.ts @@ -75,16 +75,22 @@ vi.mock('../../../src/middleware/project-auth', () => ({ requireProjectCapability: vi.fn().mockResolvedValue(undefined), })); -vi.mock('../../../src/middleware/error', () => ({ - errors: { - badRequest: (msg: string) => - Object.assign(new Error(msg), { statusCode: 400, error: 'BAD_REQUEST', message: msg }), - notFound: (msg: string) => - Object.assign(new Error(msg), { statusCode: 404, error: 'NOT_FOUND', message: msg }), - conflict: (msg: string) => - Object.assign(new Error(msg), { statusCode: 409, error: 'CONFLICT', message: msg }), - }, -})); +vi.mock('../../../src/middleware/error', async () => { + const actual = await vi.importActual( + '../../../src/middleware/error' + ); + return { + ...actual, + errors: { + badRequest: (msg: string) => + Object.assign(new Error(msg), { statusCode: 400, error: 'BAD_REQUEST', message: msg }), + notFound: (msg: string) => + Object.assign(new Error(msg), { statusCode: 404, error: 'NOT_FOUND', message: msg }), + conflict: (msg: string) => + Object.assign(new Error(msg), { statusCode: 409, error: 'CONFLICT', message: msg }), + }, + }; +}); vi.mock('drizzle-orm/d1', () => ({ drizzle: () => createMockDb(), diff --git a/apps/api/tests/unit/routes/stale-callback-guard.test.ts b/apps/api/tests/unit/routes/stale-callback-guard.test.ts new file mode 100644 index 000000000..df34ceca1 --- /dev/null +++ b/apps/api/tests/unit/routes/stale-callback-guard.test.ts @@ -0,0 +1,131 @@ +import { describe, expect, it } from 'vitest'; + +import type { Env } from '../../../src/env'; +import { + callbackTokenIssuedAtMs, + DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS, + getInstantStaleCallbackMarginMs, + isSupersededInstantCallback, +} from '../../../src/routes/_stale-callback-guard'; + +function jwtWith(payload: Record): string { + const seg = (obj: unknown) => Buffer.from(JSON.stringify(obj)).toString('base64url'); + return `${seg({ alg: 'RS256', typ: 'JWT' })}.${seg(payload)}.sig`; +} + +const IAT_SECONDS = 1_700_000_000; +const IAT_MS = IAT_SECONDS * 1000; +const MARGIN = DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS; +const iso = (ms: number) => new Date(ms).toISOString(); + +describe('callbackTokenIssuedAtMs', () => { + it('returns iat in ms for a decodable token', () => { + expect(callbackTokenIssuedAtMs(jwtWith({ iat: IAT_SECONDS }))).toBe(IAT_MS); + }); + + it('returns null for a non-JWT string', () => { + expect(callbackTokenIssuedAtMs('not-a-jwt')).toBeNull(); + }); + + it('returns null when iat is absent or non-numeric', () => { + expect(callbackTokenIssuedAtMs(jwtWith({ workspace: 'ws-1' }))).toBeNull(); + expect(callbackTokenIssuedAtMs(jwtWith({ iat: 'nope' }))).toBeNull(); + }); +}); + +describe('getInstantStaleCallbackMarginMs', () => { + it('defaults when unset', () => { + expect(getInstantStaleCallbackMarginMs({} as Env)).toBe(DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS); + }); + + it('honours a valid override', () => { + expect( + getInstantStaleCallbackMarginMs({ INSTANT_STALE_CALLBACK_MARGIN_MS: '30000' } as unknown as Env) + ).toBe(30_000); + }); + + it('falls back to default for invalid or negative values', () => { + expect( + getInstantStaleCallbackMarginMs({ INSTANT_STALE_CALLBACK_MARGIN_MS: 'abc' } as unknown as Env) + ).toBe(DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS); + expect( + getInstantStaleCallbackMarginMs({ INSTANT_STALE_CALLBACK_MARGIN_MS: '-5' } as unknown as Env) + ).toBe(DEFAULT_INSTANT_STALE_CALLBACK_MARGIN_MS); + }); +}); + +describe('isSupersededInstantCallback', () => { + it('is stale when a cf-container row is reconciled beyond the margin after the token', () => { + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: iso(IAT_MS + MARGIN + 1_000), + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(true); + }); + + it('is NOT stale for a non-cf-container runtime even when the row is far newer', () => { + expect( + isSupersededInstantCallback({ + runtime: 'vm', + rowUpdatedAt: iso(IAT_MS + MARGIN + 999_000), + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(false); + }); + + it('is NOT stale within the margin (same generation reconcile gap)', () => { + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: iso(IAT_MS + 500), + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(false); + }); + + it('is NOT stale exactly at the boundary (strictly-greater comparison)', () => { + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: iso(IAT_MS + MARGIN), + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(false); + }); + + it('fails open (false) when the token iat is unknown', () => { + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: iso(IAT_MS + MARGIN + 999_000), + tokenIssuedAtMs: null, + marginMs: MARGIN, + }) + ).toBe(false); + }); + + it('fails open (false) when the row timestamp is missing or unparseable', () => { + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: null, + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(false); + expect( + isSupersededInstantCallback({ + runtime: 'cf-container', + rowUpdatedAt: 'not-a-date', + tokenIssuedAtMs: IAT_MS, + marginMs: MARGIN, + }) + ).toBe(false); + }); +}); diff --git a/apps/api/tests/unit/routes/task-callback-recoverable-error.test.ts b/apps/api/tests/unit/routes/task-callback-recoverable-error.test.ts index aec1b45fd..56a5c1b3f 100644 --- a/apps/api/tests/unit/routes/task-callback-recoverable-error.test.ts +++ b/apps/api/tests/unit/routes/task-callback-recoverable-error.test.ts @@ -38,18 +38,26 @@ const MOCK_PROJECT_DATA_BINDING = { vi.mock('drizzle-orm/d1', () => ({ drizzle: () => ({ - select: (selection?: Record) => ({ - from: () => ({ - where: () => ({ - limit: () => { - if (selection && 'chatSessionId' in selection) { - return Promise.resolve([{ chatSessionId: 'chat-recoverable' }]); - } - return Promise.resolve([{ ...mocks.task }]); - }, - }), - }), - }), + select: (selection?: Record) => { + const rows = () => { + // S2 staleness-guard read (agent_sessions⋈nodes) — non-Instant runtime so + // the existing `failed` callback keeps processing (guard does not engage). + if (selection && 'updatedAt' in selection) { + return [{ updatedAt: null, runtime: 'vm' }]; + } + if (selection && 'chatSessionId' in selection) { + return [{ chatSessionId: 'chat-recoverable' }]; + } + return [{ ...mocks.task }]; + }; + const afterOrderBy = { limit: () => Promise.resolve(rows()) }; + const terminal = { limit: () => Promise.resolve(rows()), orderBy: () => afterOrderBy }; + const joinable: { leftJoin: () => typeof joinable; where: () => typeof terminal } = { + leftJoin: () => joinable, + where: () => terminal, + }; + return { from: () => joinable }; + }, update: () => ({ set: (values: Record) => { mocks.updateSets.push(values); diff --git a/apps/api/tests/unit/routes/task-callback-stale-guard.test.ts b/apps/api/tests/unit/routes/task-callback-stale-guard.test.ts new file mode 100644 index 000000000..473d77368 --- /dev/null +++ b/apps/api/tests/unit/routes/task-callback-stale-guard.test.ts @@ -0,0 +1,256 @@ +import { Hono } from 'hono'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { AppError } from '../../../src/middleware/error'; + +// S2: a superseded (dead) Instant container can POST a stale `toStatus:'failed'` +// callback with a still-valid token AFTER the DO recovered a NEW generation to +// running. The guard must reject that (protect the healthy task) while still +// failing the task for a genuine crash of the CURRENT container. + +const mocks = vi.hoisted(() => { + const task = { + id: 'task-stale', + projectId: 'proj-stale', + userId: 'user-1', + parentTaskId: null as string | null, + workspaceId: 'ws-stale', + status: 'in_progress', + title: 'Instant task', + description: null as string | null, + priority: 0, + taskMode: 'task', + dispatchDepth: 0, + executionStep: 'running', + errorMessage: null as string | null, + outputSummary: null as string | null, + outputBranch: null as string | null, + outputPrUrl: null as string | null, + startedAt: null as string | null, + completedAt: null as string | null, + finalizedAt: null as string | null, + createdAt: '2026-07-21T00:00:00.000Z', + updatedAt: '2026-07-21T00:00:00.000Z', + }; + return { + task, + // Row returned to the guard's agent_sessions⋈workspaces⋈nodes read. + guardRow: { updatedAt: null as string | null, runtime: 'cf-container' as string | null }, + updateSets: [] as Array>, + setTaskStatus: vi.fn(), + computeBlockedForTask: vi.fn(), + cleanupTerminalTaskResourcesOrThrow: vi.fn(), + log: { debug: vi.fn(), error: vi.fn(), info: vi.fn(), warn: vi.fn() }, + waitUntilPromises: [] as Promise[], + }; +}); + +vi.mock('drizzle-orm/d1', () => ({ + drizzle: () => ({ + select: (selection?: Record) => { + const rows = () => { + if (selection && 'updatedAt' in selection) return [{ ...mocks.guardRow }]; + if (selection && 'chatSessionId' in selection) return [{ chatSessionId: 'chat-stale' }]; + return [{ ...mocks.task }]; + }; + const afterOrderBy = { limit: () => Promise.resolve(rows()) }; + const terminal = { + limit: () => Promise.resolve(rows()), + orderBy: () => afterOrderBy, + }; + const joinable: { leftJoin: () => typeof joinable; where: () => typeof terminal } = { + leftJoin: () => joinable, + where: () => terminal, + }; + return { from: () => joinable }; + }, + update: () => ({ + set: (values: Record) => { + mocks.updateSets.push(values); + Object.assign(mocks.task, values); + return { where: () => Promise.resolve() }; + }, + }), + }), +})); + +vi.mock('../../../src/lib/logger', () => ({ + log: mocks.log, + createModuleLogger: () => mocks.log, +})); + +vi.mock('../../../src/services/jwt', () => ({ + verifyCallbackToken: vi + .fn() + .mockResolvedValue({ workspace: 'ws-stale', type: 'callback', scope: 'workspace' }), +})); + +vi.mock('../../../src/services/project-data', () => ({ + recordActivityEvent: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('../../../src/services/task-terminal-cleanup', () => ({ + cleanupTerminalTaskResourcesOrThrow: mocks.cleanupTerminalTaskResourcesOrThrow, +})); + +vi.mock('../../../src/services/notification', () => ({ + getProjectName: vi.fn().mockResolvedValue('Stale Project'), + notifyTaskComplete: vi.fn().mockResolvedValue(undefined), + notifyTaskFailed: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('../../../src/services/task-status', () => ({ + canTransitionTaskStatus: vi.fn().mockReturnValue(true), + getAllowedTaskTransitions: vi.fn().mockReturnValue(['completed', 'failed', 'cancelled']), + isTaskStatus: vi.fn().mockReturnValue(true), +})); + +vi.mock('../../../src/routes/tasks/_helpers', () => ({ + computeBlockedForTask: mocks.computeBlockedForTask.mockResolvedValue(false), + setTaskStatus: mocks.setTaskStatus.mockImplementation( + async (_db, task, toStatus, _source, _workspace, options) => ({ + ...task, + status: toStatus, + errorMessage: options?.errorMessage ?? null, + }) + ), +})); + +const TOKEN_IAT_SECONDS = 1_700_000_000; +const TOKEN_IAT_MS = TOKEN_IAT_SECONDS * 1000; + +function tokenWithIatSeconds(iatSeconds: number): string { + const seg = (obj: unknown) => Buffer.from(JSON.stringify(obj)).toString('base64url'); + return `${seg({ alg: 'RS256', typ: 'JWT' })}.${seg({ iat: iatSeconds, workspace: 'ws-stale' })}.sig`; +} + +const OLD_TOKEN = tokenWithIatSeconds(TOKEN_IAT_SECONDS); + +async function createTestApp(): Promise { + const { taskCallbackRoute } = await import('../../../src/routes/tasks/callback'); + const app = new Hono(); + app.route('/api/projects', taskCallbackRoute); + app.onError((err, c) => { + if (err instanceof AppError) { + return c.json(err.toJSON(), err.statusCode as 400 | 401 | 403 | 404 | 409 | 500); + } + return c.json( + { error: 'INTERNAL_ERROR', message: err instanceof Error ? err.message : String(err) }, + 500 + ); + }); + return app; +} + +async function postFailed(app: Hono, token = OLD_TOKEN): Promise { + return app.request( + '/api/projects/proj-stale/tasks/task-stale/status/callback', + { + method: 'POST', + headers: { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json' }, + body: JSON.stringify({ toStatus: 'failed', reason: 'agent error', errorMessage: 'fatal' }), + }, + { DATABASE: {}, PROJECT_DATA: { idFromName: (id: string) => id, get: vi.fn() } }, + { waitUntil: (p: Promise) => mocks.waitUntilPromises.push(p) } + ); +} + +describe('task callback stale Instant guard', () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.task.status = 'in_progress'; + mocks.task.errorMessage = null; + mocks.updateSets.length = 0; + mocks.waitUntilPromises.length = 0; + mocks.guardRow = { updatedAt: null, runtime: 'cf-container' }; + mocks.computeBlockedForTask.mockResolvedValue(false); + }); + + it('(a) rejects a stale failed callback after recovery completed — task NOT regressed', async () => { + // Recovery reconciled the workspace's agent session to running 180s after + // the OLD container token was issued (≫ 60s margin) ⇒ superseded generation. + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 180_000).toISOString(), + }; + const app = await createTestApp(); + + const res = await postFailed(app); + + expect(res.status).toBe(200); + expect(mocks.setTaskStatus).not.toHaveBeenCalled(); + expect(mocks.cleanupTerminalTaskResourcesOrThrow).not.toHaveBeenCalled(); + expect(mocks.task.status).toBe('in_progress'); // unchanged + const body = (await res.json()) as { status: string }; + expect(body.status).toBe('in_progress'); + expect(mocks.log.warn).toHaveBeenCalledWith( + 'task.rejected_stale_callback', + expect.objectContaining({ + projectId: 'proj-stale', + taskId: 'task-stale', + workspaceId: 'ws-stale', + runtime: 'cf-container', + toStatus: 'failed', + action: 'rejected_stale_callback', + }) + ); + }); + + it('(b) still fails the task for a genuine crash of the CURRENT container (no recovery)', async () => { + // Same generation: last reconciled ~at token issuance (gap 0.5s ≪ margin). + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 500).toISOString(), + }; + const app = await createTestApp(); + + const res = await postFailed(app); + + expect(res.status).toBe(200); + expect(mocks.setTaskStatus).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ id: 'task-stale' }), + 'failed', + 'workspace_callback', + 'ws-stale', + expect.objectContaining({ errorMessage: 'fatal' }) + ); + expect(mocks.cleanupTerminalTaskResourcesOrThrow).toHaveBeenCalled(); + expect(mocks.log.warn).not.toHaveBeenCalledWith( + 'task.rejected_stale_callback', + expect.anything() + ); + }); + + it('(c) rejects a stale failed callback arriving DURING recovery (row reconciled, not yet running)', async () => { + mocks.guardRow = { + runtime: 'cf-container', + updatedAt: new Date(TOKEN_IAT_MS + 300_000).toISOString(), + }; + const app = await createTestApp(); + + const res = await postFailed(app); + + expect(res.status).toBe(200); + expect(mocks.setTaskStatus).not.toHaveBeenCalled(); + expect(mocks.cleanupTerminalTaskResourcesOrThrow).not.toHaveBeenCalled(); + expect(mocks.log.warn).toHaveBeenCalledWith( + 'task.rejected_stale_callback', + expect.objectContaining({ action: 'rejected_stale_callback' }) + ); + }); + + it('does not engage for VM-runtime nodes even when the row is newer than the token', async () => { + mocks.guardRow = { + runtime: 'vm', + updatedAt: new Date(TOKEN_IAT_MS + 999_000).toISOString(), + }; + const app = await createTestApp(); + + const res = await postFailed(app); + + expect(res.status).toBe(200); + expect(mocks.setTaskStatus).toHaveBeenCalled(); + expect(mocks.cleanupTerminalTaskResourcesOrThrow).toHaveBeenCalled(); + }); +}); diff --git a/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts b/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts index 86ed800e2..5a5d2e686 100644 --- a/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts +++ b/apps/api/tests/unit/routes/workspaces-session-snapshots.test.ts @@ -228,6 +228,8 @@ describe('workspaces session snapshot callback routes', () => { chatSessionId: 'chat-1', workspaceId: 'WS_1', agentSessionId: 'agent-session-1', + acpSessionId: 'acp-session-1', + agentType: 'openai-codex', status: 'available', degradation: 'none', skipped: [], @@ -261,5 +263,33 @@ describe('workspaces session snapshot callback routes', () => { runtimeBindings ); expect(mismatch.status).toBe(400); + + const sessionMismatch = 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, agentSessionId: 'agent-session-other' }, + }), + }, + runtimeBindings + ); + expect(sessionMismatch.status).toBe(400); + + const incompleteHarness = 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, agentType: undefined }, + }), + }, + runtimeBindings + ); + expect(incompleteHarness.status).toBe(400); }); }); diff --git a/apps/api/tests/unit/runtime-heartbeat-policy.test.ts b/apps/api/tests/unit/runtime-heartbeat-policy.test.ts new file mode 100644 index 000000000..6fa301265 --- /dev/null +++ b/apps/api/tests/unit/runtime-heartbeat-policy.test.ts @@ -0,0 +1,98 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { shouldDeferRuntimeHeartbeatTimeout } from '../../src/durable-objects/project-data/runtime-heartbeat-policy'; +import type { Env } from '../../src/durable-objects/project-data/types'; +import type { VmAgentContainerLifecycleStatus } from '../../src/durable-objects/vm-agent-container-lifecycle'; + +function createEnv(input: { + runtime?: string; + workspaceStatus?: string; + lifecycleStatus?: VmAgentContainerLifecycleStatus | null; + binding?: boolean; + pendingProbe?: boolean; +}): Env { + const first = vi.fn().mockResolvedValue({ + workspace_status: input.workspaceStatus ?? 'sleeping', + node_runtime: input.runtime ?? 'cf-container', + }); + const inspectLifecycle = input.pendingProbe + ? vi.fn().mockReturnValue(new Promise(() => undefined)) + : vi.fn().mockResolvedValue({ + status: input.lifecycleStatus ?? 'sleeping', + recoveryPhase: null, + recoveryTrigger: null, + activeWorkStatus: null, + }); + return { + DATABASE: { + prepare: vi.fn().mockReturnValue({ + bind: vi.fn().mockReturnValue({ first }), + }), + } as unknown as D1Database, + VM_AGENT_CONTAINER: + input.binding === false + ? undefined + : ({ + idFromName: vi.fn().mockReturnValue('do-id'), + get: vi.fn().mockReturnValue({ inspectLifecycle }), + } as unknown as Env['VM_AGENT_CONTAINER']), + }; +} + +const candidate = { workspaceId: 'ws-1', nodeId: 'node-1' }; + +describe('ProjectData runtime heartbeat timeout policy', () => { + it.each(['sleeping', 'recovering', 'waking', 'restoring', 'running'] as const)( + 'defers timeout while the Instant lifecycle is %s', + async (lifecycleStatus) => { + await expect( + shouldDeferRuntimeHeartbeatTimeout(createEnv({ lifecycleStatus }), candidate) + ).resolves.toMatchObject({ defer: true, reason: `cf_container_${lifecycleStatus}` }); + } + ); + + it('defers conservatively when the container binding is unavailable', async () => { + await expect( + shouldDeferRuntimeHeartbeatTimeout(createEnv({ binding: false }), candidate) + ).resolves.toEqual({ + defer: true, + reason: 'cf_container_lifecycle_binding_unavailable', + }); + }); + + it('bounds a lifecycle probe that never settles', async () => { + await expect( + shouldDeferRuntimeHeartbeatTimeout( + { ...createEnv({ pendingProbe: true }), TASK_LIVENESS_PROBE_TIMEOUT_MS: '1' }, + candidate + ) + ).resolves.toEqual({ defer: true, reason: 'cf_container_lifecycle_timeout' }); + }); + + it.each(['stopping', 'stopped', 'expired', 'error'] as const)( + 'allows timeout after the Instant lifecycle is %s', + async (lifecycleStatus) => { + await expect( + shouldDeferRuntimeHeartbeatTimeout(createEnv({ lifecycleStatus }), candidate) + ).resolves.toMatchObject({ defer: false, reason: `cf_container_${lifecycleStatus}` }); + } + ); + + it('keeps VM/devcontainer heartbeat timeout behavior unchanged', async () => { + await expect( + shouldDeferRuntimeHeartbeatTimeout(createEnv({ runtime: 'vm' }), candidate) + ).resolves.toEqual({ defer: false, reason: 'non_container_runtime' }); + }); + + it.each(['stopping', 'stopped'])( + 'does not defer after an explicit %s workspace transition', + async (workspaceStatus) => { + await expect( + shouldDeferRuntimeHeartbeatTimeout( + createEnv({ workspaceStatus, lifecycleStatus: 'sleeping' }), + candidate + ) + ).resolves.toEqual({ defer: false, reason: `workspace_${workspaceStatus}` }); + } + ); +}); diff --git a/apps/api/tests/unit/services/node-agent-create-workspace-timeout.test.ts b/apps/api/tests/unit/services/node-agent-create-workspace-timeout.test.ts index b1a597819..0a4aa54aa 100644 --- a/apps/api/tests/unit/services/node-agent-create-workspace-timeout.test.ts +++ b/apps/api/tests/unit/services/node-agent-create-workspace-timeout.test.ts @@ -25,6 +25,7 @@ const mocks = vi.hoisted(() => ({ getVmAgentContainerConfig: vi.fn(), markVmAgentContainerActiveWorkEndedBestEffort: vi.fn(), markVmAgentContainerActiveWorkStarted: vi.fn(), + markVmAgentContainerRequestInterrupted: vi.fn(), }, drizzle: vi.fn(), })); @@ -37,6 +38,8 @@ vi.mock('drizzle-orm/d1', () => ({ drizzle: mocks.drizzle })); import { createWorkspaceOnNode, getCfContainerCreateWorkspaceTimeoutMs, + NodeAgentRequestError, + sendPromptToAgentOnNode, } from '../../../src/services/node-agent'; const cfContainerEnv = { @@ -84,6 +87,7 @@ describe('createWorkspaceOnNode cf-container timeout plumbing', () => { vmAgentPort: 8080, sleepAfter: '10m', }); + mocks.container.markVmAgentContainerRequestInterrupted.mockResolvedValue(null); mocks.drizzle.mockReturnValue({ select: () => ({ from: () => ({ @@ -103,9 +107,15 @@ describe('createWorkspaceOnNode cf-container timeout plumbing', () => { it('times out with the provided budget when the container request stalls', async () => { mocks.container.fetchVmAgentContainer.mockImplementation(() => pendingResponse(60_000)); - const createPromise = createWorkspaceOnNode('node-1', cfContainerEnv, 'user-1', workspacePayload, { - requestTimeoutMs: 50, - }); + const createPromise = createWorkspaceOnNode( + 'node-1', + cfContainerEnv, + 'user-1', + workspacePayload, + { + requestTimeoutMs: 50, + } + ); const rejection = expect(createPromise).rejects.toThrow('Request timed out after 50ms'); await vi.advanceTimersByTimeAsync(60); await rejection; @@ -114,9 +124,15 @@ describe('createWorkspaceOnNode cf-container timeout plumbing', () => { it('completes when the container responds within the provided budget', async () => { mocks.container.fetchVmAgentContainer.mockImplementation(() => pendingResponse(40_000)); - const createPromise = createWorkspaceOnNode('node-1', cfContainerEnv, 'user-1', workspacePayload, { - requestTimeoutMs: 120_000, - }); + const createPromise = createWorkspaceOnNode( + 'node-1', + cfContainerEnv, + 'user-1', + workspacePayload, + { + requestTimeoutMs: 120_000, + } + ); await vi.advanceTimersByTimeAsync(40_500); await expect(createPromise).resolves.toEqual({ workspaceId: 'ws-1' }); }); @@ -124,9 +140,80 @@ describe('createWorkspaceOnNode cf-container timeout plumbing', () => { it('keeps the interactive 30s default when no override is provided', async () => { mocks.container.fetchVmAgentContainer.mockImplementation(() => pendingResponse(60_000)); - const createPromise = createWorkspaceOnNode('node-1', cfContainerEnv, 'user-1', workspacePayload); + const createPromise = createWorkspaceOnNode( + 'node-1', + cfContainerEnv, + 'user-1', + workspacePayload + ); const rejection = expect(createPromise).rejects.toThrow('Request timed out after 30000ms'); await vi.advanceTimersByTimeAsync(30_100); await rejection; }); + it('classifies a timed-out prompt before returning and never replays it', async () => { + mocks.container.fetchVmAgentContainer.mockImplementation(() => pendingResponse(60_000)); + mocks.container.markVmAgentContainerRequestInterrupted.mockResolvedValue({ + ok: false, + status: 'recovering', + code: 'RUNTIME_REQUEST_INTERRUPTED', + message: 'internal transport detail: bearer should-not-leak', + }); + + const promptPromise = sendPromptToAgentOnNode( + 'node-1', + 'ws-1', + 'agent-1', + 'continue', + cfContainerEnv, + 'user-1', + 'message-1', + { requestTimeoutMs: 50 } + ).catch((error) => error); + await vi.advanceTimersByTimeAsync(60); + const error = await promptPromise; + + expect(error).toBeInstanceOf(NodeAgentRequestError); + expect(error).toMatchObject({ + statusCode: 409, + error: 'RUNTIME_REQUEST_INTERRUPTED', + message: expect.stringContaining('execution outcome is unknown'), + }); + expect(mocks.container.fetchVmAgentContainer).toHaveBeenCalledTimes(1); + expect(mocks.container.markVmAgentContainerRequestInterrupted).toHaveBeenCalledWith( + cfContainerEnv, + 'node-1', + { method: 'POST', errorName: 'request_timeout' } + ); + }); + + it('preserves a stable interruption response returned directly by the Durable Object', async () => { + mocks.container.fetchVmAgentContainer.mockResolvedValue( + Response.json( + { + error: 'RUNTIME_REQUEST_INTERRUPTED', + message: 'internal transport detail: bearer should-not-leak', + }, + { status: 500 } + ) + ); + + const error = await sendPromptToAgentOnNode( + 'node-1', + 'ws-1', + 'agent-1', + 'continue', + cfContainerEnv, + 'user-1' + ).catch((caught) => caught); + + expect(error).toBeInstanceOf(NodeAgentRequestError); + expect(error).toMatchObject({ + statusCode: 409, + error: 'RUNTIME_REQUEST_INTERRUPTED', + message: expect.stringContaining('execution outcome is unknown'), + }); + expect(error.message).not.toContain('bearer should-not-leak'); + expect(mocks.container.fetchVmAgentContainer).toHaveBeenCalledTimes(1); + expect(mocks.container.markVmAgentContainerRequestInterrupted).not.toHaveBeenCalled(); + }); }); diff --git a/apps/api/tests/unit/stuck-task-container-liveness.test.ts b/apps/api/tests/unit/stuck-task-container-liveness.test.ts new file mode 100644 index 000000000..2c4ea74d4 --- /dev/null +++ b/apps/api/tests/unit/stuck-task-container-liveness.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, it, vi } from 'vitest'; + +import type { Env } from '../../src/env'; +import { getTaskRuntimeLiveness } from '../../src/scheduled/stuck-tasks'; + +function createEnv(input: { + status: 'running' | 'sleeping' | 'recovering' | 'error'; + activeWorkStatus?: 'active' | 'ended' | 'expired' | null; + rejectProbe?: boolean; +}): Env { + const inspectLifecycle = input.rejectProbe + ? vi.fn().mockRejectedValue(new Error('DO unavailable')) + : vi.fn().mockResolvedValue({ + status: input.status, + recoveryPhase: input.status === 'recovering' ? 'waking' : null, + recoveryTrigger: input.status === 'recovering' ? 'stop' : null, + activeWorkStatus: input.activeWorkStatus ?? null, + }); + return { + CF_CONTAINER_ENABLED: 'true', + DATABASE: { + prepare: vi.fn().mockReturnValue({ + bind: vi.fn().mockReturnValue({ + first: vi.fn().mockResolvedValue({ + workspace_status: 'running', + chat_session_id: 'chat-1', + node_id: 'node-1', + node_status: 'running', + health_status: 'healthy', + last_heartbeat_at: new Date().toISOString(), + node_runtime: 'cf-container', + }), + }), + }), + } as unknown as D1Database, + VM_AGENT_CONTAINER: { + idFromName: vi.fn().mockReturnValue('do-id'), + get: vi.fn().mockReturnValue({ inspectLifecycle }), + } as unknown as Env['VM_AGENT_CONTAINER'], + } as Env; +} + +const task = { project_id: 'project-1', workspace_id: 'workspace-1' }; + +describe('stuck-task Instant lifecycle liveness', () => { + it('proves live active work without relying on ACP heartbeat freshness', async () => { + await expect( + getTaskRuntimeLiveness(createEnv({ status: 'running', activeWorkStatus: 'active' }), task) + ).resolves.toMatchObject({ + live: true, + conclusive: true, + reason: 'cf_container_active_work', + }); + }); + + it.each(['sleeping', 'recovering'] as const)( + 'preserves a %s Instant lifecycle as resumable', + async (status) => { + await expect(getTaskRuntimeLiveness(createEnv({ status }), task)).resolves.toMatchObject({ + live: false, + conclusive: false, + reason: `cf_container_${status}_resumable`, + }); + } + ); + + it('allows reconciliation after the Instant lifecycle reaches a true error', async () => { + await expect( + getTaskRuntimeLiveness(createEnv({ status: 'error' }), task) + ).resolves.toMatchObject({ + live: false, + conclusive: true, + reason: 'cf_container_error', + }); + }); + + it('treats a failed lifecycle probe as inconclusive', async () => { + await expect( + getTaskRuntimeLiveness(createEnv({ status: 'running', rejectProbe: true }), task) + ).resolves.toMatchObject({ + live: false, + conclusive: false, + reason: 'cf_container_lifecycle_unknown', + }); + }); +}); diff --git a/apps/api/tests/unit/stuck-tasks.test.ts b/apps/api/tests/unit/stuck-tasks.test.ts index 84d035dbd..b79e5e8a1 100644 --- a/apps/api/tests/unit/stuck-tasks.test.ts +++ b/apps/api/tests/unit/stuck-tasks.test.ts @@ -72,10 +72,28 @@ function mockPreparedStatement(results: unknown[] = [], changes = 1) { }; } +/** + * Wire a mocked env.VM_AGENT_CONTAINER namespace whose stub.inspectLifecycle() + * returns the given inspection (or a never-resolving promise for the timeout + * case). Enables cf-container liveness-probe coverage through the real sweep. + */ +function containerBinding( + inspectLifecycle: ReturnType +): Pick { + return { + CF_CONTAINER_ENABLED: 'true', + VM_AGENT_CONTAINER: { + idFromName: vi.fn().mockReturnValue('do-id'), + get: vi.fn().mockReturnValue({ inspectLifecycle }), + } as unknown as Env['VM_AGENT_CONTAINER'], + }; +} + function createMockEnv( prepareResponses: Map = new Map(), envOverrides: Partial> = {}, - taskRunnerStatus: unknown | Error = null + taskRunnerStatus: unknown | Error = null, + containerLifecycle?: ReturnType ): Env { const mockDb = { prepare: vi.fn((sql: string) => { @@ -119,6 +137,7 @@ function createMockEnv( TASK_STUCK_DELEGATED_TIMEOUT_MS: '1860000', // 31 min NODE_HEARTBEAT_STALE_SECONDS: '180', // 3 min TASK_RUNNER: mockTaskRunnerDO, + ...(containerLifecycle ? containerBinding(containerLifecycle) : {}), ...envOverrides, } as unknown as Env; } @@ -1032,6 +1051,47 @@ describe('recoverStuckTasks', () => { }); describe('prompt dead-runtime reconciliation', () => { + function resumableRuntimeResponses(workspaceStatus: 'sleeping' | 'recovery') { + const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000).toISOString(); + return new Map([ + [ + "status IN ('queued', 'delegated', 'in_progress')", + { + results: [ + { + id: `task-${workspaceStatus}`, + project_id: 'proj-instant', + user_id: 'user-instant', + status: 'in_progress', + execution_step: 'awaiting_followup', + updated_at: tenMinutesAgo, + started_at: tenMinutesAgo, + workspace_id: 'ws-instant', + auto_provisioned_node_id: 'node-instant', + }, + ], + }, + ], + [ + 'w.chat_session_id', + { + results: [ + { + workspace_status: workspaceStatus, + chat_session_id: 'chat-instant', + node_id: 'node-instant', + node_status: workspaceStatus, + health_status: 'unhealthy', + last_heartbeat_at: tenMinutesAgo, + // Explicit: the resumable short-circuit is gated on cf-container. + node_runtime: 'cf-container', + }, + ], + }, + ], + ]); + } + function deadRuntimeResponses(taskId: string) { const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000).toISOString(); return new Map([ @@ -1102,6 +1162,30 @@ describe('recoverStuckTasks', () => { ); }); + it.each(['sleeping', 'recovery'] as const)( + 'preserves an Instant task while its workspace is %s and resumable', + async (workspaceStatus) => { + const env = createMockEnv( + resumableRuntimeResponses(workspaceStatus), + { TASK_DO_MISMATCH_GRACE_MS: '60000' }, + { + completed: true, + currentStep: 'done', + retryCount: 0, + lastStepAt: Date.now() - 10 * 60 * 1000, + } + ); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(0); + expect(result.deadRuntimeReconciled).toBe(0); + expect(result.doHealthChecked).toBe(1); + expect(cleanupTaskRun).not.toHaveBeenCalled(); + expect(syncTriggerExecutionMock).not.toHaveBeenCalled(); + } + ); + it('fails a conclusively dead runtime and records the TaskRunner RPC failure', async () => { const env = createMockEnv( deadRuntimeResponses('task-do-error'), @@ -1274,6 +1358,220 @@ describe('recoverStuckTasks', () => { }); }); + describe('runtime-aware resumable gating (V1)', () => { + // A VM (non-container) workspace stuck in 'recovery' on a dead node must NOT + // be treated as resumable — 'recovery' is a legit VM /ready-callback status, + // and pre-PR such rows fell through to the node-heartbeat staleness check. + // The resumable short-circuit is gated on node_runtime === 'cf-container'. + it('conclusively reconciles a VM recovery workspace on a dead node', async () => { + const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000).toISOString(); + const responses = new Map([ + [ + "status IN ('queued', 'delegated', 'in_progress')", + { + results: [ + { + id: 'task-vm-recovery', + project_id: 'proj-vm', + user_id: 'user-vm', + status: 'in_progress', + execution_step: 'running', + updated_at: tenMinutesAgo, + started_at: tenMinutesAgo, + workspace_id: 'ws-vm', + auto_provisioned_node_id: 'node-vm', + }, + ], + }, + ], + [ + 'w.chat_session_id', + { + results: [ + { + workspace_status: 'recovery', + chat_session_id: 'chat-vm', + node_id: 'node-vm', + node_status: 'error', + health_status: 'unhealthy', + last_heartbeat_at: tenMinutesAgo, // stale + node_runtime: 'vm', + }, + ], + }, + ], + [ + 'node_id, status FROM workspaces', + { results: [{ id: 'ws-vm', node_id: 'node-vm', status: 'recovery' }] }, + ], + [ + 'status, health_status FROM nodes', + { results: [{ id: 'node-vm', status: 'error', health_status: 'unhealthy' }] }, + ], + ["UPDATE tasks SET status = 'failed'", { results: [], changes: 1 }], + ]); + const env = createMockEnv(responses, { TASK_DO_MISMATCH_GRACE_MS: '60000' }); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(1); + expect(result.deadRuntimeReconciled).toBe(1); + // Reason proves it "fell through to node-heartbeat staleness" (node_not_live), + // not a premature workspace_recovery classification. + expect(syncTriggerExecutionMock).toHaveBeenCalledWith( + env.DATABASE, + 'task-vm-recovery', + 'failed', + expect.stringContaining('node_not_live') + ); + }); + }); + + describe('cf-container lifecycle probe through the real sweep (T2)', () => { + function cfContainerResponses(workspaceStatus = 'running') { + const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000).toISOString(); + return new Map([ + [ + "status IN ('queued', 'delegated', 'in_progress')", + { + results: [ + { + id: 'task-instant', + project_id: 'proj-instant', + user_id: 'user-instant', + status: 'in_progress', + execution_step: 'running', + updated_at: tenMinutesAgo, + started_at: tenMinutesAgo, + workspace_id: 'ws-instant', + auto_provisioned_node_id: 'node-instant', + }, + ], + }, + ], + [ + 'w.chat_session_id', + { + results: [ + { + workspace_status: workspaceStatus, + chat_session_id: 'chat-instant', + node_id: 'node-instant', + node_status: 'running', + health_status: 'healthy', + last_heartbeat_at: tenMinutesAgo, + node_runtime: 'cf-container', + }, + ], + }, + ], + [ + 'node_id, status FROM workspaces', + { results: [{ id: 'ws-instant', node_id: 'node-instant', status: workspaceStatus }] }, + ], + [ + 'status, health_status FROM nodes', + { results: [{ id: 'node-instant', status: 'running', health_status: 'healthy' }] }, + ], + ["UPDATE tasks SET status = 'failed'", { results: [], changes: 1 }], + ]); + } + + function inspection(status: string, activeWorkStatus: string | null) { + return vi.fn().mockResolvedValue({ + status, + recoveryPhase: null, + recoveryTrigger: null, + activeWorkStatus, + }); + } + + it('(a) preserves a live container with active work as conclusively live', async () => { + const env = createMockEnv( + cfContainerResponses(), + { TASK_DO_MISMATCH_GRACE_MS: '60000' }, + null, + inspection('running', 'active') + ); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(0); + expect(result.deadRuntimeReconciled).toBe(0); + }); + + it.each(['sleeping', 'recovering'] as const)( + '(b) preserves a %s container with no active work as inconclusive', + async (lifecycleStatus) => { + const env = createMockEnv( + cfContainerResponses(), + { TASK_DO_MISMATCH_GRACE_MS: '60000' }, + null, + inspection(lifecycleStatus, null) + ); + + const result = await recoverStuckTasks(env); + + // Discriminating: if 'sleeping'/'recovering' were terminal in + // TERMINAL_LIFECYCLE_STATUSES, classify() would return conclusive dead and + // this would reconcile (failedInProgress === 1). + expect(result.failedInProgress).toBe(0); + expect(result.deadRuntimeReconciled).toBe(0); + } + ); + + it.each(['stopped', 'expired', 'error'] as const)( + '(c) reconciles a %s container as conclusively dead', + async (lifecycleStatus) => { + const env = createMockEnv( + cfContainerResponses(), + { TASK_DO_MISMATCH_GRACE_MS: '60000' }, + null, + inspection(lifecycleStatus, null) + ); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(1); + expect(result.deadRuntimeReconciled).toBe(1); + expect(syncTriggerExecutionMock).toHaveBeenCalledWith( + env.DATABASE, + 'task-instant', + 'failed', + expect.stringContaining(`cf_container_${lifecycleStatus}`) + ); + } + ); + + it('(d) preserves the task when the lifecycle probe throws (inconclusive)', async () => { + const env = createMockEnv( + cfContainerResponses(), + { TASK_DO_MISMATCH_GRACE_MS: '60000' }, + null, + vi.fn().mockRejectedValue(new Error('DO unavailable')) + ); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(0); + expect(result.deadRuntimeReconciled).toBe(0); + }); + + it('(d) preserves the task when the lifecycle probe times out (inconclusive)', async () => { + const env = createMockEnv( + cfContainerResponses(), + { TASK_DO_MISMATCH_GRACE_MS: '60000', TASK_LIVENESS_PROBE_TIMEOUT_MS: '1' }, + null, + vi.fn().mockImplementation(() => new Promise(() => undefined)) + ); + + const result = await recoverStuckTasks(env); + + expect(result.failedInProgress).toBe(0); + expect(result.deadRuntimeReconciled).toBe(0); + }); + }); + describe('result structure', () => { it('returns all expected counters including heartbeatSkipped', async () => { const env = createMockEnv( diff --git a/apps/api/tests/unit/vm-agent-container-lifecycle.test.ts b/apps/api/tests/unit/vm-agent-container-lifecycle.test.ts new file mode 100644 index 000000000..6b0269903 --- /dev/null +++ b/apps/api/tests/unit/vm-agent-container-lifecycle.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, it } from 'vitest'; + +import { + classifyVmAgentContainerLiveness, + type VmAgentContainerLifecycleInspection, + type VmAgentContainerLifecycleStatus, +} from '../../src/durable-objects/vm-agent-container-lifecycle'; + +function inspection( + status: VmAgentContainerLifecycleStatus | null, + activeWorkStatus: VmAgentContainerLifecycleInspection['activeWorkStatus'] = null +): VmAgentContainerLifecycleInspection { + return { status, activeWorkStatus, recoveryPhase: null, recoveryTrigger: null }; +} + +describe('VM agent container lifecycle liveness', () => { + it('proves task-scoped liveness only while active work is registered', () => { + expect(classifyVmAgentContainerLiveness(inspection('running', 'active'))).toEqual({ + live: true, + conclusive: true, + reason: 'cf_container_active_work', + }); + }); + + it.each([ + null, + 'launching', + 'running', + 'sleeping', + 'recovering', + 'waking', + 'restoring', + 'degraded', + ] as Array)( + 'keeps %s lifecycle state resumable and inconclusive', + (status) => { + expect(classifyVmAgentContainerLiveness(inspection(status, 'ended'))).toMatchObject({ + live: false, + conclusive: false, + }); + } + ); + + it.each(['stopping', 'stopped', 'expired', 'error'] as VmAgentContainerLifecycleStatus[])( + 'treats %s as conclusively terminal', + (status) => { + expect(classifyVmAgentContainerLiveness(inspection(status))).toEqual({ + live: false, + conclusive: true, + reason: `cf_container_${status}`, + }); + } + ); +}); diff --git a/apps/api/tests/workers/acp-session-do.test.ts b/apps/api/tests/workers/acp-session-do.test.ts index e1793ce58..0048111ec 100644 --- a/apps/api/tests/workers/acp-session-do.test.ts +++ b/apps/api/tests/workers/acp-session-do.test.ts @@ -9,17 +9,33 @@ * Same issue blocks existing project-data-do.test.ts. * * DOCUMENTED COVERAGE GAPS (to add when workerd issue is resolved): - * - alarm() / checkHeartbeatTimeouts: stale session → interrupted transition * - listAcpSessionsByNode: reconciliation filtering by node + statuses * - forkAcpSession: max depth rejection, fork from failed session * - updateHeartbeat: silent ignore for terminal sessions * - transitionAcpSession: nonexistent session error * - listAcpSessions: chatSessionId filter, pagination, total count */ -import { env } from 'cloudflare:test'; -import { describe, expect,it } from 'vitest'; +import { env, runInDurableObject } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; import type { ProjectData } from '../../src/durable-objects/project-data'; +import { shouldDeferRuntimeHeartbeatTimeout } from '../../src/durable-objects/project-data/runtime-heartbeat-policy'; +import type { VmAgentContainerLifecycleStatus } from '../../src/durable-objects/vm-agent-container-lifecycle'; +import type { Env } from '../../src/env'; +import { seedNode, seedUser, seedWorkspace } from './helpers/seed-d1'; +import type { VmAgentContainerTestDouble } from './support/vm-agent-container-double'; + +/** Seed the bound VM_AGENT_CONTAINER double's lifecycle status for a node. */ +function seedContainerLifecycle( + nodeId: string, + status: VmAgentContainerLifecycleStatus, +): Promise { + const ns = env.VM_AGENT_CONTAINER!; + const stub = ns.get( + ns.idFromName(nodeId.toLowerCase()), + ) as unknown as DurableObjectStub; + return stub.__seedLifecycle(status); +} function getStub(projectId: string): DurableObjectStub { const id = env.PROJECT_DATA.idFromName(projectId); @@ -225,6 +241,105 @@ describe('ACP Session Lifecycle (Spec 027)', () => { }); }); + describe('heartbeat timeout alarm — Instant container lifecycle', () => { + // Seed a cf-container node + non-terminal workspace + a running ACP session + // whose heartbeat is stale, so the alarm's heartbeat-timeout path fires and + // the container-lifecycle policy decides preserve vs terminalize. + async function setupInstantSession(prefix: string) { + const userId = `${prefix}-user`; + const nodeId = `${prefix}-node`; + const workspaceId = `${prefix}-workspace`; + const stub = getStub(`${prefix}-project`); + const { acpSession, chatSessionId } = await createSessionPair(stub); + + await seedUser(userId); + await seedNode(nodeId, userId); + await env.DATABASE.prepare(`UPDATE nodes SET runtime = 'cf-container' WHERE id = ?`) + .bind(nodeId) + .run(); + // Workspace stays non-terminal ('sleeping') so the policy does NOT + // short-circuit on workspace status and actually reaches the container RPC. + await seedWorkspace(workspaceId, nodeId, userId, { + status: 'sleeping', + chatSessionId, + }); + await stub.transitionAcpSession(acpSession.id, 'assigned', { + actorType: 'system', + workspaceId, + nodeId, + }); + await stub.transitionAcpSession(acpSession.id, 'running', { + actorType: 'vm-agent', + actorId: nodeId, + acpSdkSessionId: `${prefix}-sdk`, + }); + return { stub, acpSession, workspaceId, nodeId }; + } + + function staleHeartbeat( + stub: DurableObjectStub, + acpSessionId: string, + passes: number, + ) { + return runInDurableObject(stub, async (instance, state) => { + state.storage.sql.exec( + `UPDATE acp_sessions SET last_heartbeat_at = ? WHERE id = ?`, + Date.now() - 10 * 60 * 1000, + acpSessionId, + ); + for (let i = 0; i < passes; i++) await instance.alarm(); + return state.storage.getAlarm(); + }); + } + + it('preserves a sleeping Instant session across repeated stale-heartbeat alarms', async () => { + const prefix = `acp-instant-sleep-${Date.now()}-${crypto.randomUUID()}`; + const { stub, acpSession, workspaceId, nodeId } = await setupInstantSession(prefix); + + // The real container is asleep (idle handback), not terminated. + await seedContainerLifecycle(nodeId, 'sleeping'); + + // Direct policy assertion proves the binding is wired and the REAL + // inspectLifecycle RPC + classifier ran — the reason is lifecycle-based, + // NOT the 'cf_container_lifecycle_binding_unavailable' short-circuit that + // masked this path when VM_AGENT_CONTAINER was unbound. + const decision = await shouldDeferRuntimeHeartbeatTimeout(env as unknown as Env, { + workspaceId, + nodeId, + }); + expect(decision).toEqual({ defer: true, reason: 'cf_container_sleeping' }); + + const nextAlarm = await staleHeartbeat(stub, acpSession.id, 2); + + // End-to-end: the session survives repeated alarms and stays scheduled. + expect((await stub.getAcpSession(acpSession.id))?.status).toBe('running'); + expect(nextAlarm).not.toBeNull(); + expect(nextAlarm!).toBeGreaterThan(Date.now()); + }); + + it('terminalizes an Instant session when the container lifecycle is terminal', async () => { + const prefix = `acp-instant-stopped-${Date.now()}-${crypto.randomUUID()}`; + const { stub, acpSession, workspaceId, nodeId } = await setupInstantSession(prefix); + + // The real container is terminated ('stopped') — the classifier treats + // this as conclusively dead, so the heartbeat timeout must NOT be deferred. + await seedContainerLifecycle(nodeId, 'stopped'); + + const decision = await shouldDeferRuntimeHeartbeatTimeout(env as unknown as Env, { + workspaceId, + nodeId, + }); + expect(decision).toEqual({ defer: false, reason: 'cf_container_stopped' }); + + await staleHeartbeat(stub, acpSession.id, 1); + + // End-to-end: with a dead container the stale session is interrupted. + // This pairing is what makes the sleeping test discriminating — moving + // 'sleeping' into TERMINAL_LIFECYCLE_STATUSES would flip both outcomes. + expect((await stub.getAcpSession(acpSession.id))?.status).toBe('interrupted'); + }); + }); + describe('transitionAcpSession — invalid transitions', () => { it('rejects pending → running (must go through assigned)', async () => { const stub = getStub('acp-invalid-pending-running'); diff --git a/apps/api/tests/workers/instant-container-lifecycle-wiring.test.ts b/apps/api/tests/workers/instant-container-lifecycle-wiring.test.ts new file mode 100644 index 000000000..a7297ec2a --- /dev/null +++ b/apps/api/tests/workers/instant-container-lifecycle-wiring.test.ts @@ -0,0 +1,176 @@ +/** + * Production-incident regression: Instant (cf-container) lifecycle wiring across + * BOTH control loops that can terminate a task/session. + * + * A single seeded cf-container node/workspace/task/ACP-session row set is driven + * through the two real reconcilers against the SAME state: + * 1. ProjectData.alarm() stale-heartbeat pass (heartbeat-timeout policy) + * 2. recoverStuckTasks() sweep (stuck-task liveness gate) + * + * With the container asleep ('sleeping', a resumable state) BOTH loops must + * preserve the work: the ACP session stays running, the task stays in_progress, + * and the workspace row survives. Once the container is terminated ('stopped'), + * the sweep MUST reconcile the now-dead runtime and fail the task. + * + * This exercises the REAL heartbeat policy (shouldDeferRuntimeHeartbeatTimeout), + * the REAL container-lifecycle classifier, and the REAL sweep — with only the + * container shell replaced by VmAgentContainerTestDouble (bound as + * VM_AGENT_CONTAINER; see tests/workers/support/). Before this wiring the policy + * short-circuited with 'cf_container_lifecycle_binding_unavailable' and the + * classifier never ran. + */ +import { env, runInDurableObject } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; + +import type { ProjectData } from '../../src/durable-objects/project-data'; +import { shouldDeferRuntimeHeartbeatTimeout } from '../../src/durable-objects/project-data/runtime-heartbeat-policy'; +import type { VmAgentContainerLifecycleStatus } from '../../src/durable-objects/vm-agent-container-lifecycle'; +import type { Env } from '../../src/env'; +import { recoverStuckTasks } from '../../src/scheduled/stuck-tasks'; +import { + seedInstallation, + seedNode, + seedProject, + seedTask, + seedUser, + seedWorkspace, +} from './helpers/seed-d1'; +import type { VmAgentContainerTestDouble } from './support/vm-agent-container-double'; + +function getProjectStub(projectId: string): DurableObjectStub { + return env.PROJECT_DATA.get( + env.PROJECT_DATA.idFromName(projectId), + ) as DurableObjectStub; +} + +function seedContainerLifecycle( + nodeId: string, + status: VmAgentContainerLifecycleStatus, +): Promise { + const ns = env.VM_AGENT_CONTAINER!; + const stub = ns.get( + ns.idFromName(nodeId.toLowerCase()), + ) as unknown as DurableObjectStub; + return stub.__seedLifecycle(status); +} + +async function getTask(taskId: string) { + return env.DATABASE.prepare('SELECT status, error_message FROM tasks WHERE id = ?') + .bind(taskId) + .first<{ status: string; error_message: string | null }>(); +} + +describe('Instant container lifecycle wiring — alarm + sweep', () => { + it('preserves a sleeping cf-container task+session across both loops, then reconciles once stopped', async () => { + const prefix = `instant-wiring-${Date.now()}-${crypto.randomUUID()}`; + const userId = `${prefix}-user`; + const installationId = `${prefix}-install`; + const projectId = `${prefix}-project`; + const nodeId = `${prefix}-node`; + const workspaceId = `${prefix}-workspace`; + const taskId = `${prefix}-task`; + const cursorKey = `test:stuck-task-cursor:${prefix}`; + const oldIso = new Date(Date.now() - 10 * 60 * 1000).toISOString(); + + // --- Seed a single cf-container row set across D1 + ProjectData DO --- + await seedUser(userId); + await seedInstallation(installationId, userId, { installationIdValue: `${prefix}-ext` }); + await seedProject(projectId, userId, installationId); + await seedNode(nodeId, userId, { status: 'running' }); + await env.DATABASE.prepare(`UPDATE nodes SET runtime = 'cf-container' WHERE id = ?`) + .bind(nodeId) + .run(); + + const stub = getProjectStub(projectId); + const chatSessionId = await stub.createSession(null, 'Instant wiring'); + const acpSession = await stub.createAcpSession({ + chatSessionId, + initialPrompt: 'Do the thing', + agentType: 'codex', + }); + + // Workspace 'running' so BOTH the policy (non-terminal) and the sweep + // (LIVE_WORKSPACE_STATUSES) reach the container-lifecycle probe. + await seedWorkspace(workspaceId, nodeId, userId, { + projectId, + status: 'running', + chatSessionId, + }); + await seedTask(taskId, projectId, userId, { + status: 'in_progress', + workspaceId, + startedAt: oldIso, + updatedAt: oldIso, + taskMode: 'task', + }); + + await stub.transitionAcpSession(acpSession.id, 'assigned', { + actorType: 'system', + workspaceId, + nodeId, + }); + await stub.transitionAcpSession(acpSession.id, 'running', { + actorType: 'vm-agent', + actorId: nodeId, + acpSdkSessionId: `${prefix}-sdk`, + }); + + // The sweep gates the container probe behind CF_CONTAINER_ENABLED and uses a + // KV scan cursor; force a bounded, deterministic scan of this task. + const sweepEnv = { + ...env, + CF_CONTAINER_ENABLED: 'true', + TASK_DO_MISMATCH_GRACE_MS: '60000', + TASK_RUN_MAX_EXECUTION_MS: '3600000', + STUCK_TASK_SCAN_CURSOR_KV_KEY: cursorKey, + } as unknown as Env; + + // ================= Phase 1: container sleeping (resumable) ================= + await seedContainerLifecycle(nodeId, 'sleeping'); + + // Policy runs the REAL inspectLifecycle RPC + classifier — reason is + // lifecycle-based, NOT the binding-unavailable short-circuit. + expect( + await shouldDeferRuntimeHeartbeatTimeout(env as unknown as Env, { workspaceId, nodeId }), + ).toEqual({ defer: true, reason: 'cf_container_sleeping' }); + + // Actor 1 — ProjectData alarm stale-heartbeat pass: session stays running. + await runInDurableObject(stub, async (instance, state) => { + state.storage.sql.exec( + `UPDATE acp_sessions SET last_heartbeat_at = ? WHERE id = ?`, + Date.now() - 10 * 60 * 1000, + acpSession.id, + ); + await instance.alarm(); + }); + expect((await stub.getAcpSession(acpSession.id))?.status).toBe('running'); + + // Actor 2 — stuck-task sweep: a resumable runtime is NOT reconciled. + await env.KV.delete(cursorKey); + const firstSweep = await recoverStuckTasks(sweepEnv); + expect(firstSweep.deadRuntimeReconciled).toBe(0); + expect((await getTask(taskId))?.status).toBe('in_progress'); + + // Actor 3 — workspace row survives. + const wsAfterSleep = await env.DATABASE.prepare(`SELECT status FROM workspaces WHERE id = ?`) + .bind(workspaceId) + .first<{ status: string }>(); + expect(wsAfterSleep?.status).toBe('running'); + + // ================= Phase 2: container stopped (terminal) ================= + await seedContainerLifecycle(nodeId, 'stopped'); + + expect( + await shouldDeferRuntimeHeartbeatTimeout(env as unknown as Env, { workspaceId, nodeId }), + ).toEqual({ defer: false, reason: 'cf_container_stopped' }); + + // The sweep now sees a conclusively-dead runtime and reconciles the task. + await env.KV.delete(cursorKey); + const secondSweep = await recoverStuckTasks(sweepEnv); + expect(secondSweep.deadRuntimeReconciled).toBeGreaterThanOrEqual(1); + + const taskAfterStop = await getTask(taskId); + expect(taskAfterStop?.status).toBe('failed'); + expect(taskAfterStop?.error_message).toContain('cf_container_stopped'); + }); +}); diff --git a/apps/api/tests/workers/instant-runtime-recovery-failure.test.ts b/apps/api/tests/workers/instant-runtime-recovery-failure.test.ts new file mode 100644 index 000000000..cee77d365 --- /dev/null +++ b/apps/api/tests/workers/instant-runtime-recovery-failure.test.ts @@ -0,0 +1,119 @@ +import { env } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; + +import { + persistRuntimeRecoveryFailed, + RUNTIME_RECOVERY_DEGRADED_MESSAGE, +} from '../../src/durable-objects/vm-agent-container-recovery'; +import type { Env } from '../../src/env'; +import { + seedInstallation, + seedNode, + seedProject, + seedTask, + seedUser, + seedWorkspace, +} from './helpers/seed-d1'; + +describe('Instant runtime terminal reconciliation with Miniflare D1', () => { + it('fails every related runtime row and the active task after recovery exhaustion', async () => { + const prefix = `runtime-failure-${Date.now()}-${crypto.randomUUID()}`; + const userId = `${prefix}-user`; + const installationId = `${prefix}-installation`; + const projectId = `${prefix}-project`; + const nodeId = `${prefix}-node`; + const workspaceId = `${prefix}-workspace`; + const chatSessionId = `${prefix}-chat`; + const agentSessionId = `${prefix}-agent`; + const taskId = `${prefix}-task`; + + await seedUser(userId); + await seedInstallation(installationId, userId, { + installationIdValue: `${prefix}-external`, + }); + await seedProject(projectId, userId, installationId); + await seedNode(nodeId, userId, { status: 'recovery', healthStatus: 'unhealthy' }); + await env.DATABASE.prepare(`UPDATE nodes SET runtime = 'cf-container' WHERE id = ?`) + .bind(nodeId) + .run(); + await seedWorkspace(workspaceId, nodeId, userId, { + projectId, + status: 'recovery', + chatSessionId, + }); + await env.DATABASE.prepare( + `INSERT INTO agent_sessions + (id, workspace_id, user_id, status, agent_type, created_at, updated_at) + VALUES (?, ?, ?, 'recovery', 'codex', datetime('now'), datetime('now'))` + ) + .bind(agentSessionId, workspaceId, userId) + .run(); + await seedTask(taskId, projectId, userId, { + status: 'in_progress', + workspaceId, + autoProvisionedNodeId: nodeId, + executionStep: 'agent_running', + }); + + await persistRuntimeRecoveryFailed(env as unknown as Env, { + nodeId, + workspaceId, + projectId, + chatSessionId, + agentSessionId, + }); + + const node = await env.DATABASE.prepare( + `SELECT status, health_status, error_message FROM nodes WHERE id = ?` + ) + .bind(nodeId) + .first>(); + const workspace = await env.DATABASE.prepare( + `SELECT status, error_message FROM workspaces WHERE id = ?` + ) + .bind(workspaceId) + .first>(); + const agent = await env.DATABASE.prepare( + `SELECT status, stopped_at, error_message FROM agent_sessions WHERE id = ?` + ) + .bind(agentSessionId) + .first>(); + const task = await env.DATABASE.prepare( + `SELECT status, execution_step, error_message FROM tasks WHERE id = ?` + ) + .bind(taskId) + .first>(); + const event = await env.DATABASE.prepare( + `SELECT from_status, to_status, actor_type, reason + FROM task_status_events WHERE task_id = ? ORDER BY created_at DESC LIMIT 1` + ) + .bind(taskId) + .first>(); + + expect(node).toEqual({ + status: 'error', + health_status: 'unhealthy', + error_message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }); + expect(workspace).toEqual({ + status: 'error', + error_message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }); + expect(agent).toMatchObject({ + status: 'error', + error_message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }); + expect(agent?.stopped_at).toEqual(expect.any(String)); + expect(task).toEqual({ + status: 'failed', + execution_step: null, + error_message: RUNTIME_RECOVERY_DEGRADED_MESSAGE, + }); + expect(event).toEqual({ + from_status: 'in_progress', + to_status: 'failed', + actor_type: 'system', + reason: 'Instant runtime recovery exhausted', + }); + }); +}); diff --git a/apps/api/tests/workers/instant-runtime-recovery-persistence.test.ts b/apps/api/tests/workers/instant-runtime-recovery-persistence.test.ts new file mode 100644 index 000000000..0ecf9c3ec --- /dev/null +++ b/apps/api/tests/workers/instant-runtime-recovery-persistence.test.ts @@ -0,0 +1,116 @@ +import { env } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; + +import { + persistRuntimeRecovered, + persistRuntimeRecovering, + RUNTIME_RECOVERING_MESSAGE, + RUNTIME_REQUEST_INTERRUPTED_MESSAGE, +} from '../../src/durable-objects/vm-agent-container-recovery'; +import type { Env } from '../../src/env'; +import { + seedInstallation, + seedNode, + seedProject, + seedUser, + seedWorkspace, +} from './helpers/seed-d1'; + +describe('Instant runtime status reconciliation with Miniflare D1', () => { + it('moves related rows through recovery to running without losing manual-retry state', async () => { + const prefix = `runtime-persistence-${Date.now()}-${crypto.randomUUID()}`; + const userId = `${prefix}-user`; + const installationId = `${prefix}-installation`; + const projectId = `${prefix}-project`; + const nodeId = `${prefix}-node`; + const workspaceId = `${prefix}-workspace`; + const chatSessionId = `${prefix}-chat`; + const agentSessionId = `${prefix}-agent`; + + await seedUser(userId); + await seedInstallation(installationId, userId, { + installationIdValue: `${prefix}-external`, + }); + await seedProject(projectId, userId, installationId); + await seedNode(nodeId, userId, { status: 'running' }); + await env.DATABASE.prepare(`UPDATE nodes SET runtime = 'cf-container' WHERE id = ?`) + .bind(nodeId) + .run(); + await seedWorkspace(workspaceId, nodeId, userId, { + projectId, + status: 'running', + chatSessionId, + }); + await env.DATABASE.prepare( + `INSERT INTO agent_sessions + (id, workspace_id, user_id, status, agent_type, created_at, updated_at) + VALUES (?, ?, ?, 'running', 'codex', datetime('now'), datetime('now'))` + ) + .bind(agentSessionId, workspaceId, userId) + .run(); + + const bindings = env as unknown as Env; + const target = { nodeId, workspaceId, projectId, chatSessionId, agentSessionId }; + await persistRuntimeRecovering(bindings, target); + + const recoveringNode = await env.DATABASE.prepare( + `SELECT status, health_status, error_message FROM nodes WHERE id = ?` + ) + .bind(nodeId) + .first>(); + const recoveringWorkspace = await env.DATABASE.prepare( + `SELECT status, error_message FROM workspaces WHERE id = ?` + ) + .bind(workspaceId) + .first>(); + const recoveringAgent = await env.DATABASE.prepare( + `SELECT status, error_message FROM agent_sessions WHERE id = ?` + ) + .bind(agentSessionId) + .first>(); + + expect(recoveringNode).toMatchObject({ + status: 'recovery', + health_status: 'unhealthy', + error_message: RUNTIME_RECOVERING_MESSAGE, + }); + expect(recoveringWorkspace).toMatchObject({ + status: 'recovery', + error_message: RUNTIME_RECOVERING_MESSAGE, + }); + expect(recoveringAgent).toMatchObject({ + status: 'recovery', + error_message: RUNTIME_RECOVERING_MESSAGE, + }); + + await persistRuntimeRecovered(bindings, target, 'manual_retry'); + + const recoveredNode = await env.DATABASE.prepare( + `SELECT status, health_status, error_message FROM nodes WHERE id = ?` + ) + .bind(nodeId) + .first>(); + const recoveredWorkspace = await env.DATABASE.prepare( + `SELECT status, error_message FROM workspaces WHERE id = ?` + ) + .bind(workspaceId) + .first>(); + const recoveredAgent = await env.DATABASE.prepare( + `SELECT status, stopped_at, error_message FROM agent_sessions WHERE id = ?` + ) + .bind(agentSessionId) + .first>(); + + expect(recoveredNode).toEqual({ + status: 'running', + health_status: 'healthy', + error_message: null, + }); + expect(recoveredWorkspace).toEqual({ status: 'running', error_message: null }); + expect(recoveredAgent).toEqual({ + status: 'running', + stopped_at: null, + error_message: RUNTIME_REQUEST_INTERRUPTED_MESSAGE, + }); + }); +}); diff --git a/apps/api/tests/workers/instant-runtime-recovery-resolver.test.ts b/apps/api/tests/workers/instant-runtime-recovery-resolver.test.ts new file mode 100644 index 000000000..d796f2311 --- /dev/null +++ b/apps/api/tests/workers/instant-runtime-recovery-resolver.test.ts @@ -0,0 +1,102 @@ +import { env } from 'cloudflare:test'; +import { drizzle } from 'drizzle-orm/d1'; +import { beforeAll, describe, expect, it } from 'vitest'; + +import * as schema from '../../src/db/schema'; +import { + resolveLiveAgentSessionForChat, + resolveLiveWorkspaceForSession, +} from '../../src/routes/chat-workspace-resolver'; +import { + seedInstallation, + seedNode, + seedProject, + seedUser, + seedWorkspace, +} from './helpers/seed-d1'; + +const PREFIX = `instant-recovery-${Date.now()}`; +const USER_ID = `${PREFIX}-user`; +const INSTALLATION_ID = `${PREFIX}-installation`; +const PROJECT_ID = `${PREFIX}-project`; +const CONTAINER_NODE_ID = `${PREFIX}-container-node`; +const VM_NODE_ID = `${PREFIX}-vm-node`; +const CONTAINER_WORKSPACE_ID = `${PREFIX}-container-workspace`; +const VM_WORKSPACE_ID = `${PREFIX}-vm-workspace`; +const CONTAINER_CHAT_ID = `${PREFIX}-container-chat`; +const VM_CHAT_ID = `${PREFIX}-vm-chat`; +const AGENT_SESSION_ID = `${PREFIX}-agent-session`; + +describe('Instant recovery resolver with real D1 state', () => { + beforeAll(async () => { + await seedUser(USER_ID); + await seedInstallation(INSTALLATION_ID, USER_ID, { + installationIdValue: `${PREFIX}-external`, + }); + await seedProject(PROJECT_ID, USER_ID, INSTALLATION_ID); + await seedNode(CONTAINER_NODE_ID, USER_ID, { status: 'error', healthStatus: 'unhealthy' }); + await seedNode(VM_NODE_ID, USER_ID, { status: 'error', healthStatus: 'unhealthy' }); + await env.DATABASE.prepare(`UPDATE nodes SET runtime = 'cf-container' WHERE id = ?`) + .bind(CONTAINER_NODE_ID) + .run(); + await seedWorkspace(CONTAINER_WORKSPACE_ID, CONTAINER_NODE_ID, USER_ID, { + projectId: PROJECT_ID, + status: 'error', + chatSessionId: CONTAINER_CHAT_ID, + }); + await seedWorkspace(VM_WORKSPACE_ID, VM_NODE_ID, USER_ID, { + projectId: PROJECT_ID, + status: 'error', + chatSessionId: VM_CHAT_ID, + }); + await env.DATABASE.prepare( + `INSERT INTO agent_sessions + (id, workspace_id, user_id, status, agent_type, created_at, updated_at) + VALUES (?, ?, ?, 'error', 'codex', datetime('now'), datetime('now'))` + ) + .bind(AGENT_SESSION_ID, CONTAINER_WORKSPACE_ID, USER_ID) + .run(); + }); + + function db() { + return drizzle(env.DATABASE, { schema }); + } + + it('allows a cf-container legacy error row to reach the Durable Object recovery path', async () => { + const resolved = await resolveLiveAgentSessionForChat(db(), { + projectId: PROJECT_ID, + sessionId: CONTAINER_CHAT_ID, + userId: USER_ID, + }); + + expect(resolved).toEqual({ + workspace: { + id: CONTAINER_WORKSPACE_ID, + nodeId: CONTAINER_NODE_ID, + nodeStatus: 'error', + nodeRuntime: 'cf-container', + }, + agentSession: { id: AGENT_SESSION_ID }, + }); + }); + + it('does not weaken the dead-node guard for a VM workspace in error', async () => { + const resolved = await resolveLiveWorkspaceForSession(db(), { + projectId: PROJECT_ID, + sessionId: VM_CHAT_ID, + userId: USER_ID, + }); + + expect(resolved).toBeNull(); + }); + + it('keeps project and user ownership predicates on the recovery path', async () => { + const resolved = await resolveLiveWorkspaceForSession(db(), { + projectId: `${PROJECT_ID}-other`, + sessionId: CONTAINER_CHAT_ID, + userId: USER_ID, + }); + + expect(resolved).toBeNull(); + }); +}); diff --git a/apps/api/tests/workers/session-snapshot-wiring.test.ts b/apps/api/tests/workers/session-snapshot-wiring.test.ts index 48fb1e39f..cffffa498 100644 --- a/apps/api/tests/workers/session-snapshot-wiring.test.ts +++ b/apps/api/tests/workers/session-snapshot-wiring.test.ts @@ -70,6 +70,8 @@ describe('session snapshot D1/R2 worker wiring', () => { chatSessionId, workspaceId, agentSessionId, + acpSessionId: 'acp-session-1', + agentType: 'openai-codex', baseCommit: 'base-commit', status: 'available', degradation: 'none', @@ -116,6 +118,8 @@ describe('session snapshot D1/R2 worker wiring', () => { await expect(manifestObject.json()).resolves.toMatchObject({ chatSessionId, workspaceId, + acpSessionId: 'acp-session-1', + agentType: 'openai-codex', artifacts: { home: { sizeBytes: 4 }, wip: { sizeBytes: 3 }, diff --git a/apps/api/tests/workers/support/test-worker-entry.ts b/apps/api/tests/workers/support/test-worker-entry.ts new file mode 100644 index 000000000..fa438d087 --- /dev/null +++ b/apps/api/tests/workers/support/test-worker-entry.ts @@ -0,0 +1,19 @@ +/** + * Test-only worker entry for the Miniflare workers pool + * (see `vitest.workers.config.ts` -> `main`). + * + * It re-exports the real API worker UNCHANGED — its default fetch/scheduled + * handler and every Durable Object class — so `SELF`, all routes, and every + * existing DO binding behave exactly as when `main` pointed at `src/index.ts`. + * + * It additionally exports `VmAgentContainerTestDouble` so the pool can bind + * `VM_AGENT_CONTAINER` to a container-less stand-in. The real `VmAgentContainer` + * cannot be instantiated under vitest-pool-workers (its `Container` base throws + * when `ctx.container === undefined`); see `vm-agent-container-double.ts`. This + * wrapper is required because vitest-pool-workers only resolves DO classes and + * the default handler from the `main` module (and only `main` is transformed by + * Vite, so the double can import the real TypeScript lifecycle helpers). + */ +export * from '../../../src/index'; +export { default } from '../../../src/index'; +export { VmAgentContainerTestDouble } from './vm-agent-container-double'; diff --git a/apps/api/tests/workers/support/vm-agent-container-double.ts b/apps/api/tests/workers/support/vm-agent-container-double.ts new file mode 100644 index 000000000..e8ea3653d --- /dev/null +++ b/apps/api/tests/workers/support/vm-agent-container-double.ts @@ -0,0 +1,70 @@ +/** + * Test-only Durable Object that stands in for VmAgentContainer under the + * Miniflare workers pool. + * + * WHY A DOUBLE (constructor investigation): the real VmAgentContainer extends + * `Container` from `@cloudflare/containers`, whose constructor throws when + * `ctx.container === undefined`: + * + * // @cloudflare/containers dist/lib/container.js (constructor) + * if (ctx.container === undefined) { + * throw new Error('Containers have not been enabled for this Durable Object + * class. Have you correctly setup your Wrangler config? ...'); + * } + * + * Under `@cloudflare/vitest-pool-workers` there is no container service attached + * to the DO, so `ctx.container` is always `undefined` and constructing the real + * class throws on instantiation. The Container base additionally creates SQLite + * schedule tables, schedules alarms, and wires a container monitor — none of + * which can run in workerd-without-containers. Binding the real class is + * therefore infeasible. + * + * This double mocks ONLY that container shell. It is a plain Durable Object + * that stores the same `lifecycleStatus` key the real container persists, and + * its `inspectLifecycle()` RPC calls the REAL shared read helper + * (`inspectStoredVmAgentContainerLifecycle`) over its own `ctx.storage`. The + * production heartbeat policy (`shouldDeferRuntimeHeartbeatTimeout`), the RPC + * shape (`VmAgentContainerLifecycleInspection`), and the terminal classifier + * (`isVmAgentContainerLifecycleTerminal`) all run on the real code path — only + * the container runtime itself is replaced. + */ +import { DurableObject } from 'cloudflare:workers'; + +import { ACTIVE_WORK_KEY } from '../../../src/durable-objects/vm-agent-container-active-work'; +import { + inspectStoredVmAgentContainerLifecycle, + type VmAgentContainerLifecycleInspection, + type VmAgentContainerLifecycleStatus, +} from '../../../src/durable-objects/vm-agent-container-lifecycle'; +import type { Env } from '../../../src/env'; + +// These keys MUST match VmAgentContainer's private storage keys so the shared +// read helper observes the same state the production DO would. `lifecycleStatus` +// and RECOVERY_STATE_KEY (`'runtimeRecovery'`) are the container's private +// constants in vm-agent-container.ts; ACTIVE_WORK_KEY is imported from its +// source module to stay drift-safe. +const LIFECYCLE_STATUS_KEY = 'lifecycleStatus'; +const RECOVERY_STATE_KEY = 'runtimeRecovery'; + +export class VmAgentContainerTestDouble extends DurableObject { + /** + * Mirrors VmAgentContainer.inspectLifecycle() exactly: same shared read + * helper, same storage keys. Keeps the real RPC contract on the code path. + */ + async inspectLifecycle(): Promise { + return inspectStoredVmAgentContainerLifecycle( + this.ctx.storage, + RECOVERY_STATE_KEY, + ACTIVE_WORK_KEY, + ); + } + + /** + * Test-only seeding RPC: writes the lifecycle status the real container would + * have persisted (e.g. 'sleeping' after idle handback, 'stopped'/'error' after + * termination) so the real classifier decides the outcome. + */ + async __seedLifecycle(status: VmAgentContainerLifecycleStatus): Promise { + await this.ctx.storage.put(LIFECYCLE_STATUS_KEY, status); + } +} diff --git a/apps/api/vitest.workers.config.ts b/apps/api/vitest.workers.config.ts index e5d9deca8..9eda4ed41 100644 --- a/apps/api/vitest.workers.config.ts +++ b/apps/api/vitest.workers.config.ts @@ -15,7 +15,12 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ plugins: [ cloudflareTest({ - main: './src/index.ts', + // Test-only entry: re-exports the real worker (default handler + every DO + // class) AND VmAgentContainerTestDouble so VM_AGENT_CONTAINER can be bound + // to a container-less stand-in. The real VmAgentContainer's Container base + // throws under vitest-pool-workers (no container service -> ctx.container + // is undefined). See tests/workers/support/*. + main: './tests/workers/support/test-worker-entry.ts', miniflare: { // 2024-04-03+ required for DO RPC (calling methods directly on stubs) compatibilityDate: '2024-04-03', @@ -60,6 +65,15 @@ export default defineConfig({ className: 'ProjectOrchestrator', useSQLite: true, }, + // Container-less stand-in for the Instant runtime container DO. The + // real VmAgentContainer cannot be instantiated here (Container base + // throws without a container service), so the heartbeat policy + stuck + // sweep exercise the real inspectLifecycle RPC + classifier against + // this double. See tests/workers/support/vm-agent-container-double.ts. + VM_AGENT_CONTAINER: { + className: 'VmAgentContainerTestDouble', + useSQLite: true, + }, }, bindings: { BASE_DOMAIN: 'test.example.com', diff --git a/apps/web/src/components/project-message-view/index.tsx b/apps/web/src/components/project-message-view/index.tsx index 464cbbc85..ac763f1ef 100644 --- a/apps/web/src/components/project-message-view/index.tsx +++ b/apps/web/src/components/project-message-view/index.tsx @@ -6,7 +6,12 @@ * TypewriterText animates the latest assistant message; historical messages * render instantly. */ -import type { ConversationItem, PlanItem, SlashCommand, ToolCallContentItem } from '@simple-agent-manager/acp-client'; +import type { + ConversationItem, + PlanItem, + SlashCommand, + ToolCallContentItem, +} from '@simple-agent-manager/acp-client'; import { mapToolCallContent, PlanModal } from '@simple-agent-manager/acp-client'; import type { AgentProfile } from '@simple-agent-manager/shared'; import { Button, Spinner } from '@simple-agent-manager/ui'; @@ -76,7 +81,15 @@ function useFloatingHeaderHeight(): [(el: HTMLDivElement | null) => void, number /** Floating session header with optional error banner and summary. */ function FloatingHeader({ - projectId, lc, onSessionMutated, onRetry, onFork, onOpenTimeline, sourceContext, onShowHierarchy, containerRef, + projectId, + lc, + onSessionMutated, + onRetry, + onFork, + onOpenTimeline, + sourceContext, + onShowHierarchy, + containerRef, }: { projectId: string; lc: ReturnType; @@ -90,7 +103,7 @@ function FloatingHeader({ }) { if (!lc.session) return null; const initialPromptFallback = !lc.hasMore - ? lc.messages.find((msg) => msg.role === 'user')?.content ?? null + ? (lc.messages.find((msg) => msg.role === 'user')?.content ?? null) : null; const taskStatus = lc.taskEmbed?.status; const hasRecoverableTaskError = Boolean( @@ -149,7 +162,9 @@ function ErrorBanner({ message, recoverable }: { message: string; recoverable: b @@ -172,7 +188,11 @@ function currentPlanToPlanItem(plan: Array<{ content: string; status: string }>) entries: plan.map((e) => ({ content: e.content, priority: 'medium' as const, - status: (e.status === 'completed' ? 'completed' : e.status === 'in_progress' ? 'in_progress' : 'pending') as 'pending' | 'in_progress' | 'completed', + status: (e.status === 'completed' + ? 'completed' + : e.status === 'in_progress' + ? 'in_progress' + : 'pending') as 'pending' | 'in_progress' | 'completed', })), timestamp: Date.now(), }; @@ -192,7 +212,11 @@ const ElapsedTime: FC<{ startedAt: number }> = ({ startedAt }) => { const interval = setInterval(update, 1000); return () => clearInterval(interval); }, [startedAt]); - return ; + return ( + + ); }; interface ProjectMessageViewProps { @@ -279,26 +303,32 @@ export const ProjectMessageView: FC = ({ const [pendingJump, setPendingJump] = useState(null); const [highlightedItemId, setHighlightedItemId] = useState(null); - const scrollAndHighlight = useCallback((itemId: string): boolean => { - const index = itemIndexById.get(itemId); - if (index === undefined) return false; - virtuosoRef.current?.scrollToIndex({ index, behavior: 'smooth', align: 'center' }); - setHighlightedItemId(itemId); - return true; - }, [itemIndexById]); - - const handleTimelineJump = useCallback((target: TimelineJumpTarget) => { - setShowTimeline(false); - // Fast path: exact message already loaded. - if (target.messageId && itemIndexById.has(target.messageId)) { - scrollAndHighlight(target.messageId); - return; - } - // Otherwise resolve via the pending-jump effect, loading older pages toward - // the target timestamp first (no-op when the history is already fully loaded). - setPendingJump(target); - void lc.loadUntil(target.timestamp); - }, [itemIndexById, scrollAndHighlight, lc]); + const scrollAndHighlight = useCallback( + (itemId: string): boolean => { + const index = itemIndexById.get(itemId); + if (index === undefined) return false; + virtuosoRef.current?.scrollToIndex({ index, behavior: 'smooth', align: 'center' }); + setHighlightedItemId(itemId); + return true; + }, + [itemIndexById] + ); + + const handleTimelineJump = useCallback( + (target: TimelineJumpTarget) => { + setShowTimeline(false); + // Fast path: exact message already loaded. + if (target.messageId && itemIndexById.has(target.messageId)) { + scrollAndHighlight(target.messageId); + return; + } + // Otherwise resolve via the pending-jump effect, loading older pages toward + // the target timestamp first (no-op when the history is already fully loaded). + setPendingJump(target); + void lc.loadUntil(target.timestamp); + }, + [itemIndexById, scrollAndHighlight, lc] + ); // Resolve a pending jump once the target (or the nearest message, after // loading settles) is available in the rendered list. @@ -336,10 +366,15 @@ export const ProjectMessageView: FC = ({ const prevMsgCountRef = useRef(0); /** Lazy-load tool content for a compact-mode tool call card. */ - const handleLoadToolContent = useCallback(async (messageId: string): Promise => { - const { content } = await getMessageToolContent(projectId, sessionId, messageId); - return (content as Array<{ type: string } & Record>).map((c) => mapToolCallContent(c)); - }, [projectId, sessionId]); + const handleLoadToolContent = useCallback( + async (messageId: string): Promise => { + const { content } = await getMessageToolContent(projectId, sessionId, messageId); + return (content as Array<{ type: string } & Record>).map((c) => + mapToolCallContent(c) + ); + }, + [projectId, sessionId] + ); // Detect newly added optimistic user messages for fade animation useEffect(() => { @@ -351,7 +386,9 @@ export const ProjectMessageView: FC = ({ if (msg && msg.role === 'user' && msg.id.startsWith('optimistic-')) { animatedUserMsgIds.add(msg.id); // Remove from set after animation completes (max 1.5s + buffer) - setTimeout(() => { animatedUserMsgIds.delete(msg.id); }, 2000); + setTimeout(() => { + animatedUserMsgIds.delete(msg.id); + }, 2000); } } } @@ -368,16 +405,19 @@ export const ProjectMessageView: FC = ({ }, [conversationItems]); const planItem = useMemo( - () => lc.currentPlan && lc.currentPlan.length > 0 ? currentPlanToPlanItem(lc.currentPlan) : null, - [lc.currentPlan], + () => + lc.currentPlan && lc.currentPlan.length > 0 ? currentPlanToPlanItem(lc.currentPlan) : null, + [lc.currentPlan] ); const canWriteSession = lc.session?.isMine !== false; - const sessionOwnerLabel = lc.session?.createdBy?.name?.trim() - || lc.session?.createdBy?.email?.split('@')[0] - || 'the creator'; + const sessionOwnerLabel = + lc.session?.createdBy?.name?.trim() || + lc.session?.createdBy?.email?.split('@')[0] || + 'the creator'; const canArchiveSession = Boolean( onCloseConversation && - (lc.taskEmbed?.taskMode === 'conversation' || (!lc.taskEmbed?.id && lc.session?.status === 'active')) + (lc.taskEmbed?.taskMode === 'conversation' || + (!lc.taskEmbed?.id && lc.session?.status === 'active')) ); // Initial load — only show full spinner when no data exists yet @@ -390,11 +430,7 @@ export const ProjectMessageView: FC = ({ } if (lc.error && !lc.session) { - return ( -
- {lc.error} -
- ); + return
{lc.error}
; } const isActive = lc.sessionState === 'active' || lc.sessionState === 'idle'; @@ -409,35 +445,40 @@ export const ProjectMessageView: FC = ({ )} {/* Connection indicator (DO WebSocket) */} - {lc.sessionState === 'active' && lc.connectionState !== 'connected' && lc.showConnectionBanner && ( - - )} + {lc.sessionState === 'active' && + lc.connectionState !== 'connected' && + lc.showConnectionBanner && ( + + )} {/* Resuming agent banner */} {lc.isResuming && ( -
+
- Resuming agent... - + Waking and restoring Instant session... + {lc.resumeStartedAt != null && }
)} {/* Resume error banner */} {lc.resumeError && ( -
- {lc.resumeError} +
+ + {lc.resumeError} +
)} @@ -445,16 +486,46 @@ export const ProjectMessageView: FC = ({ {/* Messages area — virtualized, DO-only */} {conversationItems.length === 0 ? (
- setShowTimeline(true)} sourceContext={sourceContext} onShowHierarchy={onShowHierarchy} containerRef={floatingHeaderRef} /> -
+ setShowTimeline(true)} + sourceContext={sourceContext} + onShowHierarchy={onShowHierarchy} + containerRef={floatingHeaderRef} + /> +
- {lc.sessionState === 'active' ? 'Waiting for messages...' : 'No messages in this session.'} + {lc.sessionState === 'active' + ? 'Waiting for messages...' + : 'No messages in this session.'}
) : ( -
- setShowTimeline(true)} sourceContext={sourceContext} onShowHierarchy={onShowHierarchy} containerRef={floatingHeaderRef} /> +
+ setShowTimeline(true)} + sourceContext={sourceContext} + onShowHierarchy={onShowHierarchy} + containerRef={floatingHeaderRef} + />
= ({ data={conversationItems} firstItemIndex={lc.firstItemIndex} initialTopMostItemIndex={conversationItems.length - 1} - followOutput={(isAtBottom: boolean) => isAtBottom ? 'smooth' : false} + followOutput={(isAtBottom: boolean) => (isAtBottom ? 'smooth' : false)} alignToBottom atBottomThreshold={50} atBottomStateChange={(atBottom) => lc.setShowScrollButton(!atBottom)} overscan={200} itemContent={(index, item) => ( -
+
)} @@ -489,7 +572,12 @@ export const ProjectMessageView: FC = ({
{lc.hasMore && (
-
@@ -538,11 +626,7 @@ export const ProjectMessageView: FC = ({ /> )} {planItem && ( - setShowPlanModal(false)} - /> + setShowPlanModal(false)} /> )} {/* Input area */} @@ -550,15 +634,21 @@ export const ProjectMessageView: FC = ({ { void lc.handleSendFollowUp(); }} - onUploadFiles={(files) => { void lc.handleUploadFiles(files); }} + onSend={() => { + void lc.handleSendFollowUp(); + }} + onUploadFiles={(files) => { + void lc.handleUploadFiles(files); + }} sending={lc.sendingFollowUp} uploading={lc.uploading} - placeholder={lc.agentActivity === 'prompting' || lc.agentActivity === 'responding' - ? 'Agent is working...' - : lc.sessionState === 'idle' - ? 'Send a message to resume the agent...' - : 'Send a message...'} + placeholder={ + lc.agentActivity === 'prompting' || lc.agentActivity === 'responding' + ? 'Agent is working...' + : lc.sessionState === 'idle' + ? 'Send a message to resume the agent...' + : 'Send a message...' + } transcribeApiUrl={lc.transcribeApiUrl} agentProfiles={agentProfiles} slashCommands={slashCommands} @@ -569,9 +659,7 @@ export const ProjectMessageView: FC = ({ )} {lc.sessionState === 'terminated' && (
- - This session has ended. - + This session has ended.
)} diff --git a/apps/web/src/components/project-message-view/runtimeRecoveryMessages.ts b/apps/web/src/components/project-message-view/runtimeRecoveryMessages.ts new file mode 100644 index 000000000..c58c841a8 --- /dev/null +++ b/apps/web/src/components/project-message-view/runtimeRecoveryMessages.ts @@ -0,0 +1,47 @@ +import { ApiClientError } from '../../lib/api'; + +const RUNTIME_RECOVERY_CODES = new Set([ + 'RUNTIME_RECOVERING', + 'RUNTIME_REQUEST_INTERRUPTED', + 'RUNTIME_RECOVERY_DEGRADED', + 'RUNTIME_STOPPED', +]); + +/** + * Terminal runtime code (HTTP 410): the Instant runtime is permanently stopped + * and the agent session can never be resumed. Callers must reflect termination + * in local session state (status → 'stopped') so the existing terminated + * presentation takes over — never surface a dismissible "retry" banner, which + * only invites endless futile retries against a dead runtime. + */ +export function isRuntimeStoppedError(error: unknown): boolean { + return error instanceof ApiClientError && error.code === 'RUNTIME_STOPPED'; +} + +/** + * Fallback shown when a follow-up could not be confirmed delivered and the error + * carries no recognized runtime-recovery code. + */ +export const DEFAULT_DELIVERY_ERROR_MESSAGE = + 'Your message is saved, but delivery could not be confirmed. Check the transcript and partial output before deciding whether to send it again.'; + +export function getRuntimeRecoveryMessage(error: unknown): string | null { + if (!(error instanceof ApiClientError) || !RUNTIME_RECOVERY_CODES.has(error.code)) { + return null; + } + if (error.code === 'RUNTIME_RECOVERING') { + return 'Waking and restoring the Instant session. Wait for restore to finish, then send your message.'; + } + return error.message; +} + +export function getResumeFailureMessage(error: unknown): string { + const runtimeMessage = getRuntimeRecoveryMessage(error); + if (runtimeMessage) return runtimeMessage; + + const message = error instanceof Error ? error.message : String(error); + if (message.includes('404') || /not found/i.test(message)) { + return 'Could not resume agent — workspace may have been cleaned up.'; + } + return 'Could not resume agent — please try again.'; +} diff --git a/apps/web/src/components/project-message-view/useConnectionRecovery.ts b/apps/web/src/components/project-message-view/useConnectionRecovery.ts index 37b37eda5..7692aaaa1 100644 --- a/apps/web/src/components/project-message-view/useConnectionRecovery.ts +++ b/apps/web/src/components/project-message-view/useConnectionRecovery.ts @@ -11,12 +11,28 @@ import { useCallback, useEffect, useRef, useState } from 'react'; import type { ChatConnectionState } from '../../hooks/useChatWebSocket'; import type { ChatSessionResponse } from '../../lib/api'; import { resumeAgentSession, sendFollowUpPrompt } from '../../lib/api'; -import type { SessionState } from './types'; import { - AUTO_RESUME_DELAY_MS, - DEFAULT_IDLE_TIMEOUT_MS, - RECONNECT_BANNER_DELAY_MS, -} from './types'; + DEFAULT_DELIVERY_ERROR_MESSAGE, + getResumeFailureMessage, + getRuntimeRecoveryMessage, + isRuntimeStoppedError, +} from './runtimeRecoveryMessages'; +import type { SessionState } from './types'; +import { AUTO_RESUME_DELAY_MS, DEFAULT_IDLE_TIMEOUT_MS, RECONNECT_BANNER_DELAY_MS } from './types'; + +/** Outcome callbacks for a follow-up dispatched through the resume path. */ +export interface ResumeSendHandlers { + /** Fired once the follow-up was resumed AND delivered successfully. */ + onDelivered?: () => void; + /** Fired when the resume or the delivery failed (including terminal stop). */ + onFailed?: () => void; +} + +/** A follow-up queued to send after the in-flight/next resume succeeds. */ +interface PendingFollowUp { + text: string; + handlers?: ResumeSendHandlers; +} export interface UseConnectionRecoveryOptions { sessionId: string; @@ -30,14 +46,20 @@ export interface UseConnectionRecoveryOptions { export interface UseConnectionRecoveryResult { isResuming: boolean; + /** Timestamp (ms) the current resume began, for live elapsed-time display. */ + resumeStartedAt: number | null; resumeError: string | null; showConnectionBanner: boolean; idleCountdownMs: number | null; + clearResumeError: () => void; + reportDeliveryError: (error: unknown) => void; /** Resume an idle session and optionally send a follow-up prompt. */ - resumeAndSend: (followUpText?: string) => void; + resumeAndSend: (followUpText?: string, handlers?: ResumeSendHandlers) => void; } -export function useConnectionRecovery(opts: UseConnectionRecoveryOptions): UseConnectionRecoveryResult { +export function useConnectionRecovery( + opts: UseConnectionRecoveryOptions +): UseConnectionRecoveryResult { const { sessionId, projectId, @@ -50,20 +72,81 @@ export function useConnectionRecovery(opts: UseConnectionRecoveryOptions): UseCo // Resume state const [isResuming, setIsResuming] = useState(false); + const [resumeStartedAt, setResumeStartedAt] = useState(null); const [resumeError, setResumeError] = useState(null); const hasAttemptedAutoResumeRef = useRef(false); + // Re-entrancy + ordering guards (UX4): + // - isResumingRef mirrors `isResuming` synchronously so overlapping callers + // (auto-resume timer vs. a user-initiated send) never launch two resumes. + // - resumeAttemptRef is a monotonic token; only the newest attempt's + // settlement is allowed to mutate state, so a stale (superseded or + // session-switched) resume can never clobber newer state. + // - pendingFollowUpRef carries the user's follow-up so it is delivered when + // the in-flight resume settles (piggyback instead of a second resume). + const isResumingRef = useRef(false); + const resumeAttemptRef = useRef(0); + const pendingFollowUpRef = useRef(null); + + const clearResumeError = useCallback(() => setResumeError(null), []); + + /** + * Reflect a terminal RUNTIME_STOPPED in local session state so the existing + * terminated presentation (composer disabled) takes over. Reuses the same + * `status: 'stopped'` transition the WebSocket "session stopped" path uses. + */ + const markTerminated = useCallback(() => { + setSession((prev) => (prev ? ({ ...prev, status: 'stopped' } as ChatSessionResponse) : prev)); + }, [setSession]); + + // Delivery failure: terminal stop → terminate (no banner); otherwise banner. + const applyDeliveryError = useCallback( + (error: unknown) => { + if (isRuntimeStoppedError(error)) { + setResumeError(null); + markTerminated(); + return; + } + setResumeError(getRuntimeRecoveryMessage(error) ?? DEFAULT_DELIVERY_ERROR_MESSAGE); + }, + [markTerminated] + ); + + // Resume failure: terminal stop → terminate (no banner); otherwise banner. + const applyResumeError = useCallback( + (error: unknown) => { + if (isRuntimeStoppedError(error)) { + setResumeError(null); + markTerminated(); + return; + } + setResumeError(getResumeFailureMessage(error)); + }, + [markTerminated] + ); + + const reportDeliveryError = useCallback( + (error: unknown) => applyDeliveryError(error), + [applyDeliveryError] + ); + // Connection banner debounce const [showConnectionBanner, setShowConnectionBanner] = useState(false); // Idle timer const [idleCountdownMs, setIdleCountdownMs] = useState(null); - // Reset state when switching sessions + // Reset state when switching sessions. Bumping the attempt token invalidates + // any in-flight resume from the previous session so its settlement cannot + // mutate the new session's state. useEffect(() => { setIsResuming(false); + setResumeStartedAt(null); setResumeError(null); hasAttemptedAutoResumeRef.current = false; + isResumingRef.current = false; + resumeAttemptRef.current += 1; + pendingFollowUpRef.current = null; setShowConnectionBanner(false); }, [sessionId]); @@ -91,47 +174,111 @@ export function useConnectionRecovery(opts: UseConnectionRecoveryOptions): UseCo return; }, [sessionState, cleanupAt, agentCompletedAt, isResuming]); - // Auto-resume for idle sessions (one attempt only) const agentSessionId = session?.agentSessionId ?? null; - useEffect(() => { - if ( - sessionState !== 'idle' || - isResuming || - isProvisioning || - !session?.workspaceId || - !agentSessionId || - hasAttemptedAutoResumeRef.current - ) return; + const workspaceId = session?.workspaceId ?? null; + + // Shared resume core used by BOTH the auto-resume timer and user-initiated + // sends. Guarantees a single in-flight resume (re-entrancy guard) and that + // only the newest attempt mutates state (monotonic token). + const beginResume = useCallback( + (followUp?: PendingFollowUp) => { + if (!workspaceId || !agentSessionId) { + followUp?.handlers?.onFailed?.(); + return; + } + + // Re-entrancy: a resume is already running — piggyback the follow-up onto + // it rather than launching an overlapping resume. + if (isResumingRef.current) { + if (followUp) pendingFollowUpRef.current = followUp; + return; + } - const timer = setTimeout(() => { - if (hasAttemptedAutoResumeRef.current) return; hasAttemptedAutoResumeRef.current = true; + const attempt = (resumeAttemptRef.current += 1); + isResumingRef.current = true; setIsResuming(true); + setResumeStartedAt(Date.now()); setResumeError(null); + if (followUp) pendingFollowUpRef.current = followUp; - resumeAgentSession(session.workspaceId!, agentSessionId) - .then(() => { - setSession((prev) => { - if (!prev) return prev; - return { ...prev, isIdle: false, agentCompletedAt: null } as ChatSessionResponse; - }); + resumeAgentSession(workspaceId, agentSessionId) + .then(async () => { + // Stale (superseded by a newer attempt or a session switch) — do not + // touch state or send, and leave `isResuming` to the newest attempt. + if (attempt !== resumeAttemptRef.current) return; + isResumingRef.current = false; setIsResuming(false); + setResumeStartedAt(null); + setSession((prev) => + prev ? ({ ...prev, isIdle: false, agentCompletedAt: null } as ChatSessionResponse) : prev + ); + + const pending = pendingFollowUpRef.current; + pendingFollowUpRef.current = null; + if (!pending) return; + try { + await sendFollowUpPrompt(projectId, sessionId, pending.text); + if (attempt !== resumeAttemptRef.current) return; + setResumeError(null); + pending.handlers?.onDelivered?.(); + } catch (err) { + if (attempt !== resumeAttemptRef.current) return; + applyDeliveryError(err); + pending.handlers?.onFailed?.(); + } }) .catch((err) => { + if (attempt !== resumeAttemptRef.current) return; + isResumingRef.current = false; setIsResuming(false); - const msg = err instanceof Error ? err.message : String(err); - if (msg.includes('404') || msg.includes('not found') || msg.includes('Not Found')) { - setResumeError('Could not resume agent \u2014 workspace may have been cleaned up.'); - } else { - console.error('Auto-resume failed:', msg); - setResumeError('Could not resume agent \u2014 please try again.'); - } + setResumeStartedAt(null); + const pending = pendingFollowUpRef.current; + pendingFollowUpRef.current = null; + applyResumeError(err); + pending?.handlers?.onFailed?.(); }); + }, + [ + workspaceId, + agentSessionId, + projectId, + sessionId, + setSession, + applyDeliveryError, + applyResumeError, + ] + ); + + // Auto-resume reads `beginResume` through a ref so the timer effect does NOT + // list it as a dependency — otherwise a `session` object identity change + // during the 2s window would clear and restart the timer, delaying (or + // starving) auto-resume indefinitely. + const beginResumeRef = useRef(beginResume); + useEffect(() => { + beginResumeRef.current = beginResume; + }, [beginResume]); + + // Auto-resume for idle sessions (one attempt only) + useEffect(() => { + if ( + sessionState !== 'idle' || + isResuming || + isProvisioning || + !workspaceId || + !agentSessionId || + hasAttemptedAutoResumeRef.current + ) + return; + + const timer = setTimeout(() => { + if (hasAttemptedAutoResumeRef.current || isResumingRef.current) return; + beginResumeRef.current(); }, AUTO_RESUME_DELAY_MS); return () => clearTimeout(timer); - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [sessionState, isResuming, isProvisioning, session?.workspaceId, agentSessionId]); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [sessionState, isResuming, isProvisioning, workspaceId, agentSessionId]); // Debounced connection banner useEffect(() => { @@ -143,45 +290,22 @@ export function useConnectionRecovery(opts: UseConnectionRecoveryOptions): UseCo return () => clearTimeout(timer); }, [connectionState]); - // Resume an idle session and optionally send a follow-up prompt via REST API - const resumeAndSend = useCallback((followUpText?: string) => { - if (!session?.workspaceId || !agentSessionId) return; - hasAttemptedAutoResumeRef.current = true; - setIsResuming(true); - setResumeError(null); - - resumeAgentSession(session.workspaceId, agentSessionId) - .then(async () => { - setSession((prev) => { - if (!prev) return prev; - return { ...prev, isIdle: false, agentCompletedAt: null } as ChatSessionResponse; - }); - setIsResuming(false); - if (followUpText) { - try { - await sendFollowUpPrompt(projectId, sessionId, followUpText); - } catch { - // Agent may be offline — message is still persisted via DO - } - } - }) - .catch((err) => { - setIsResuming(false); - const msg = err instanceof Error ? err.message : String(err); - if (msg.includes('404') || msg.includes('not found') || msg.includes('Not Found')) { - setResumeError('Could not resume agent \u2014 workspace may have been cleaned up.'); - } else { - console.error('Agent resume failed:', msg); - setResumeError('Could not resume agent \u2014 please try again.'); - } - }); - }, [session?.workspaceId, agentSessionId, projectId, sessionId, setSession]); + // Resume an idle session and optionally send a follow-up prompt via REST API. + const resumeAndSend = useCallback( + (followUpText?: string, handlers?: ResumeSendHandlers) => { + beginResume(followUpText ? { text: followUpText, handlers } : undefined); + }, + [beginResume] + ); return { isResuming, + resumeStartedAt, resumeError, showConnectionBanner, idleCountdownMs, + clearResumeError, + reportDeliveryError, resumeAndSend, }; } diff --git a/apps/web/src/components/project-message-view/useSessionLifecycle.ts b/apps/web/src/components/project-message-view/useSessionLifecycle.ts index 953f4f9d7..54d573cf0 100644 --- a/apps/web/src/components/project-message-view/useSessionLifecycle.ts +++ b/apps/web/src/components/project-message-view/useSessionLifecycle.ts @@ -1,22 +1,51 @@ import type { NodeResponse, WorkspaceResponse } from '@simple-agent-manager/shared'; -import { DEFAULT_CHAT_LOAD_UNTIL_MAX_PAGES, DEFAULT_CHAT_SESSION_MESSAGE_LIMIT, DEFAULT_CHAT_SESSION_MESSAGE_MAX } from '@simple-agent-manager/shared'; +import { + DEFAULT_CHAT_LOAD_UNTIL_MAX_PAGES, + DEFAULT_CHAT_SESSION_MESSAGE_LIMIT, + DEFAULT_CHAT_SESSION_MESSAGE_MAX, +} from '@simple-agent-manager/shared'; import { useCallback, useEffect, useRef, useState } from 'react'; import type { ChatConnectionState } from '../../hooks/useChatWebSocket'; import { useChatWebSocket } from '../../hooks/useChatWebSocket'; import { useTokenRefresh } from '../../hooks/useTokenRefresh'; import { useWorkspacePorts } from '../../hooks/useWorkspacePorts'; -import type { ChatMessageResponse, ChatSessionDetailResponse, ChatSessionResponse, SessionStateSnapshot } from '../../lib/api'; -import { cancelAgentPrompt, getChatSession, getNode, getTerminalToken, getTranscribeApiUrl, getWorkspace, resetIdleTimer, sendFollowUpPrompt, uploadSessionFiles } from '../../lib/api'; +import type { + ChatMessageResponse, + ChatSessionDetailResponse, + ChatSessionResponse, + SessionStateSnapshot, +} from '../../lib/api'; +import { + cancelAgentPrompt, + getChatSession, + getNode, + getTerminalToken, + getTranscribeApiUrl, + getWorkspace, + resetIdleTimer, + sendFollowUpPrompt, + uploadSessionFiles, +} from '../../lib/api'; import { mergeMessages } from '../../lib/merge-messages'; import { isWorkspaceOperational } from '../../lib/workspace-status-utils'; import type { SessionState } from './types'; import type { AgentActivityState } from './types'; -import { CHAT_FALLBACK_POLL_MS, deriveSessionState, IDLE_TIMEOUT_MS, isWorkingActivity, VIRTUAL_START } from './types'; +import { + CHAT_FALLBACK_POLL_MS, + deriveSessionState, + IDLE_TIMEOUT_MS, + isWorkingActivity, + VIRTUAL_START, +} from './types'; import { useActivityVerifyTimer } from './useActivityVerifyTimer'; import { useConnectionRecovery } from './useConnectionRecovery'; -type FilePanelState = { mode: 'browse' | 'view' | 'diff' | 'git-status'; path?: string; line?: number | null } | null; +type FilePanelState = { + mode: 'browse' | 'view' | 'diff' | 'git-status'; + path?: string; + line?: number | null; +} | null; function parsePlanContent(content: string): SessionStateSnapshot['currentPlan'] | null { try { @@ -39,7 +68,9 @@ function hashPlanContent(plan: SessionStateSnapshot['currentPlan']): string { function getPlanFingerprint(state: SessionStateSnapshot | null | undefined): string { if (!state) return 'no-state'; - return state.planUpdatedAt ? `updated:${state.planUpdatedAt}` : `content:${hashPlanContent(state.currentPlan)}`; + return state.planUpdatedAt + ? `updated:${state.planUpdatedAt}` + : `content:${hashPlanContent(state.currentPlan)}`; } export interface UseSessionLifecycleResult { @@ -59,7 +90,9 @@ export interface UseSessionLifecycleResult { sendingFollowUp: boolean; uploading: boolean; isResuming: boolean; + resumeStartedAt: number | null; resumeError: string | null; + clearResumeError: () => void; connectionState: ChatConnectionState; showConnectionBanner: boolean; retryWs: () => void; @@ -89,7 +122,7 @@ export function useSessionLifecycle( projectId: string, sessionId: string, isProvisioning: boolean, - _onSessionMutated?: () => void, + _onSessionMutated?: () => void ): UseSessionLifecycleResult { const [session, setSession] = useState(null); const [taskEmbed, setTaskEmbed] = useState(null); @@ -103,8 +136,12 @@ export function useSessionLifecycle( // can read current state without stale closures. const messagesRef = useRef([]); const hasMoreRef = useRef(false); - useEffect(() => { messagesRef.current = messages; }, [messages]); - useEffect(() => { hasMoreRef.current = hasMore; }, [hasMore]); + useEffect(() => { + messagesRef.current = messages; + }, [messages]); + useEffect(() => { + hasMoreRef.current = hasMore; + }, [hasMore]); const [workspace, setWorkspace] = useState(null); const [node, setNode] = useState(null); @@ -115,7 +152,10 @@ export function useSessionLifecycle( const [agentActivity, setAgentActivity] = useState('idle'); const [currentPlan, setCurrentPlan] = useState(null); const [promptStartedAt, setPromptStartedAt] = useState(null); - const clearActivity = useCallback(() => { setAgentActivity('idle'); setPromptStartedAt(null); }, []); + const clearActivity = useCallback(() => { + setAgentActivity('idle'); + setPromptStartedAt(null); + }, []); const hydratePlan = useCallback((s: SessionStateSnapshot | null | undefined) => { if (!s) return; @@ -131,18 +171,21 @@ export function useSessionLifecycle( onStateSnapshot: hydratePlan, }); - const hydrateState = useCallback((s: SessionStateSnapshot | null | undefined) => { - if (!s) return; - if (isWorkingActivity(s.activity)) { - setAgentActivity(s.activity); - setPromptStartedAt(s.promptStartedAt ?? null); - startVerifyDecayTimer(); - } else { - clearActivity(); - stopVerifyDecayTimer(); - } - hydratePlan(s); - }, [clearActivity, hydratePlan, startVerifyDecayTimer, stopVerifyDecayTimer]); + const hydrateState = useCallback( + (s: SessionStateSnapshot | null | undefined) => { + if (!s) return; + if (isWorkingActivity(s.activity)) { + setAgentActivity(s.activity); + setPromptStartedAt(s.promptStartedAt ?? null); + startVerifyDecayTimer(); + } else { + clearActivity(); + stopVerifyDecayTimer(); + } + hydratePlan(s); + }, + [clearActivity, hydratePlan, startVerifyDecayTimer, stopVerifyDecayTimer] + ); const [filePanel, setFilePanel] = useState(null); @@ -162,59 +205,87 @@ export function useSessionLifecycle( const transcribeApiUrl = getTranscribeApiUrl(); // ── DO WebSocket (sole message source) ── - const { connectionState, wsRef, retry: retryWs } = useChatWebSocket({ + const { + connectionState, + wsRef, + retry: retryWs, + } = useChatWebSocket({ projectId, sessionId, enabled: session?.status === 'active', - onMessage: useCallback((msg: ChatMessageResponse) => { - setMessages((prev) => mergeMessages(prev, [msg], 'append')); + onMessage: useCallback( + (msg: ChatMessageResponse) => { + setMessages((prev) => mergeMessages(prev, [msg], 'append')); - if (msg.role === 'plan' && msg.content) { - const parsed = parsePlanContent(msg.content); - if (parsed) setCurrentPlan(parsed); - } - // Streaming agent output: show 'responding' heuristic, but arm the SHARED - // verify-before-decay timer instead of a blind decay. The blind timer used to - // clobber onAgentActivity's verified timer and flip to idle during long tool calls. - if (msg.role !== 'user') { - setAgentActivity('responding'); - startVerifyDecayTimer(); - } - }, [startVerifyDecayTimer]), + if (msg.role === 'plan' && msg.content) { + const parsed = parsePlanContent(msg.content); + if (parsed) setCurrentPlan(parsed); + } + // Streaming agent output: show 'responding' heuristic, but arm the SHARED + // verify-before-decay timer instead of a blind decay. The blind timer used to + // clobber onAgentActivity's verified timer and flip to idle during long tool calls. + if (msg.role !== 'user') { + setAgentActivity('responding'); + startVerifyDecayTimer(); + } + }, + [startVerifyDecayTimer] + ), onSessionStopped: useCallback(() => { - setSession((prev) => prev ? { ...prev, status: 'stopped' } : prev); + setSession((prev) => (prev ? { ...prev, status: 'stopped' } : prev)); setAgentActivity('idle'); setPromptStartedAt(null); // Stop any pending verify timer so it can't re-arm and flash the bar back on. stopVerifyDecayTimer(); }, [stopVerifyDecayTimer]), - onCatchUp: useCallback((catchUpMessages: ChatMessageResponse[], catchUpSession: ChatSessionResponse, state?: SessionStateSnapshot | null) => { - setSession(catchUpSession); - setMessages((prev) => mergeMessages(prev, catchUpMessages, 'replace')); - hydrateState(state); - }, [hydrateState]), - onAgentCompleted: useCallback((agentCompletedAt: number) => { - setSession((prev) => prev ? { ...prev, agentCompletedAt, isIdle: true } as ChatSessionResponse : prev); - setAgentActivity('idle'); - setPromptStartedAt(null); - // Stop any pending verify timer so it can't re-arm and flash the bar back on. - stopVerifyDecayTimer(); - }, [stopVerifyDecayTimer]), - onAgentActivity: useCallback((activity: 'prompting' | 'idle' | 'recovering' | 'error', promptStartedAt?: number | null) => { - const working = activity === 'prompting' || activity === 'recovering'; - setAgentActivity(working ? activity : 'idle'); - setPromptStartedAt(working ? (promptStartedAt ?? Date.now()) : null); - if (working) { - // Arm the shared verify-before-decay timer (prevents false idle during long tool calls). - startVerifyDecayTimer(); - } else { - // Authoritative idle from the DO: stop any pending verify timer. + onCatchUp: useCallback( + ( + catchUpMessages: ChatMessageResponse[], + catchUpSession: ChatSessionResponse, + state?: SessionStateSnapshot | null + ) => { + setSession(catchUpSession); + setMessages((prev) => mergeMessages(prev, catchUpMessages, 'replace')); + hydrateState(state); + }, + [hydrateState] + ), + onAgentCompleted: useCallback( + (agentCompletedAt: number) => { + setSession((prev) => + prev ? ({ ...prev, agentCompletedAt, isIdle: true } as ChatSessionResponse) : prev + ); + setAgentActivity('idle'); + setPromptStartedAt(null); + // Stop any pending verify timer so it can't re-arm and flash the bar back on. stopVerifyDecayTimer(); - } - }, [startVerifyDecayTimer, stopVerifyDecayTimer]), - onSessionUpdated: useCallback((updates: Partial>) => { - setSession((prev) => prev ? { ...prev, ...updates } : prev); - }, []), + }, + [stopVerifyDecayTimer] + ), + onAgentActivity: useCallback( + ( + activity: 'prompting' | 'idle' | 'recovering' | 'error', + promptStartedAt?: number | null + ) => { + const working = activity === 'prompting' || activity === 'recovering'; + setAgentActivity(working ? activity : 'idle'); + setPromptStartedAt(working ? (promptStartedAt ?? Date.now()) : null); + if (working) { + // Arm the shared verify-before-decay timer (prevents false idle during long tool calls). + startVerifyDecayTimer(); + } else { + // Authoritative idle from the DO: stop any pending verify timer. + stopVerifyDecayTimer(); + } + }, + [startVerifyDecayTimer, stopVerifyDecayTimer] + ), + onSessionUpdated: useCallback( + (updates: Partial>) => { + setSession((prev) => (prev ? { ...prev, ...updates } : prev)); + }, + [] + ), }); // Connection recovery (banner debounce, idle timer, auto-resume) @@ -259,7 +330,9 @@ export function useSessionLifecycle( } }, [projectId, sessionId, hydrateState]); - useEffect(() => { void loadSession(); }, [loadSession]); + useEffect(() => { + void loadSession(); + }, [loadSession]); // Fetch workspace and node details useEffect(() => { @@ -312,7 +385,7 @@ export function useSessionLifecycle( workspace?.url ?? undefined, session?.workspaceId ?? undefined, terminalToken ?? undefined, - isWorkspaceRunning, + isWorkspaceRunning ); // Degraded fallback while the DO WebSocket is unavailable. Connected active @@ -331,9 +404,10 @@ export function useSessionLifecycle( try { // Poll only the most-recent window — mergeReplace preserves the fully // loaded history, so polling must NOT re-fetch the whole conversation. - const data: ChatSessionDetailResponse = await getChatSession( - projectId, sessionId, { signal: abortController.signal, limit: DEFAULT_CHAT_SESSION_MESSAGE_LIMIT }, - ); + const data: ChatSessionDetailResponse = await getChatSession(projectId, sessionId, { + signal: abortController.signal, + limit: DEFAULT_CHAT_SESSION_MESSAGE_LIMIT, + }); if (data.session.id !== sessionId) return; const newLastId = data.messages[data.messages.length - 1]?.id ?? ''; const taskStatus = data.session.task?.status ?? ''; @@ -353,7 +427,9 @@ export function useSessionLifecycle( pollInFlight = false; } }; - const pollInterval = setInterval(() => { void pollActiveSession(); }, CHAT_FALLBACK_POLL_MS); + const pollInterval = setInterval(() => { + void pollActiveSession(); + }, CHAT_FALLBACK_POLL_MS); return () => { clearInterval(pollInterval); @@ -375,7 +451,12 @@ export function useSessionLifecycle( if (result.cleanupAt) { setSession((prev) => { if (!prev) return prev; - return { ...prev, cleanupAt: result.cleanupAt, isIdle: false, agentCompletedAt: null } as ChatSessionResponse; + return { + ...prev, + cleanupAt: result.cleanupAt, + isIdle: false, + agentCompletedAt: null, + } as ChatSessionResponse; }); } }) @@ -384,66 +465,91 @@ export function useSessionLifecycle( // Optimistic user message const optimisticId = `optimistic-${crypto.randomUUID()}`; - setMessages((prev) => [...prev, { - id: optimisticId, - sessionId, - role: 'user', - content: trimmed, - toolMetadata: null, - createdAt: Date.now(), - }]); + setMessages((prev) => [ + ...prev, + { + id: optimisticId, + sessionId, + role: 'user', + content: trimmed, + toolMetadata: null, + createdAt: Date.now(), + }, + ]); // Persist via DO WebSocket if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) { - wsRef.current.send(JSON.stringify({ - type: 'message.send', - sessionId, - content: trimmed, - role: 'user', - })); + wsRef.current.send( + JSON.stringify({ + type: 'message.send', + sessionId, + content: trimmed, + role: 'user', + }) + ); } - // For idle sessions, resume first then send the prompt + // For idle sessions, resume first then send the prompt. The composer is + // cleared only when delivery is confirmed (onDelivered); a failed resume + // or delivery resets the working state and keeps the typed text as the + // manual-retry affordance (onFailed). if (sessionState === 'idle' && session?.workspaceId && session?.agentSessionId) { - recovery.resumeAndSend(trimmed); + recovery.resumeAndSend(trimmed, { + onDelivered: () => { + setFollowUp(''); + }, + onFailed: () => { + setAgentActivity('idle'); + }, + }); } else { // Forward prompt to the running agent via REST API try { await sendFollowUpPrompt(projectId, sessionId, trimmed); - } catch { - // Agent may be offline — message is still persisted via DO. + // Delivery confirmed — clear any stale recovery banner and the composer. + recovery.clearResumeError(); + setFollowUp(''); + } catch (err) { + // reportDeliveryError terminates the session on a terminal RUNTIME_STOPPED + // (composer disabled) or shows the recovery banner otherwise. Reset the + // working state and keep the composer text so the user can retry. + recovery.reportDeliveryError(err); setAgentActivity('idle'); } } - - setFollowUp(''); } finally { setSendingFollowUp(false); } }; // Upload files - const handleUploadFiles = useCallback(async (files: FileList | File[]) => { - const fileArray = Array.from(files); - if (fileArray.length === 0) return; - setUploading(true); - try { - const result = await uploadSessionFiles(projectId, sessionId, fileArray); - const names = result.files.map((f) => f.name).join(', '); - setMessages((prev) => [...prev, { - id: `optimistic-upload-${crypto.randomUUID()}`, - sessionId, - role: 'user' as const, - content: `Uploaded ${result.files.length} file${result.files.length > 1 ? 's' : ''}: ${names}`, - toolMetadata: null, - createdAt: Date.now(), - }]); - } catch (err) { - console.error('File upload failed:', err); - } finally { - setUploading(false); - } - }, [projectId, sessionId]); + const handleUploadFiles = useCallback( + async (files: FileList | File[]) => { + const fileArray = Array.from(files); + if (fileArray.length === 0) return; + setUploading(true); + try { + const result = await uploadSessionFiles(projectId, sessionId, fileArray); + const names = result.files.map((f) => f.name).join(', '); + setMessages((prev) => [ + ...prev, + { + id: `optimistic-upload-${crypto.randomUUID()}`, + sessionId, + role: 'user' as const, + content: `Uploaded ${result.files.length} file${result.files.length > 1 ? 's' : ''}: ${names}`, + toolMetadata: null, + createdAt: Date.now(), + }, + ]); + } catch (err) { + console.error('File upload failed:', err); + } finally { + setUploading(false); + } + }, + [projectId, sessionId] + ); // Cancel the current in-flight prompt via REST API const cancellingRef = useRef(false); @@ -488,58 +594,94 @@ export function useSessionLifecycle( // Load older pages until a target timestamp is covered (or no more history). // Used by timeline jump-to-message for the rare oversized/guard-trimmed session // where the target predates the loaded window, so a jump never dead-clicks. - const loadUntil = useCallback(async (targetTimestamp: number) => { - let oldest = messagesRef.current[0]?.createdAt ?? Infinity; - if (oldest <= targetTimestamp) return; - let more = hasMoreRef.current; - if (!more) return; - - setLoadingMore(true); - try { - let before: number | undefined = oldest === Infinity ? undefined : oldest; - const accumulated: ChatMessageResponse[] = []; - // Safety bound: never loop unbounded even if the server misreports hasMore. - const maxPages = Number.parseInt( - import.meta.env.VITE_CHAT_LOAD_UNTIL_MAX_PAGES || '', - 10, - ) || DEFAULT_CHAT_LOAD_UNTIL_MAX_PAGES; - let pages = 0; - while (more && oldest > targetTimestamp && pages++ < maxPages) { - const data = await getChatSession(projectId, sessionId, { - before, - limit: DEFAULT_CHAT_SESSION_MESSAGE_LIMIT, - }); - if (data.messages.length === 0) { more = false; break; } - accumulated.unshift(...data.messages); - oldest = data.messages[0]!.createdAt; - before = oldest; - more = data.hasMore; - } - if (accumulated.length > 0) { - setMessages((prev) => { - const merged = mergeMessages(prev, accumulated, 'prepend'); - const actualAdded = merged.length - prev.length; - setFirstItemIndex((fi) => fi - actualAdded); - return merged; - }); + const loadUntil = useCallback( + async (targetTimestamp: number) => { + let oldest = messagesRef.current[0]?.createdAt ?? Infinity; + if (oldest <= targetTimestamp) return; + let more = hasMoreRef.current; + if (!more) return; + + setLoadingMore(true); + try { + let before: number | undefined = oldest === Infinity ? undefined : oldest; + const accumulated: ChatMessageResponse[] = []; + // Safety bound: never loop unbounded even if the server misreports hasMore. + const maxPages = + Number.parseInt(import.meta.env.VITE_CHAT_LOAD_UNTIL_MAX_PAGES || '', 10) || + DEFAULT_CHAT_LOAD_UNTIL_MAX_PAGES; + let pages = 0; + while (more && oldest > targetTimestamp && pages++ < maxPages) { + const data = await getChatSession(projectId, sessionId, { + before, + limit: DEFAULT_CHAT_SESSION_MESSAGE_LIMIT, + }); + if (data.messages.length === 0) { + more = false; + break; + } + accumulated.unshift(...data.messages); + oldest = data.messages[0]!.createdAt; + before = oldest; + more = data.hasMore; + } + if (accumulated.length > 0) { + setMessages((prev) => { + const merged = mergeMessages(prev, accumulated, 'prepend'); + const actualAdded = merged.length - prev.length; + setFirstItemIndex((fi) => fi - actualAdded); + return merged; + }); + } + setHasMore(more); + } finally { + setLoadingMore(false); } - setHasMore(more); - } finally { - setLoadingMore(false); - } - }, [projectId, sessionId]); + }, + [projectId, sessionId] + ); return { - session, messages, hasMore, loading, error, setError, sessionState, taskEmbed, - workspace, node, detectedPorts, - followUp, setFollowUp, sendingFollowUp, uploading, - isResuming: recovery.isResuming, resumeError: recovery.resumeError, - connectionState, showConnectionBanner: recovery.showConnectionBanner, retryWs, - agentActivity, currentPlan, promptStartedAt, - firstItemIndex, showScrollButton, setShowScrollButton, + session, + messages, + hasMore, + loading, + error, + setError, + sessionState, + taskEmbed, + workspace, + node, + detectedPorts, + followUp, + setFollowUp, + sendingFollowUp, + uploading, + isResuming: recovery.isResuming, + resumeStartedAt: recovery.resumeStartedAt, + resumeError: recovery.resumeError, + clearResumeError: recovery.clearResumeError, + connectionState, + showConnectionBanner: recovery.showConnectionBanner, + retryWs, + agentActivity, + currentPlan, + promptStartedAt, + firstItemIndex, + showScrollButton, + setShowScrollButton, idleCountdownMs: recovery.idleCountdownMs, - filePanel, setFilePanel, handleFileClick, handleOpenFileBrowser, handleOpenGitChanges, - handleCancelPrompt, handleSendFollowUp, handleUploadFiles, - loadMore, loadUntil, loadingMore, transcribeApiUrl, wsRef, + filePanel, + setFilePanel, + handleFileClick, + handleOpenFileBrowser, + handleOpenGitChanges, + handleCancelPrompt, + handleSendFollowUp, + handleUploadFiles, + loadMore, + loadUntil, + loadingMore, + transcribeApiUrl, + wsRef, }; } diff --git a/apps/web/src/lib/session-utils.ts b/apps/web/src/lib/session-utils.ts index 9c34fd31e..ab77ea8c8 100644 --- a/apps/web/src/lib/session-utils.ts +++ b/apps/web/src/lib/session-utils.ts @@ -7,6 +7,11 @@ const ACTIVE_HOST_STATUSES: ReadonlySet = new Set([ 'ready', 'prompting', ]); +const RECOVERABLE_SESSION_STATUSES: ReadonlySet = new Set([ + 'running', + 'recovery', + 'sleeping', +]); /** * Determines if a session should be treated as active (visible in tabs). @@ -14,7 +19,7 @@ const ACTIVE_HOST_STATUSES: ReadonlySet = new Set([ * is still alive despite a non-running status (orphan recovery). */ export function isSessionActive(session: AgentSession): boolean { - if (session.status === 'running') return true; + if (RECOVERABLE_SESSION_STATUSES.has(session.status)) return true; if (!session.hostStatus) return false; return ACTIVE_HOST_STATUSES.has(session.hostStatus); } @@ -24,7 +29,7 @@ export function isSessionActive(session: AgentSession): boolean { * These are "orphaned" sessions -- alive on the VM but previously hidden. */ export function isOrphanedSession(session: AgentSession): boolean { - if (session.status === 'running') return false; + if (RECOVERABLE_SESSION_STATUSES.has(session.status)) return false; if (!session.hostStatus) return false; return ACTIVE_HOST_STATUSES.has(session.hostStatus); } diff --git a/apps/web/tests/playwright/instant-runtime-recovery-audit.spec.ts b/apps/web/tests/playwright/instant-runtime-recovery-audit.spec.ts new file mode 100644 index 000000000..277d1c311 --- /dev/null +++ b/apps/web/tests/playwright/instant-runtime-recovery-audit.spec.ts @@ -0,0 +1,291 @@ +import { expect, type Page, type Route, test } from '@playwright/test'; + +import { assertNoOverflow, makeMockUser, screenshot, seedTheme } from './audit-helpers'; + +const PROJECT_ID = 'proj-instant-recovery-audit'; +const SESSION_ID = 'chat-instant-recovery-audit'; +const WORKSPACE_ID = 'workspace-instant-recovery-audit'; +const AGENT_SESSION_ID = 'agent-instant-recovery-audit'; +const NOW = Date.now(); + +type RecoveryScenario = 'normal' | 'waking' | 'degraded' | 'interrupted' | 'stopped'; + +const PROJECT = { + id: PROJECT_ID, + name: 'Instant Recovery Audit', + repository: 'test-user/instant-recovery-audit', + defaultBranch: 'main', + createdAt: new Date(NOW - 86_400_000).toISOString(), + updatedAt: new Date(NOW).toISOString(), +}; + +const MESSAGES = [ + { + id: 'message-before-break', + sessionId: SESSION_ID, + role: 'user', + content: 'Keep the staged, unstaged, and untracked markers while I step away.', + toolMetadata: null, + createdAt: NOW - 1_200_000, + sequence: 1, + }, + { + id: 'message-before-idle', + sessionId: SESSION_ID, + role: 'assistant', + content: + 'The work is partway complete. I saved the safe checkpoint and will continue from this same chat when you return.', + toolMetadata: null, + createdAt: NOW - 1_190_000, + sequence: 2, + }, +]; + +function makeSession(scenario: RecoveryScenario) { + const idle = scenario === 'waking'; + return { + id: SESSION_ID, + workspaceId: WORKSPACE_ID, + taskId: null, + topic: 'Resume the exact Instant session after a short break', + status: 'active', + messageCount: MESSAGES.length, + createdAt: NOW - 1_300_000, + updatedAt: NOW - 1_190_000, + endedAt: null, + cleanupAt: NOW + 3_600_000, + isIdle: idle, + agentCompletedAt: idle ? NOW - 60_000 : null, + agentSessionId: AGENT_SESSION_ID, + agentType: 'codex', + }; +} + +function respond(route: Route, status: number, body: unknown) { + return route.fulfill({ + status, + contentType: 'application/json', + body: JSON.stringify(body), + }); +} + +async function setupMocks(page: Page, scenario: RecoveryScenario) { + const session = makeSession(scenario); + const auth = makeMockUser({ + email: 'instant-recovery@example.com', + name: 'Instant Recovery Auditor', + role: 'superadmin', + sessionId: 'browser-session-instant-recovery', + userId: 'user-instant-recovery', + }); + + await page.addInitScript(() => { + window.localStorage.setItem('sam-onboarding-wizard-dismissed-user-instant-recovery', 'true'); + }); + + await page.route('**/api/**', async (route) => { + const request = route.request(); + const path = new URL(request.url()).pathname; + + if (path === '/api/auth/get-session') return respond(route, 200, auth); + if (path.startsWith('/api/notifications')) { + return respond(route, 200, { notifications: [], unreadCount: 0 }); + } + if (path === '/api/github/installations') return respond(route, 200, []); + if (path === '/api/credentials') return respond(route, 200, []); + if (path === '/api/trial/status') return respond(route, 200, { available: false }); + if (path === '/api/agents') return respond(route, 200, { agents: [] }); + if (path === '/api/projects') { + return respond(route, 200, { projects: [PROJECT], nextCursor: null }); + } + if (path === '/api/workspaces') return respond(route, 200, []); + if (path === `/api/projects/${PROJECT_ID}`) return respond(route, 200, PROJECT); + if (path === `/api/projects/${PROJECT_ID}/sessions/${SESSION_ID}`) { + return respond(route, 200, { + session, + messages: MESSAGES, + hasMore: false, + state: { + activity: 'idle', + activityAt: NOW - 1_190_000, + statusError: null, + currentPlan: [ + { content: 'Preserve the current filesystem checkpoint', status: 'completed' }, + { content: 'Continue after the user returns', status: 'pending' }, + ], + planUpdatedAt: NOW - 1_190_000, + promptStartedAt: null, + agentType: 'codex', + lastStopReason: null, + }, + }); + } + if (path === `/api/projects/${PROJECT_ID}/sessions`) { + return respond(route, 200, { sessions: [session], total: 1, hasMore: false }); + } + if (path === `/api/projects/${PROJECT_ID}/tasks`) { + return respond(route, 200, { tasks: [], total: 0 }); + } + if (path === `/api/projects/${PROJECT_ID}/agent-profiles`) { + return respond(route, 200, { items: [] }); + } + if (path.startsWith(`/api/projects/${PROJECT_ID}/commands`)) { + return respond(route, 200, { commands: [] }); + } + if (path === `/api/workspaces/${WORKSPACE_ID}`) { + return respond(route, 200, { + id: WORKSPACE_ID, + name: 'instant-recovery-audit', + nodeId: 'node-instant-recovery-audit', + projectId: PROJECT_ID, + status: scenario === 'waking' ? 'sleeping' : 'running', + }); + } + if (path === '/api/nodes/node-instant-recovery-audit') { + return respond(route, 200, { + id: 'node-instant-recovery-audit', + name: 'Instant runtime', + status: scenario === 'waking' ? 'sleeping' : 'ready', + }); + } + if (path === '/api/terminal/token') { + return respond(route, 200, { token: 'visual-audit-token' }); + } + if (path.endsWith('/idle-reset')) { + return respond(route, 200, { cleanupAt: NOW + 3_600_000 }); + } + if (path === `/api/workspaces/${WORKSPACE_ID}/agent-sessions/${AGENT_SESSION_ID}/resume`) { + if (scenario === 'waking') { + // Intentionally keep the real resume request pending so the visible + // waking/restoring state can be audited without timing races. + return; + } + return respond(route, 200, { id: AGENT_SESSION_ID, status: 'running' }); + } + if (path === `/api/projects/${PROJECT_ID}/sessions/${SESSION_ID}/prompt`) { + if (scenario === 'degraded') { + return respond(route, 503, { + error: 'RUNTIME_RECOVERY_DEGRADED', + message: + 'The Instant runtime woke without a usable checkpoint. Your transcript is safe, but verify local files before continuing — résumé 🚧 ' + + 'A_very_long_unbroken_marker_name_that_must_wrap_without_pushing_the_Dismiss_button_off_screen_0123456789.', + }); + } + if (scenario === 'interrupted') { + return respond(route, 409, { + error: 'RUNTIME_REQUEST_INTERRUPTED', + message: + 'The request was interrupted before the agent confirmed it. Your message is saved — résumé 🚧 ' + + 'An_extremely_long_unbroken_disposition_token_that_must_wrap_without_pushing_Dismiss_offscreen_0123456789. ' + + 'Resend when you are ready.', + }); + } + if (scenario === 'stopped') { + // HTTP 410 — terminal. The runtime is permanently stopped and can never + // be resumed; the client reflects this by terminating the session. + return respond(route, 410, { + error: 'RUNTIME_STOPPED', + message: 'The Instant runtime has stopped and cannot be resumed.', + }); + } + return respond(route, 200, { status: 'accepted', sessionId: AGENT_SESSION_ID }); + } + + return respond(route, 200, {}); + }); +} + +async function openConversation(page: Page) { + await page.goto(`/projects/${PROJECT_ID}/chat/${SESSION_ID}`); + await expect(page.getByRole('log', { name: 'Conversation' })).toBeVisible(); + await expect(page.getByText('The work is partway complete.')).toBeVisible(); +} + +async function sendFollowUp(page: Page, content: string) { + const composer = page.getByPlaceholder(/send a message/i); + await composer.fill(content); + await page.getByRole('button', { name: /^send$/i }).click(); +} + +for (const theme of ['dark', 'light'] as const) { + test.describe(`Instant runtime recovery — ${theme}`, () => { + test.beforeEach(async ({ page }) => { + await seedTheme(page, theme); + }); + + test('normal conversation preserves context', async ({ page }) => { + await setupMocks(page, 'normal'); + await openConversation(page); + await screenshot(page, `instant-recovery-normal-${theme}`); + await assertNoOverflow(page); + }); + + test('waking and restoring remains inline with the transcript', async ({ page }) => { + await setupMocks(page, 'waking'); + await openConversation(page); + await sendFollowUp(page, 'Continue from the exact checkpoint.'); + await expect(page.getByText('Waking and restoring Instant session...')).toBeVisible(); + await screenshot(page, `instant-recovery-waking-${theme}`); + await assertNoOverflow(page); + }); + + test('long degraded message wraps and remains dismissible', async ({ page }) => { + await setupMocks(page, 'degraded'); + await openConversation(page); + await sendFollowUp(page, 'Continue after the replacement.'); + const recoveryAlert = page.getByRole('alert').filter({ hasText: 'transcript is safe' }); + await expect(recoveryAlert).toBeVisible(); + await expect(page.getByRole('button', { name: 'Dismiss' })).toBeVisible(); + await expect(page.getByText(/