diff --git a/doc/execution-semantics.md b/doc/execution-semantics.md index 0253f8cb459..12afad1ec74 100644 --- a/doc/execution-semantics.md +++ b/doc/execution-semantics.md @@ -2075,6 +2075,22 @@ the duration of remote cleanup. Existing Stop receipts and process cleanup still destructive lifecycle or binding changes. Operator plugin restarts and external machine changes can still interrupt startup. +A contextual replacement task can retain creation recovery when its selected +worker supports inspection and its driver declares `creationOperationVersion: 1`. +Core saves that opt-in and the prior opaque binding in the existing creation +journal before acquisition, and sends acquisition to that selected worker. +An ambiguous failure makes existing recovery work due. The inspection result's +optional `creationLease` recovers identity and binding for the requested +`creationOperation`; it is not termination evidence. Contextual task recovery +requires this field; a generic sandbox cleanup envelope is not a substitute. +The existing exact-task Stop, status and completion receipts still gate ownership release. Older journal +entries without the recorded opt-in do not infer a resource identity from the +creation operation ID. Once recovered identity is durably queued for cleanup, +Core can discharge that acquisition's missing-response placeholder. Accepted +host handlers and received result observations still drain before verified +retirement; a late reply re-enters observation tracking without restoring the +worker's expired invocation authority. + Cancellation records a preparation fence while holding the run row lock. Native runtime selection checks that fence, the running status, and the current startup controller lease in the same transaction that creates the native coordinator. diff --git a/packages/plugins/sdk/src/protocol.ts b/packages/plugins/sdk/src/protocol.ts index 77189311219..bf43887f757 100644 --- a/packages/plugins/sdk/src/protocol.ts +++ b/packages/plugins/sdk/src/protocol.ts @@ -709,6 +709,7 @@ export interface PluginEnvironmentInspectLeaseParams extends PluginEnvironmentDr } export interface PluginEnvironmentLeaseObservation { + creationLease?: PluginEnvironmentLease; creationCleanup?: import("./environment-creation-cleanup.js").PluginEnvironmentCreationCleanup; providerLeaseId: string; state: "active" | "stopped" | "destroyed" | "unknown"; diff --git a/server/src/__tests__/task-creation-recovery.integration.test.ts b/server/src/__tests__/task-creation-recovery.integration.test.ts new file mode 100644 index 00000000000..ad2659f90a7 --- /dev/null +++ b/server/src/__tests__/task-creation-recovery.integration.test.ts @@ -0,0 +1,247 @@ +import { randomUUID } from "node:crypto"; +import { mkdtemp, writeFile, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { and, eq } from "drizzle-orm"; +import { describe, expect, it, vi } from "vitest"; +import { agents, environmentLeases, environments, heartbeatRuns, issues, nativeRunFinalizations, plugins, sandboxCreationOperations, type Db } from "@paperclipai/db"; +import type { PluginEnvironmentLease } from "@paperclipai/plugin-sdk"; +import { setupRunControlDatabase, callbacks, runControlPostgresSupport } from "../services/heartbeat/__tests__/run-control-fixtures.js"; +import { createHeartbeatRecovery, type HeartbeatRecoveryDependencies } from "../services/heartbeat/recovery.js"; +import { createHeartbeatRunControl } from "../services/heartbeat/run-control.js"; +import { environmentRuntimeService } from "../services/environment-runtime.js"; +import { environmentService } from "../services/environments.js"; +import { NativeExecutionContext } from "../services/native-runtime/native-execution-context.js"; +import { requestNativeStartupCleanup, verifyNativeStartupCleanup } from "../services/native-runtime/native-startup-cleanup.js"; +import { prepareNativeHeartbeatRun } from "../services/native-runtime/prepare-native-run.js"; +import { hasRemoteTerminationReceipt, remoteTerminationReceipt } from "../services/remote-execution-termination.js"; +import { createPluginWorkerManager, type PluginWorkerHandle, type PluginWorkerManager } from "../services/plugin-worker-manager.js"; +import { nextEnvironmentCleanupAt } from "../services/environment-cleanup-state.js"; + +vi.mock("../services/local-service-supervisor.js", async importOriginal => ({ + ...await importOriginal(), terminateLocalService: vi.fn(), +})); + +describe.skipIf(!runControlPostgresSupport.supported)("task creation recovery", () => { +let db: Db; +const database = setupRunControlDatabase("task-creation-recovery-", value => { db = value; }); + +async function fixture(version = true, inspector = true) { + const { company, agent } = await database.fixture(); + const companyId = company.id, agentId = agent.id, runId = randomUUID(), issueId = randomUUID(); + const pluginId = randomUUID(), environmentId = randomUUID(), leaseId = randomUUID(), taskId = randomUUID(), leaseOwner = randomUUID(); + await db.update(agents).set({ adapterType: "paperclip_runner" }).where(eq(agents.id, agentId)); + await db.insert(issues).values({ id: issueId, companyId, title: "Recover task creation", status: "in_progress", assigneeAgentId: agentId }); + const [run] = await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, issueId, nativeIssueId: issueId, + runtimeMode: "native", invocationSource: "manual", status: "running", contextSnapshot: { issueId } }).returning(); + await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId)); + await prepareNativeHeartbeatRun({ db, run, issue: { id: issueId, title: "Recover task creation", description: null }, environmentLeaseId: leaseId }); + await db.insert(nativeRunFinalizations).values({ runId, companyId, issueId, phase: "observed", attempt: 1, + controllerGeneration: 1, leaseOwner, leaseExpiresAt: new Date(Date.now() + 60_000) }); + await db.insert(plugins).values({ id: pluginId, pluginKey: pluginId, packageName: "task-recovery-fixture", version: "1", status: "ready", + manifestJson: { capabilities: ["environment.drivers.register"], environmentDrivers: [{ driverKey: "tasks", supportsTasks: true, + ...(version ? { creationOperationVersion: 1 } : {}) }] } as never }); + await db.insert(environments).values({ id: environmentId, companyId, name: environmentId, driver: "plugin", + config: { pluginKey: pluginId, driverKey: "tasks", driverConfig: {} } }); + const binding = { agentId, runId: "predecessor", computerId: "saved-computer", opaque: "saved-binding" }; + await db.insert(environmentLeases).values({ id: leaseId, companyId, environmentId, heartbeatRunId: runId, issueId, + provider: "plugin:fixture:tasks", providerLeaseId: taskId, status: "active", leasePolicy: "reuse_by_environment", + metadata: { driver: "plugin", pluginId, pluginKey: pluginId, driverKey: "tasks", agentId, + runnerTaskLease: true, runnerTaskReused: true, providerMetadata: binding } }); + const context = new NativeExecutionContext({ binding: { companyId, agentId, runId } } as never, 30_000); + const stopped = new Set([taskId]), calls: string[] = []; + let acquireSuccess = false, mismatched = false, inspectUnknown = false, genericCleanup = false; + let acquiredId = "", observed: ((lease: PluginEnvironmentLease) => Promise) | undefined; + const recoveredBinding = () => ({ ...binding, runId }); + const acquire = vi.fn(async (_method: string, params: any, _timeout: unknown, _sink: unknown, settled?: () => void, + _context?: string, observe?: (lease: PluginEnvironmentLease) => Promise) => { + acquiredId = params.creationOperationId; observed = observe; + if (!acquireSuccess) throw new Error("Acquire reply was lost"); + const result = { providerLeaseId: acquiredId, metadata: recoveredBinding() }; + await observe?.(result); settled?.(); return result; + }); + const originalWorker = { supportedMethods: ["environmentTask", "environmentAcquireLease", "environmentReleaseLease", ...(inspector ? ["environmentInspectLease"] : [])], call: acquire } as unknown as PluginWorkerHandle; + const replacementAcquire = vi.fn(async () => ({ providerLeaseId: randomUUID() })); + const oldWorker = { supportedMethods: ["environmentTask", "environmentAcquireLease"], call: replacementAcquire } as unknown as PluginWorkerHandle; + let currentWorker = originalWorker; + const workers = { getWorker: () => currentWorker, isRunning: () => true, drainStartup: vi.fn(async () => {}), retireStartup: vi.fn(async () => {}), recoverStartupAcquisition: vi.fn(), + call: vi.fn(async (_plugin: string, method: string, params: any, ...rest: any[]) => { + calls.push(method === "environmentTask" ? params.operation.kind : method); + if (method === "environmentAcquireLease") return currentWorker.call(method, params, ...rest as []); + if (method === "environmentInspectLease") { + expect(params).toMatchObject({ companyId, environmentId, creationOperation: { id: acquiredId, runId }, + leaseMetadata: { creationOperationVersion: 1, runnerTaskCreationContextId: context.id, agentId, runId, providerMetadata: binding } }); + return { providerLeaseId: acquiredId, state: "unknown", ...(!inspectUnknown ? { + creationLease: { providerLeaseId: mismatched ? "foreign-task" : acquiredId, metadata: recoveredBinding() }, + } : {}), ...(genericCleanup ? { creationCleanup: { providerLeaseId: acquiredId, companyId, environmentId, + runId, attemptId: acquiredId, accountFingerprint: "a".repeat(64), labels: {} } } : {}) }; + } + if (method === "environmentReleaseLease") return undefined; + if (params.operation.kind === "stop") stopped.add(params.taskId); + if (params.operation.kind === "status") return { kind: "status", taskId: params.taskId, + phase: stopped.has(params.taskId) ? "cancelled" : "running", executionStopped: stopped.has(params.taskId) }; + if (params.operation.kind === "describe") return { kind: "description", taskId: params.taskId, + runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }; + return { kind: "accepted", taskId: params.taskId }; + }), + } as unknown as PluginWorkerManager; + const runtime = environmentRuntimeService(db, { pluginWorkerManager: workers }); + const reload = async () => (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0]!; + const owner = async () => (await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, runId)))[0]!; + const creation = async () => (await db.select().from(sandboxCreationOperations).where(eq(sandboxCreationOperations.runId, runId)))[0]!; + const start = () => runtime.task({ companyId, leaseId, operation: { kind: "describe" }, executionContext: context }); + const request = async () => { + await requestNativeStartupCleanup(db, { companyId, agentId, runId, leaseId, leaseOwner, + contextId: context.id, deadlineAtMs: context.deadlineAtMs, requiresPreparationStop: false }); + const lease = (await environmentService(db).getLeaseById(leaseId))!; + await environmentService(db).releaseLease(leaseId, "released", { cleanupStatus: "success", + remoteExecutionTermination: remoteTerminationReceipt(lease, { providerLeaseId: taskId, state: "stopped" }) }); + }; + const deps = callbacks(db); + deps.verifyStartupCleanup = run => verifyNativeStartupCleanup(db, run, async () => {}); + const control = createHeartbeatRunControl(db, deps); + const recovery = createHeartbeatRecovery(db, { ...deps, environmentRuntime: runtime, environmentsSvc: environmentService(db), + runtimeEnv: {}, activeRunExecutionPromises: new Set>(), acknowledgeRemoteStop: async () => {}, + resumeRemoteStopComments: async () => {} } as unknown as HeartbeatRecoveryDependencies); + return { companyId, pluginId, deps, runId, context, calls, acquired: () => acquiredId, binding, creation, start, request, runtime, recovery, control, + reload, owner, originalWorker, replacementAcquire, acquire, workers, + succeed: () => { acquireSuccess = true; }, replaceWorker: () => { currentWorker = oldWorker; }, + mismatch: () => { mismatched = true; }, unknown: () => { inspectUnknown = true; }, + legacyCleanup: () => { inspectUnknown = true; genericCleanup = true; }, + lateReply: () => observed!({ providerLeaseId: acquiredId, metadata: recoveredBinding() }) }; +} + +it.each(["pending", "needs_reconciliation", "late_first", "inspection_first"])("recovers an opted-in lost task reply through verified cleanup: %s", async mode => { + const f = await fixture(); + try { + await expect(f.start()).rejects.toThrow(); + expect(await f.creation()).toMatchObject({ metadata: { creationOperationVersion: 1, providerMetadata: f.binding } }); + expect((await f.creation()).nextAttemptAt.getTime()).toBeLessThanOrEqual(Date.now()); + await f.request(); + await expect(f.control.cancelRunInternal(f.runId, "Startup cleanup retry")).rejects.toThrow("cleanup is still pending"); + if (mode === "needs_reconciliation") await db.update(sandboxCreationOperations).set({ state: "needs_reconciliation", + createdAt: new Date(Date.now() - 31 * 60_000), nextAttemptAt: new Date(0) }) + .where(eq(sandboxCreationOperations.id, (await f.creation()).id)); + if (mode === "late_first") await f.lateReply(); + const restartedRuntime = environmentRuntimeService(db, { pluginWorkerManager: f.workers }); + await restartedRuntime.flushDeferredOrphanCleanups(); + if (mode === "inspection_first") await f.lateReply(); + const creation = await f.creation(); + if (mode !== "late_first") expect(f.workers.recoverStartupAcquisition).toHaveBeenCalledWith(f.context.id, creation.id); + expect(creation.state).toBe("cleanup_pending"); + expect(await db.select().from(environmentLeases).where(and(eq(environmentLeases.heartbeatRunId, f.runId), + eq(environmentLeases.providerLeaseId, f.acquired())))).toHaveLength(1); + expect((await f.reload()).status).toBe("running"); + expect((await f.recovery.sweepPendingCleanupLeases({ explicitRetry: { companyId: f.companyId, runId: f.runId, actorId: "operator" } })).destroyed).toBe(1); + expect(f.calls).toContain("stop"); expect(f.calls).toContain("status"); expect(f.calls).toContain("complete"); + expect(await f.control.cancelRunInternal(f.runId, "Startup cleanup retry")).toMatchObject({ status: "cancelled" }); + expect((await f.owner()).leaseOwner).toBeNull(); + expect((await f.creation()).state).toBe("resolved"); + expect(await nextEnvironmentCleanupAt(db)).toBeNull(); + } finally { f.context.finish(); } +}); + +it("cancels through actual worker IPC after reconstructing an acquire that never replies", async () => { + const f = await fixture(), manager = createPluginWorkerManager(); + const directory = await mkdtemp(path.join(tmpdir(), "task-creation-recovery-ipc-")), entrypointPath = path.join(directory, "worker.cjs"); + await writeFile(entrypointPath, ` +const readline = require("node:readline"); +const send = (id, result) => process.stdout.write(JSON.stringify({jsonrpc:"2.0",id,result})+"\\n"); +readline.createInterface({input:process.stdin,crlfDelay:Infinity}).on("line", line => { + const {id,method,params:p} = JSON.parse(line); + if (method === "initialize") send(id,{ok:true,supportedMethods:["environmentTask","environmentAcquireLease","environmentReleaseLease","environmentInspectLease"]}); + else if (method === "environmentAcquireLease") return; + else if (method === "environmentInspectLease") send(id,{providerLeaseId:p.creationOperation.id,state:"unknown",creationLease:{providerLeaseId:p.creationOperation.id,metadata:{...p.leaseMetadata.providerMetadata,runId:p.creationOperation.runId}}}); + else if (method === "environmentTask") send(id,p.operation.kind === "status" ? {kind:"status",taskId:p.taskId,phase:"cancelled",executionStopped:true} : {kind:"accepted",taskId:p.taskId}); + else if (method === "shutdown") {send(id,{});setImmediate(()=>process.exit(0));} + else send(id,{}); +}); +`); + try { + const handle = await manager.startWorker(f.pluginId, { entrypointPath, manifest: { id: f.pluginId, apiVersion: 1, version: "1.0.0", + displayName: "Recovery fixture", description: "Recovery fixture", author: "Paperclip", categories: ["automation"], capabilities: [], + entrypoints: { worker: "worker.cjs" } }, config: {}, instanceInfo: { instanceId: "fixture", hostVersion: "1.0.0" }, apiVersion: 1, + hostHandlers: {}, autoRestart: false }); + const call = handle.call.bind(handle); + vi.spyOn(handle, "call").mockImplementation(((method: string, params: unknown, timeout: number, ...rest: unknown[]) => + (call as (...args: unknown[]) => Promise)(method, params, method === "environmentAcquireLease" ? 10 : timeout, ...rest)) as typeof handle.call); + Object.assign(f.workers, { getWorker: manager.getWorker.bind(manager), isRunning: manager.isRunning.bind(manager), + call: manager.call.bind(manager), drainStartup: manager.drainStartup!.bind(manager), retireStartup: manager.retireStartup!.bind(manager), + recoverStartupAcquisition: manager.recoverStartupAcquisition!.bind(manager) }); + f.deps.verifyStartupCleanup = run => verifyNativeStartupCleanup(db, run, async () => manager.retireStartup!(f.context.id)); + await expect(f.start()).rejects.toThrow("reconcile the same task"); + expect(handle.call).toHaveBeenCalledWith("environmentAcquireLease", expect.anything(), 30_000, undefined, expect.any(Function), f.context.id, expect.any(Function)); + await f.request(); + await f.runtime.flushDeferredOrphanCleanups(); + expect((await f.creation()).state).toBe("cleanup_pending"); + expect((await f.recovery.sweepPendingCleanupLeases({ backoffMs: 0 })).destroyed).toBe(1); + expect(await f.control.cancelRunInternal(f.runId, "Startup cleanup retry")).toMatchObject({ status: "cancelled" }); + expect((await f.owner()).leaseOwner).toBeNull(); + expect((await f.creation()).state).toBe("resolved"); + expect(handle.status).toBe("running"); + } finally { f.context.finish(); await manager.stopAll(); await rm(directory, { recursive: true, force: true }); } +}, 25_000); + +it.each(["legacy", "manifest_only"])("does not infer a deterministic identity for %s task creation", async mode => { + const f = await fixture(mode !== "legacy", false); + try { + await expect(f.start()).rejects.toThrow(); + const operation = await f.creation(); + expect(operation.metadata.creationOperationVersion).toBeUndefined(); + await db.update(sandboxCreationOperations).set({ nextAttemptAt: new Date(0) }).where(eq(sandboxCreationOperations.id, operation.id)); + await f.request(); await f.runtime.flushDeferredOrphanCleanups(); + expect((await f.creation()).state).toBe("needs_reconciliation"); + expect(f.calls).not.toContain("environmentInspectLease"); + await expect(f.control.cancelRunInternal(f.runId, "Startup cleanup retry")).rejects.toThrow("cleanup is still pending"); + expect((await f.owner()).leaseOwner).not.toBeNull(); + } finally { f.context.finish(); } +}); + +it.each(["mismatched", "unknown"])("retains ownership when inspection is %s", async mode => { + const f = await fixture(); + try { + await expect(f.start()).rejects.toThrow(); await f.request(); + if (mode === "mismatched") f.mismatch(); else f.unknown(); + await f.runtime.flushDeferredOrphanCleanups(); + expect((await f.creation()).state).toBe("pending"); + expect((await f.creation()).nextAttemptAt.getTime()).toBeGreaterThan(Date.now()); + await expect(f.control.cancelRunInternal(f.runId, "Startup cleanup retry")).rejects.toThrow("cleanup is still pending"); + expect((await f.owner()).leaseOwner).not.toBeNull(); + } finally { f.context.finish(); } +}); + +it.each(["pending", "needs_reconciliation"] as const)("does not substitute a generic cleanup envelope for a contextual task identity: %s", async initialState => { + const f = await fixture(); + try { + await expect(f.start()).rejects.toThrow(); await f.request(); f.legacyCleanup(); + if (initialState === "needs_reconciliation") await db.update(sandboxCreationOperations).set({ state: initialState, nextAttemptAt: new Date(0) }) + .where(eq(sandboxCreationOperations.id, (await f.creation()).id)); + await f.runtime.flushDeferredOrphanCleanups(); + const swept = await f.recovery.sweepPendingCleanupLeases({ backoffMs: 0 }); + const cancelled = await f.control.cancelRunInternal(f.runId, "Startup cleanup retry").then(() => true, () => false); + const operation = await f.creation(), lease = operation.leaseId ? await environmentService(db).getLeaseById(operation.leaseId) : null; + expect({ state: operation.state, cleanupLease: operation.leaseId, destroyed: swept.destroyed, + stoppedReceipt: lease ? hasRemoteTerminationReceipt(lease) : false, cancelled, ownerRetained: (await f.owner()).leaseOwner !== null }) + .toEqual({ state: initialState, cleanupLease: null, destroyed: 0, stoppedReceipt: false, cancelled: false, ownerRetained: true }); + expect(operation.nextAttemptAt.getTime()).toBeGreaterThan(Date.now()); + } finally { f.context.finish(); } +}); + +it("pins the opted-in acquisition worker across the journal commit", async () => { + const f = await fixture(); f.succeed(); + const transaction = db.transaction.bind(db); + const spy = vi.spyOn(db, "transaction").mockImplementation(((work: any) => transaction(async tx => { + const result = await work(tx); + if (result?.metadata?.runnerTaskCreationContextId === f.context.id) f.replaceWorker(); + return result; + })) as typeof db.transaction); + try { + const result = await f.start(); + expect(f.acquire).toHaveBeenCalledTimes(1); + expect(f.replacementAcquire).not.toHaveBeenCalled(); + expect(result).toMatchObject({ kind: "description", taskId: f.acquired() }); + expect((await f.creation()).state).toBe("attached"); + } finally { spy.mockRestore(); f.context.finish(); } +}); + +}); diff --git a/server/src/services/environment-cleanup-state.ts b/server/src/services/environment-cleanup-state.ts index e9d1fa6bf5f..444bc090600 100644 --- a/server/src/services/environment-cleanup-state.ts +++ b/server/src/services/environment-cleanup-state.ts @@ -10,7 +10,9 @@ export function recoverableSandboxCreationCondition() { sql`coalesce(${sandboxCreationOperations.metadata}->>'creationReconciliationReason', '') <> 'company_deleted'`, sql`${sandboxCreationOperations.metadata}->>'runnerTaskLease' = 'true'`, sql`jsonb_typeof(${sandboxCreationOperations.metadata}->'runnerTaskCreationContextId') = 'string' and ${sandboxCreationOperations.metadata}->>'runnerTaskCreationContextId' <> ''`, - sql`jsonb_typeof(${sandboxCreationOperations.metadata}->'creationProviderLeaseId') = 'string' and ${sandboxCreationOperations.metadata}->>'creationProviderLeaseId' <> ''`, + or(sql`jsonb_typeof(${sandboxCreationOperations.metadata}->'creationProviderLeaseId') = 'string' and ${sandboxCreationOperations.metadata}->>'creationProviderLeaseId' <> ''`, + and(sql`${sandboxCreationOperations.metadata}->>'creationOperationVersion' = '1'`, + sql`jsonb_typeof(${sandboxCreationOperations.metadata}->'providerMetadata') = 'object'`)), )); } diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index 7ffdb2c379a..bb3c56d268f 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -1,6 +1,6 @@ import { HttpError } from "../errors.js"; import { captureEnvironmentAcquisitionDiagnostic } from "./environment-acquisition-diagnostics.js"; -import { createOwnedSandbox, recoverSandboxCreations } from "./sandbox-creation-operations.js"; +import { createOwnedSandbox, recoverSandboxCreations, recoverableTaskCreation } from "./sandbox-creation-operations.js"; import { dueRunLeaseReleaseCondition, pendingRunLeaseReleaseCondition } from "./environment-lease-release-state.js"; import { flushDeferredEnvironmentCleanup, registerDeferredEnvironmentCleanup, writeEnvironmentCleanup } from "./environment-cleanup-notifications.js"; import { executeEnvironmentTask, pluginEnvironmentSupportsTasks, settleDispatchedEnvironmentTask } from "./environment-task-runtime.js"; @@ -4100,7 +4100,8 @@ export function environmentRuntimeService( async flushDeferredOrphanCleanups(): Promise<{ recovered: number; pending: number }> { orphanRecoveryFlush ??= (async () => { await recoverSandboxCreations(db, async operation => { - if (!operation.pluginId || operation.metadata.sandboxProviderPlugin !== true) return null; + const taskCreation = recoverableTaskCreation(operation); + if (!operation.pluginId || (operation.metadata.sandboxProviderPlugin !== true && !taskCreation)) return null; if (!options.pluginWorkerManager) throw new Error("Sandbox creation recovery worker is unavailable"); const plugin = await pluginRegistryService(db).getById(operation.pluginId); const driver = plugin?.manifestJson.environmentDrivers?.find(value => value.driverKey === operation.driverKey); @@ -4109,6 +4110,11 @@ export function environmentRuntimeService( !options.pluginWorkerManager.getWorker(operation.pluginId)?.supportedMethods.includes("environmentInspectLease")) { throw new Error("Sandbox creation recovery worker is unavailable"); } + if (taskCreation) return options.pluginWorkerManager.call(operation.pluginId, "environmentInspectLease", { + driverKey: operation.driverKey, companyId: operation.companyId, environmentId: operation.originalEnvironmentId, + config: operation.config, providerLeaseId: operation.id, creationOperation: { id: operation.id, runId: operation.runId }, + leaseMetadata: { ...operation.metadata, runId: operation.runId }, + }); let resolved: SandboxEnvironmentConfig; try { resolved = await resolveSandboxCleanupConfigSecrets(db, operation.companyId, @@ -4136,6 +4142,9 @@ export function environmentRuntimeService( config: stripSandboxProviderEnvelope(resolved), providerLeaseId: operation.id, creationOperation: { id: operation.id, runId: operation.runId }, }); + }, operation => { + const contextId = operation.metadata.runnerTaskCreationContextId; + if (typeof contextId === "string") options.pluginWorkerManager?.recoverStartupAcquisition?.(contextId, operation.id); }); return flushDeferredEnvironmentCleanup(db, Array.from(drivers.values()) .flatMap(driver => driver.flushDeferredOrphanCleanups ? [driver.flushDeferredOrphanCleanups] : [])); diff --git a/server/src/services/environment-task-runtime.ts b/server/src/services/environment-task-runtime.ts index a74aa9bf924..f7b89c69628 100644 --- a/server/src/services/environment-task-runtime.ts +++ b/server/src/services/environment-task-runtime.ts @@ -259,6 +259,10 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage runnerTaskSessionRetained:state.sessionStateAvailable === true})}::jsonb`, }).where(and(eq(environmentLeases.id, metadata.runnerTaskPreviousLeaseId), eq(environmentLeases.companyId, input.companyId))); const replacementDriverKey = metadata.driverKey; + const creationWorker = context && driver.creationOperationVersion === 1 + ? workers.getWorker(plugin.id) : undefined; + const recoverableCreation = creationWorker?.supportedMethods.includes("environmentInspectLease") === true && + metadata.providerMetadata && typeof metadata.providerMetadata === "object" && !Array.isArray(metadata.providerMetadata); let abandoned = false; const fresh = await createOwnedSandbox(db, { companyId: input.companyId, environmentId: environment!.id, heartbeatRunId: lease.heartbeatRunId, @@ -266,7 +270,8 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage pluginId: plugin.id, driverKey: metadata.driverKey, config: {}, requestKey: taskId, metadata: { driver: "plugin", pluginId: plugin.id, pluginKey: plugin.pluginKey, driverKey: metadata.driverKey, runnerTaskLease: true, agentId: run!.agentId, - ...(context ? { runnerTaskCreationContextId: context.id } : {}) }, + ...(context ? { runnerTaskCreationContextId: context.id } : {}), + ...(recoverableCreation ? { creationOperationVersion: 1, providerMetadata: metadata.providerMetadata } : {}) }, }, async identity => { const params = { driverKey: replacementDriverKey, companyId: input.companyId, environmentId: environment!.id, issueId: lease.issueId, config: driverConfig, runId: identity.runId, creationOperationId: identity.id, @@ -279,8 +284,7 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage try { return await context.wait(() => { dispatched = true; - return workers.call(plugin.id, "environmentAcquireLease", params, 30_000, undefined, settled, context.id, - async result => { + const observe = async (result: HostToWorkerMethods["environmentAcquireLease"][1]) => { if (!result.providerLeaseId || result.providerLeaseId === taskId) throw new Error("Plugin replacement task identity unavailable"); const resource = { id: identity.id, companyId: input.companyId, environmentId: environment!.id, providerLeaseId: result.providerLeaseId, providerMetadata: result.metadata }; @@ -289,7 +293,9 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage if (abandoned || !context.active || context.signal.aborted || Date.now() >= context.deadlineAtMs) { await queueSandboxCreationCleanup(db, resource); } - }); + }; + return recoverableCreation ? creationWorker!.call("environmentAcquireLease", params, 30_000, undefined, settled, context.id, observe) + : workers.call(plugin.id, "environmentAcquireLease", params, 30_000, undefined, settled, context.id, observe); }); } catch (error) { if (!dispatched) settled(); diff --git a/server/src/services/native-runtime/native-execution-context-worker.test.ts b/server/src/services/native-runtime/native-execution-context-worker.test.ts index 2c511ce108f..389d6d77bf8 100644 --- a/server/src/services/native-runtime/native-execution-context-worker.test.ts +++ b/server/src/services/native-runtime/native-execution-context-worker.test.ts @@ -59,6 +59,64 @@ async function startWorker(hostHandlers: Record = {}) { } describe("Core startup responsibility across real worker IPC", () => { + it.each([false, true])("reconstructs only a missing acquire reply while draining actual work (late reply=%s)", async lateReply => { + let releaseWrite!: () => void, writing!: () => void, releaseObservation!: () => void, observing!: () => void; + const writeGate = new Promise(resolve => { releaseWrite = resolve; }), entered = new Promise(resolve => { writing = resolve; }); + const observationGate = new Promise(resolve => { releaseObservation = resolve; }), observed = new Promise(resolve => { observing = resolve; }); + const manager = await startWorker({ "companies.get": async () => { writing(); await writeGate; return {}; } }); + const settled = vi.fn(), unrelatedSettled = vi.fn(); + const observe = vi.fn(async () => { observing(); await observationGate; }); + const recover = (operationId: string) => manager.recoverStartupAcquisition?.("reconstructed-acquire", operationId); + let draining: Promise | undefined; + try { + const acquire = manager.call(manifest.id, "environmentAcquireLease", { companyId: "company", environmentId: "environment", + driverKey: "tasks", config: { hostWrite: true }, runId: "run", agentId: "agent", adapterType: "paperclip_runner", + creationOperationId: "recovered-operation" }, 10, undefined, settled, "reconstructed-acquire", observe); + const rejected = expect(acquire).rejects.toThrow("timed out"); + await entered; await rejected; + await expect(manager.call(manifest.id, "environmentAcquireLease", { companyId: "company", environmentId: "environment", + driverKey: "tasks", config: {}, runId: "run", agentId: "agent", adapterType: "paperclip_runner", + creationOperationId: "unrelated-operation" }, 10, undefined, unrelatedSettled, "unrelated-acquire", async () => {})).rejects.toThrow("timed out"); + recover("unrelated-operation"); + vi.useFakeTimers(); + const unmatched = manager.drainStartup!("reconstructed-acquire").catch(error => error.message); + await vi.advanceTimersByTimeAsync(15_000); + expect(await unmatched).toBe("Task cleanup is still pending"); + vi.useRealTimers(); + recover("recovered-operation"); + let drained = false; + draining = manager.drainStartup!("reconstructed-acquire").then(() => { drained = true; }); + void draining.catch(() => undefined); + await new Promise(resolve => setImmediate(resolve)); + expect(drained).toBe(false); + expect(settled).not.toHaveBeenCalled(); + expect(unrelatedSettled).not.toHaveBeenCalled(); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ status: "ok" }); + let retired = false; + const retiring = manager.retireStartup!("reconstructed-acquire").then(() => { retired = true; }); + if (lateReply) { + await manager.call(manifest.id, "configChanged", { config: { release: true, duplicate: true } }); + await observed; + } + releaseWrite(); + await new Promise(resolve => setImmediate(resolve)); + if (lateReply) { expect(retired).toBe(false); expect(drained).toBe(false); releaseObservation(); } + expect(await Promise.race([draining.then(() => true), new Promise(resolve => setTimeout(() => resolve(false), 200))])).toBe(true); + await retiring; + expect(settled).toHaveBeenCalledOnce(); + expect(observe).toHaveBeenCalledTimes(lateReply ? 1 : 0); + if (!lateReply) { + expect(unrelatedSettled).not.toHaveBeenCalled(); + releaseObservation(); + await manager.call(manifest.id, "configChanged", { config: { release: true, duplicate: true } }); + } + await manager.drainStartup!("reconstructed-acquire"); + await manager.drainStartup!("unrelated-acquire"); + expect(observe).toHaveBeenCalledOnce(); + expect(unrelatedSettled).toHaveBeenCalledOnce(); + } finally { vi.useRealTimers(); releaseWrite(); releaseObservation(); await manager.stopAll(); await draining?.catch(() => undefined); } + }); + it("drops only a dead acquisition reply placeholder while retaining real writes until verified retirement", async () => { let release!: () => void, writing!: () => void; const gate = new Promise(resolve => { release = resolve; }), entered = new Promise(resolve => { writing = resolve; }); diff --git a/server/src/services/plugin-worker-manager.ts b/server/src/services/plugin-worker-manager.ts index 2fa59b62573..3a25a8dd28d 100644 --- a/server/src/services/plugin-worker-manager.ts +++ b/server/src/services/plugin-worker-manager.ts @@ -683,6 +683,7 @@ interface ExecuteLogRoute { export interface PluginWorkerHandle { hasStartupWork?(): boolean; drainStartup?(contextId: string): Promise; + recoverStartupAcquisition?(contextId: string, creationOperationId: string): void; retireStartup?(contextId: string): Promise; prepareIdleSleep?(hold: { ownerId: string; expiresAt: number }): Promise<"none" | "present" | "unknown">; releaseIdleSleep?(): void; @@ -807,6 +808,7 @@ export interface WorkerDiagnostics { */ export interface PluginWorkerManager { drainStartup?(contextId: string): Promise; + recoverStartupAcquisition?(contextId: string, creationOperationId: string): void; retireStartup?(contextId: string): Promise; /** Subscribe after a worker successfully applies live company configuration. */ onConfigApplied(listener: () => void): () => void; @@ -918,6 +920,7 @@ export function createPluginWorkerHandle( const unsettledCalls = new Map void>(); const observedResults = new Map; + creationOperationId?: string; track: () => void; observe: (result: unknown) => Promise; finish: () => void; }>(); const startupCalls = new Map void>>(); @@ -935,7 +938,7 @@ export function createPluginWorkerHandle( let timer: ReturnType | undefined; try { await Promise.race([ - Promise.all([...startupHostHandlers.get(contextId)!]), + (async () => { while (startupHostHandlers.get(contextId)?.size) await Promise.all([...startupHostHandlers.get(contextId)!]); })(), new Promise((_resolve, reject) => { timer = setTimeout(() => reject(new Error("Task cleanup is still pending")), 15_000); timer.unref?.(); }), ]); } finally { clearTimeout(timer); } @@ -1201,6 +1204,7 @@ export function createPluginWorkerHandle( const observer = observedResults.get(id); if (observer) { if (observer.response) return; + observer.track(); observer.response = response; void processObservedResult(id).catch(() => log.warn({ id }, "plugin result recording will retry during cleanup")); return; @@ -3524,17 +3528,25 @@ export function createPluginWorkerHandle( if (method === "environmentInspectLease") unsettledInspections.add(finish); } if (observeResult && startupContextId) { - let finish!: () => void; - const processing = new Promise(resolve => { finish = resolve; }); - let handlers = startupHostHandlers.get(startupContextId); - if (!handlers) { handlers = new Set(); startupHostHandlers.set(startupContextId, handlers); } - handlers.add(processing); - observedResults.set(id, { contextId: startupContextId, observe: observeResult as (result: unknown) => Promise, finish: () => { - startupHostHandlers.get(startupContextId)?.delete(processing); - if (startupHostHandlers.get(startupContextId)?.size === 0) startupHostHandlers.delete(startupContextId); - finish(); - onStartupWorkChanged?.(); - } }); + let processing: Promise | undefined, finish!: () => void; + const track = () => { + if (processing) return; + processing = new Promise(resolve => { finish = resolve; }); + let handlers = startupHostHandlers.get(startupContextId); + if (!handlers) { handlers = new Set(); startupHostHandlers.set(startupContextId, handlers); } + handlers.add(processing); + }; + track(); + observedResults.set(id, { contextId: startupContextId, + creationOperationId: method === "environmentAcquireLease" && isRecord(params) ? readNonEmptyString((params as Record).creationOperationId) ?? undefined : undefined, + track, observe: observeResult as (result: unknown) => Promise, finish: () => { + if (!processing) return; + startupHostHandlers.get(startupContextId)?.delete(processing); + if (startupHostHandlers.get(startupContextId)?.size === 0) startupHostHandlers.delete(startupContextId); + processing = undefined; + finish(); + onStartupWorkChanged?.(); + } }); } const timeout = resolveRpcCallTimeoutMs(timeoutMs, rpcTimeoutMs); const invocationScope = deriveInvocationScope(method, params); @@ -3689,6 +3701,11 @@ export function createPluginWorkerHandle( await startInternal(); }, hasStartupWork() { return startupCalls.size > 0 || startupHostHandlers.size > 0 || startupRetirements.size > 0; }, + recoverStartupAcquisition(contextId: string, creationOperationId: string) { + for (const observer of observedResults.values()) { + if (observer.contextId === contextId && observer.creationOperationId === creationOperationId && !observer.response) observer.finish(); + } + }, async drainStartup(contextId: string) { revokeStartupInvocations(contextId); for (const [id, observer] of observedResults) if (observer.contextId === contextId) void processObservedResult(id).catch(() => undefined); @@ -4043,6 +4060,9 @@ export function createPluginWorkerManager( async retireStartup(contextId: string) { await Promise.all([...startupWorkers()].map(worker => worker.retireStartup?.(contextId))); }, + recoverStartupAcquisition(contextId: string, creationOperationId: string) { + for (const worker of startupWorkers()) worker.recoverStartupAcquisition?.(contextId, creationOperationId); + }, async drainStartup(contextId: string) { await Promise.all([...startupWorkers()].map(worker => worker.drainStartup?.(contextId))); }, diff --git a/server/src/services/sandbox-creation-operations.ts b/server/src/services/sandbox-creation-operations.ts index 7fab2432fe6..d17300c7d4e 100644 --- a/server/src/services/sandbox-creation-operations.ts +++ b/server/src/services/sandbox-creation-operations.ts @@ -10,6 +10,13 @@ import { notifyEnvironmentCleanup, writeEnvironmentCleanup } from "./environment export type SandboxCreationOperation = typeof sandboxCreationOperations.$inferSelect; +export function recoverableTaskCreation(operation: Pick): boolean { + const metadata = operation.metadata; + return metadata.runnerTaskLease === true && metadata.creationOperationVersion === 1 && + typeof metadata.runnerTaskCreationContextId === "string" && metadata.runnerTaskCreationContextId.length > 0 && + metadata.providerMetadata !== null && typeof metadata.providerMetadata === "object" && !Array.isArray(metadata.providerMetadata); +} + type CreationOwner = { companyId: string; environmentId: string; @@ -70,8 +77,22 @@ export async function createOwnedSandbox( return { providerLeaseId: existing.providerLeaseId, expiresAt: existing.expiresAt?.toISOString(), metadata: { ...existing.metadata, creationOperationId: operation.id } }; } - const lease = await create({ id, runId }); - return { ...lease, metadata: { ...lease.metadata, creationOperationId: id } }; + try { + const lease = await create({ id, runId }); + return { ...lease, metadata: { ...lease.metadata, creationOperationId: id } }; + } catch (error) { + if (recoverableTaskCreation(operation)) { + try { + await writeEnvironmentCleanup(db, tx => tx.update(sandboxCreationOperations).set({ nextAttemptAt: new Date(), updatedAt: new Date() }) + .where(and(eq(sandboxCreationOperations.id, id), recoverableSandboxCreationCondition())), + { companyId: operation.companyId, runId: operation.runId }); + } catch { + logger.warn({ errorKind: "task_creation_recovery_wakeup_failed", operationId: id, companyId: operation.companyId }, + "Task creation recovery retains its saved deadline"); + } + } + throw error; + } } /** Keep an already returned resource identity even when task attachment fails. @@ -129,6 +150,7 @@ export async function queueSandboxCreationCleanup(db: Db, input: { export async function recoverSandboxCreations(db: Db, inspect: (operation: SandboxCreationOperation) => Promise, + onTaskRecovered?: (operation: SandboxCreationOperation) => void | Promise, ): Promise { const due = await db.select().from(sandboxCreationOperations).where(and( recoverableSandboxCreationCondition(), lte(sandboxCreationOperations.nextAttemptAt, new Date()), @@ -139,10 +161,25 @@ export async function recoverSandboxCreations(db: Db, typeof candidate.metadata.creationProviderLeaseId === "string") { await queueSandboxCreationCleanup(db, { id: candidate.id, companyId: candidate.companyId, environmentId: candidate.originalEnvironmentId, providerLeaseId: candidate.metadata.creationProviderLeaseId }); + if (recoverableTaskCreation(candidate)) await onTaskRecovered?.(candidate); continue; } const [existingCompany] = await db.select({ id: companies.id }).from(companies).where(eq(companies.id, candidate.companyId)).limit(1); const observation = existingCompany ? await inspect(candidate) : null; + if (existingCompany && recoverableTaskCreation(candidate) && observation?.creationLease) { + const lease = observation.creationLease; + if (typeof lease.providerLeaseId !== "string" || lease.providerLeaseId.length === 0 || lease.providerLeaseId !== observation.providerLeaseId || + !lease.metadata || typeof lease.metadata !== "object" || Array.isArray(lease.metadata)) { + throw new Error("Task creation recovery returned foreign ownership"); + } + const identity = { id: candidate.id, companyId: candidate.companyId, environmentId: candidate.originalEnvironmentId, + providerLeaseId: lease.providerLeaseId, providerMetadata: lease.metadata }; + await recordSandboxCreationResource(db, identity); + await queueSandboxCreationCleanup(db, identity); + await onTaskRecovered?.(candidate); + continue; + } + if (existingCompany && recoverableTaskCreation(candidate)) throw new Error("Task creation identity is not available yet"); const cleanup = observation?.creationCleanup ? readEnvironmentCreationCleanupError({ data: { schema: "paperclip/environment-creation-cleanup/v1", cleanup: observation.creationCleanup, } }) : null;