From 03a25996cfcb4b78241232974270f0e744dc24cb Mon Sep 17 00:00:00 2001 From: Dotta Date: Sat, 12 Sep 2026 21:18:56 -0500 Subject: [PATCH] fix: stop stalled sandbox startup before draining remote commands Co-Authored-By: Paperclip --- doc/connections/AI-CONNECTIONS.md | 14 ++ .../src/acpx-engine/execute.test.ts | 57 ++++++++ .../adapter-utils/src/acpx-engine/execute.ts | 40 +++--- .../acpx-engine/startup-cancellation.test.ts | 124 ++++++++++++++++++ .../src/acpx-engine/startup-cancellation.ts | 88 +++++++++++++ packages/adapter-utils/src/types.ts | 3 + .../daytona/src/plugin.test.ts | 46 +++++++ .../sandbox-providers/daytona/src/plugin.ts | 21 ++- packages/plugins/sdk/src/protocol.ts | 3 + .../environment-run-orchestrator.test.ts | 14 ++ .../src/__tests__/environment-runtime.test.ts | 29 ++++ .../services/environment-run-orchestrator.ts | 3 + server/src/services/environment-runtime.ts | 6 + server/src/services/heartbeat.ts | 20 +++ 14 files changed, 450 insertions(+), 18 deletions(-) create mode 100644 packages/adapter-utils/src/acpx-engine/startup-cancellation.test.ts create mode 100644 packages/adapter-utils/src/acpx-engine/startup-cancellation.ts diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 3db5d77a08..4a0cb77adf 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -127,6 +127,20 @@ process may already hold credentials. The revoke confirmation lists attributed active runs and exposes the existing Stop action; it does not promise immediate provider-side revocation. +### Stop during sandbox preparation + +ACPX startup registers cancellation while it materializes the remote auth home +and stages files. Stop requests termination of that run's sandbox. Daytona closes +admission and stops the sandbox before waiting for outstanding setup commands. +The host requires a receipt for the exact company, run, and provider lease before +abandoning the blocked setup RPC. Normal completion still drains work gracefully. + +Late setup responses cannot launch the agent. A cancelled sandbox cannot resume +while its old provider requests are still settling; a retry receives an explicit +error instead. If termination cannot be verified, the adapter keeps ownership +until the outstanding operation settles, and Stop is not acknowledged as complete. +Local execution and cancellation of an already-running agent turn are unchanged. + ## Legacy adoption Migration `0273` indexes only explicitly owned personal secrets with a recognized diff --git a/packages/adapter-utils/src/acpx-engine/execute.test.ts b/packages/adapter-utils/src/acpx-engine/execute.test.ts index a91503b9c9..ec080e174a 100644 --- a/packages/adapter-utils/src/acpx-engine/execute.test.ts +++ b/packages/adapter-utils/src/acpx-engine/execute.test.ts @@ -3793,6 +3793,63 @@ describe("ACPX engine remote session-lifecycle re-staging (PR 3: stage once / re }; } + it("cancels a stalled remote startup without waiting for the command or launching the provider", async () => { + const { stateDir, localCwd, executionTarget } = await setupRemoteSandbox(); + const controller = new AbortController(); + let entered!: () => void; + const started = new Promise(resolve => { entered = resolve; }); + let release!: () => void; + const blocked = new Promise(resolve => { release = resolve; }); + const stopRemoteStartup = vi.fn(async () => {}); + const createRuntime = vi.fn(() => recordingRuntime({ ensureInputs: [] }) as never); + const originalExecute = executionTarget.runner.execute; + executionTarget.runner.execute = async input => { + if (input.command === "stalled-auth-setup") { + entered(); + await blocked; + return { exitCode: 0, signal: null, timedOut: false, stdout: "", stderr: "", pid: null, startedAt: new Date().toISOString() }; + } + return originalExecute(input); + }; + const execute = createAcpxEngineExecutor({ + createRuntime, + prepareRemoteManagedHome: async input => { + const target = input.executionTarget; + if (target?.kind !== "remote" || target.transport !== "sandbox" || !target.runner) { + throw new Error("Expected sandbox target"); + } + await target.runner.execute({ command: "stalled-auth-setup" }); + return { stagedRuntime: await input.stage([]) }; + }, + }); + let settled = false; + const run = execute({ + runId: "cancel-startup", runtime: {}, + ...baseExecuteArgs({ stateDir, localCwd, executionTarget }), + signal: controller.signal, stopRemoteStartup, + } as never).catch(error => error).finally(() => { settled = true; }); + await started; + controller.abort(new Error("Stopped by user")); + try { + await vi.waitFor(() => expect(stopRemoteStartup).toHaveBeenCalledTimes(1), { timeout: 500 }); + await vi.waitFor(() => expect(settled).toBe(true), { timeout: 500 }); + expect(createRuntime).not.toHaveBeenCalled(); + } finally { + release(); + await run; + } + // A provider RPC completing late must not resume the cancelled startup. + expect(createRuntime).not.toHaveBeenCalled(); + // A subsequent run can acquire the same local staging/auth preparation + // resources; cancellation must not leave the staging lease held. + const retry = await execute({ + runId: "retry-after-cancel", runtime: {}, + ...baseExecuteArgs({ stateDir, localCwd, executionTarget }), + } as never); + expect(retry.exitCode).toBe(0); + expect(createRuntime).toHaveBeenCalledOnce(); + }); + it("test_acp_resume_compatible_session_does_not_restage", async () => { const { stateDir, localCwd, remoteCwd, executionTarget } = await setupRemoteSandbox(); const ensureInputs: Array> = []; diff --git a/packages/adapter-utils/src/acpx-engine/execute.ts b/packages/adapter-utils/src/acpx-engine/execute.ts index e61bf94ae2..580e19c0f7 100644 --- a/packages/adapter-utils/src/acpx-engine/execute.ts +++ b/packages/adapter-utils/src/acpx-engine/execute.ts @@ -1,3 +1,4 @@ +import { cancellableSandboxStartup } from "./startup-cancellation.js"; import fs from "node:fs/promises"; import fsSync from "node:fs"; import os from "node:os"; @@ -4115,23 +4116,28 @@ export function createAcpxEngineExecutor(deps: AcpxEngineExecutorOptions = {}) { // run inside this wrap and overrides the store, so an in-step exec still // parents to its step span. On a local or SSH target // `spanParent.parentContext` is a no-op token, so the wrap is inert. - prepared = await runWithRuntimeParent(spanParent.parentContext, () => - buildRuntime({ - ctx, - engine, - deps, - ledger: runResourceLedger, - stagedIdleMs: warmIdleMs, - spanParent, - getRuntimeParentContext, - runtimeSpan: runRuntimeSpan, - stageRuntimeSpan: runStageSpan, - }), - ); - buildRuntimeSettled = true; - // Capture the run's staging lease release now that the runtime built. The - // run root `finally` releases it as the final settlement act. - releaseStagingLease = prepared.sessionStagingLeaseRelease; + const startupCancellation = cancellableSandboxStartup(ctx); + try { + prepared = await runWithRuntimeParent(spanParent.parentContext, () => + buildRuntime({ + ctx: startupCancellation.context, + engine, + deps, + ledger: runResourceLedger, + stagedIdleMs: warmIdleMs, + spanParent, + getRuntimeParentContext, + runtimeSpan: runRuntimeSpan, + stageRuntimeSpan: runStageSpan, + }), + ); + buildRuntimeSettled = true; + // Capture acquired resources before the cancellation boundary so the + // normal settlement path also releases a just-completed build. + releaseStagingLease = prepared.sessionStagingLeaseRelease; + } finally { + await startupCancellation.finish(); + } // Per-project staging outcomes for the referenced (mentioned) projects, surfaced back to the // server on the run result. A referenced project that failed to stage into the sandbox is a // first-class, counted failure in the requested-vs-synced observability, not only a warning. The diff --git a/packages/adapter-utils/src/acpx-engine/startup-cancellation.test.ts b/packages/adapter-utils/src/acpx-engine/startup-cancellation.test.ts new file mode 100644 index 0000000000..b02bbe59b6 --- /dev/null +++ b/packages/adapter-utils/src/acpx-engine/startup-cancellation.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, it, vi } from "vitest"; +import type { AdapterExecutionContext } from "../types.js"; +import type { CommandManagedDuplexChannel, CommandManagedRuntimeRunner } from "../command-managed-runtime.js"; +import { cancellableSandboxStartup } from "./startup-cancellation.js"; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((yes, no) => { resolve = yes; reject = no; }); + return { promise, resolve, reject }; +} +const result = { exitCode: 0, signal: null, timedOut: false, stdout: "", stderr: "" }; +function fixture(stop = vi.fn(async () => {}), extra: Partial = {}) { + const controller = new AbortController(); + const execute = vi.fn(async () => result); + const ctx = { + signal: controller.signal, stopRemoteStartup: stop, + executionTarget: { kind: "remote", transport: "sandbox", providerKey: "daytona", runner: { execute, ...extra } }, + } as unknown as AdapterExecutionContext; + const startup = cancellableSandboxStartup(ctx); + const runner = (startup.context.executionTarget as { runner: CommandManagedRuntimeRunner }).runner; + return { controller, execute, stop, startup, runner }; +} + +describe("sandbox startup cancellation boundary", () => { + it("keeps ownership until termination is confirmed, including an early command error", async () => { + const receipt = deferred(); + const command = deferred(); + const f = fixture(vi.fn(() => receipt.promise)); + f.execute.mockReturnValue(command.promise); + let settled = false; + const outcome = f.runner.execute({ command: "setup" }).catch(error => error).finally(() => { settled = true; }); + await vi.waitFor(() => expect(f.execute).toHaveBeenCalledOnce()); + f.controller.abort(new Error("Stopped")); + command.reject(new Error("socket closed")); + await vi.waitFor(() => expect(f.stop).toHaveBeenCalledOnce()); + expect(settled).toBe(false); + receipt.resolve(); + expect(await outcome).toEqual(new Error("Stopped")); + await expect(f.startup.finish()).rejects.toThrow("Stopped"); + await expect(f.runner.execute({ command: "late setup" })).rejects.toThrow("Stopped"); + expect(f.execute).toHaveBeenCalledOnce(); + }); + + it("does not abandon a hung command when the provider cannot confirm stop", async () => { + const command = deferred(); + const f = fixture(vi.fn(async () => { throw new Error("stop unverified"); })); + f.execute.mockReturnValue(command.promise); + let settled = false; + const outcome = f.runner.execute({ command: "setup" }).catch(error => error).finally(() => { settled = true; }); + await vi.waitFor(() => expect(f.execute).toHaveBeenCalledOnce()); + f.controller.abort(); + await vi.waitFor(() => expect(f.stop).toHaveBeenCalledOnce()); + expect(settled).toBe(false); + command.resolve(result); + expect(await outcome).toEqual(new Error("stop unverified")); + await expect(f.startup.finish()).rejects.toThrow("stop unverified"); + }); + + it("waits for all parallel setup requests after an unverified stop", async () => { + const first = deferred(); + const second = deferred(); + const f = fixture(vi.fn(async () => { throw new Error("stop unverified"); })); + f.execute.mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise); + let settled = false; + const a = f.runner.execute({ command: "first" }).catch(error => error).finally(() => { settled = true; }); + const b = f.runner.execute({ command: "second" }).catch(error => error); + await vi.waitFor(() => expect(f.execute).toHaveBeenCalledTimes(2)); + f.controller.abort(); + first.reject(new Error("first failed")); + await vi.waitFor(() => expect(f.stop).toHaveBeenCalledOnce()); + expect(settled).toBe(false); + second.resolve(result); + expect(await a).toEqual(new Error("stop unverified")); + expect(await b).toEqual(new Error("stop unverified")); + await expect(f.startup.finish()).rejects.toThrow("stop unverified"); + }); + + it("disarms after successful startup and leaves ordinary turn cancellation alone", async () => { + const f = fixture(); + await f.runner.execute({ command: "setup" }); + await f.startup.finish(); + f.controller.abort(); + await expect(f.runner.execute({ command: "normal turn cleanup" })).resolves.toEqual(result); + expect(f.stop).not.toHaveBeenCalled(); + expect(f.execute).toHaveBeenCalledTimes(2); + }); + + it("never starts a command when cancellation arrives before setup", async () => { + const f = fixture(); + f.controller.abort(new Error("Stopped")); + await expect(f.runner.execute({ command: "setup" })).rejects.toThrow("Stopped"); + await expect(f.startup.finish()).rejects.toThrow("Stopped"); + expect(f.execute).not.toHaveBeenCalled(); + expect(f.stop).toHaveBeenCalledOnce(); + }); + + it("closes a duplex route returned after startup was cancelled", async () => { + const opening = deferred(); + const channel = { + write: vi.fn(), onData: vi.fn(), onExit: vi.fn(), stop: vi.fn(), + close: vi.fn(async () => {}), + }; + const openDuplexChannel = vi.fn(() => opening.promise); + const f = fixture(vi.fn(async () => {}), { openDuplexChannel }); + const outcome = f.runner.openDuplexChannel!({ command: ["agent"] }).catch(error => error); + await vi.waitFor(() => expect(openDuplexChannel).toHaveBeenCalledOnce()); + f.controller.abort(new Error("Stopped")); + expect(await outcome).toEqual(new Error("Stopped")); + opening.resolve(channel); + await vi.waitFor(() => expect(channel.close).toHaveBeenCalledOnce()); + await expect(f.startup.finish()).rejects.toThrow("Stopped"); + }); + + it("does not stop another run's runner", async () => { + const a = fixture(); + const b = fixture(); + a.controller.abort(new Error("Stopped")); + await expect(a.startup.finish()).rejects.toThrow("Stopped"); + await expect(b.runner.execute({ command: "setup" })).resolves.toEqual(result); + await b.startup.finish(); + expect(b.stop).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/adapter-utils/src/acpx-engine/startup-cancellation.ts b/packages/adapter-utils/src/acpx-engine/startup-cancellation.ts new file mode 100644 index 0000000000..ff41468a50 --- /dev/null +++ b/packages/adapter-utils/src/acpx-engine/startup-cancellation.ts @@ -0,0 +1,88 @@ +import type { AdapterExecutionContext } from "../types.js"; +import type { CommandManagedRuntimeRunner } from "../command-managed-runtime.js"; + +/** Stop owns the sandbox before abandoning an outstanding startup RPC. The + * original RPC may complete late; its result cannot restart preparation. */ +export function cancellableSandboxStartup(ctx: AdapterExecutionContext) { + const target = ctx.executionTarget; + const signal = ctx.signal; + const stop = ctx.stopRemoteStartup; + if (!signal || !stop || target?.kind !== "remote" || target.transport !== "sandbox" || !target.runner) { + return { context: ctx, finish: async () => {} }; + } + let armed = true; + let stopping: Promise | undefined; + const inFlight = new Set>(); + let rejectStopped!: (error: unknown) => void; + const stopped = new Promise((_, reject) => { rejectStopped = reject; }); + // Cancellation can arrive between RPCs, when nobody is awaiting this yet. + void stopped.catch(() => {}); + const onAbort = () => { + if (stopping) return; + stopping = Promise.resolve().then(stop); + void stopping.then( + () => rejectStopped(signal.reason ?? new Error("Stopped during sandbox startup")), + // Without proof, the original operation still owns its resources. Do + // not abandon it or release credentials while it could be running. + () => {}, + ); + }; + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) onAbort(); + const joinStop = async () => { + try { + await stopping; + } catch (error) { + // Parallel setup may fail fast. Without a receipt, retain ownership of + // every outstanding RPC, not only the first one that returns an error. + await Promise.allSettled([...inFlight]); + throw error; + } + }; + const guard = async (call: () => Promise): Promise => { + if (!armed && !stopping) return call(); + if (stopping) { + await joinStop(); + return stopped; + } + const operation = Promise.resolve().then(() => { + if (stopping) signal.throwIfAborted(); + return call(); + }); + inFlight.add(operation); + void operation.then(() => inFlight.delete(operation), () => inFlight.delete(operation)); + try { + return await Promise.race([operation, stopped]); + } finally { + if (stopping) { + await joinStop(); + signal.throwIfAborted(); + } + } + }; + const original = target.runner; + const runner: CommandManagedRuntimeRunner = { + ...original, + execute: input => guard(() => original.execute(input)), + ...(original.syncIn ? { syncIn: (input: Parameters>[0]) => guard(() => original.syncIn!(input)) } : {}), + ...(original.syncOut ? { syncOut: (input: Parameters>[0]) => guard(() => original.syncOut!(input)) } : {}), + ...(original.openDuplexChannel ? { + openDuplexChannel: (input: Parameters>[0]) => guard(async () => { + const channel = await original.openDuplexChannel!(input); + // An open RPC can return after confirmed termination. Its host route + // must not survive just because the caller already abandoned the RPC. + if (stopping) await channel.close(); + return channel; + }), + } : {}), + }; + return { + context: { ...ctx, executionTarget: { ...target, runner } }, + async finish() { + signal.removeEventListener("abort", onAbort); + armed = false; + if (stopping) await joinStop(); + signal.throwIfAborted(); + }, + }; +} diff --git a/packages/adapter-utils/src/types.ts b/packages/adapter-utils/src/types.ts index 7540874d71..4a166341b2 100644 --- a/packages/adapter-utils/src/types.ts +++ b/packages/adapter-utils/src/types.ts @@ -199,6 +199,9 @@ export interface AdapterExecutionContext { signal?: AbortSignal; /** Opt in to signal-based cancellation before starting provider work. */ onCancellationReady?: () => Promise; + /** Host-owned stop of this run's sandbox during setup. Resolves only after + * provider termination is verified; never accepts an agent-selected lease. */ + stopRemoteStartup?: () => Promise; /** Server-owned, actor-attributed snapshot also rendered by legacy wake prompts. */ executionContinuation?: ExecutionContinuationEnvelope | null; runId: string; diff --git a/packages/plugins/sandbox-providers/daytona/src/plugin.test.ts b/packages/plugins/sandbox-providers/daytona/src/plugin.test.ts index 4a08bef53b..cfa47b467f 100644 --- a/packages/plugins/sandbox-providers/daytona/src/plugin.test.ts +++ b/packages/plugins/sandbox-providers/daytona/src/plugin.test.ts @@ -2645,6 +2645,52 @@ describe("Daytona sandbox provider plugin", () => { expect(sandbox.process.executeCommand).toHaveBeenCalledTimes(1); }); + it("cancels active work before waiting for a stalled execute to drain", async () => { + process.env.DAYTONA_API_KEY = "host-key"; + const sandbox = createMockSandbox({ id: "lease-a" }); + let release!: () => void; + sandbox.process.executeCommand.mockImplementation(async () => { + await new Promise(resolve => { release = resolve; }); + return { exitCode: 0, result: "", artifacts: { stdout: "" } }; + }); + mockGet.mockResolvedValue(sandbox); + const execute = plugin.definition.onEnvironmentExecute!(execParams("lease-a")); + await vi.waitFor(() => expect(release).toBeTypeOf("function")); + const cancellation = plugin.definition.onEnvironmentReleaseLease!({ + driverKey: "daytona", companyId: "company-1", environmentId: "env-1", + providerLeaseId: "lease-a", config: { timeoutMs: 300000, reuseLease: true }, + cancelActiveWork: true, + }); + try { + await vi.waitFor(() => expect(sandbox.stop).toHaveBeenCalledTimes(1), { timeout: 500 }); + await expect(cancellation).resolves.toEqual({ providerLeaseId: "lease-a", state: "stopped" }); + await expect(plugin.definition.onEnvironmentExecute!(execParams("lease-a"))).rejects.toThrow(/no longer active/); + await expect(plugin.definition.onEnvironmentResumeLease!({ + driverKey: "daytona", companyId: "company-1", environmentId: "env-1", + providerLeaseId: "lease-a", config: { timeoutMs: 300000, reuseLease: true }, + })).rejects.toThrow(/still settling cancelled work/); + expect(sandbox.start).not.toHaveBeenCalled(); + } finally { + release(); + await execute; + await cancellation; + } + }); + + it("does not report termination when the provider rejects Stop", async () => { + process.env.DAYTONA_API_KEY = "host-key"; + const sandbox = createMockSandbox({ id: "lease-a" }); + sandbox.stop.mockRejectedValue(new Error("provider unavailable")); + mockGet.mockResolvedValue(sandbox); + await expect(plugin.definition.onEnvironmentReleaseLease!({ + driverKey: "daytona", companyId: "company-1", environmentId: "env-1", + providerLeaseId: "lease-a", config: { timeoutMs: 300000, reuseLease: true }, + cancelActiveWork: true, + })).rejects.toThrow("provider unavailable"); + expect(sandbox.delete).not.toHaveBeenCalled(); + await expect(plugin.definition.onEnvironmentExecute!(execParams("lease-a"))).rejects.toThrow(/no longer active/); + }); + it("waits for an in-flight execute before teardown cleanup starts", async () => { process.env.DAYTONA_API_KEY = "host-key"; const sandbox = createMockSandbox({ id: "lease-a" }); diff --git a/packages/plugins/sandbox-providers/daytona/src/plugin.ts b/packages/plugins/sandbox-providers/daytona/src/plugin.ts index cc202abee1..21fc96e429 100644 --- a/packages/plugins/sandbox-providers/daytona/src/plugin.ts +++ b/packages/plugins/sandbox-providers/daytona/src/plugin.ts @@ -1126,7 +1126,11 @@ const sandboxHandleActivityGates = (() => { gates.clear(); } - return { begin, waitForIdle, end, reset }; + function isActive(scope: SandboxScope): boolean { + return gates.has(sandboxHandleCacheKey(scope)); + } + + return { begin, waitForIdle, isActive, end, reset }; })(); type SandboxLeaseAdmissionOptions = { @@ -2187,6 +2191,11 @@ const plugin = definePlugin({ providerLeaseId: params.providerLeaseId, config, }; + // A confirmed stop may precede completion of an old SDK request. Do not + // restart its sandbox under that request; it could still write or execute. + if (sandboxHandleLeaseAdmissionStates.isClosed(scope) && sandboxHandleActivityGates.isActive(scope)) { + throw new Error("The stopped Daytona sandbox is still settling cancelled work. Retry shortly."); + } return await withSandboxActivityGate(scope, async () => { const sandbox = await getSandboxOrNull(scope, { bypassTeardownGate: true }); if (!sandbox) { @@ -2283,6 +2292,16 @@ const plugin = definePlugin({ if (!sandbox) return { providerLeaseId: params.providerLeaseId, state: "destroyed" }; evictSandboxHandle(scope); + if (params.cancelActiveWork) { + // Graceful release waits for activity below. Stop must not wait on the + // command it is cancelling. Admission is already closed for this lease. + const timeoutSeconds = Math.min(toTimeoutSeconds(config.timeoutMs), 30); + await withLivenessTimeout("sandbox.refreshData", Math.min(config.livenessTimeoutMs, 30_000), () => sandbox.refreshData()); + if (sandbox.state !== "stopped") await sandbox.stop(timeoutSeconds); + sandboxHandleSessionStore.clear(scope); + await closeDaytonaDuplexChannelsForLease(params.providerLeaseId); + return { providerLeaseId: params.providerLeaseId, state: "stopped" }; + } await sandboxHandleActivityGates.waitForIdle(scope); await teardownSession(sandbox, scope); // Close every duplex channel on this lease before the stop or the delete, diff --git a/packages/plugins/sdk/src/protocol.ts b/packages/plugins/sdk/src/protocol.ts index 36c1189fb7..a599a8fc7e 100644 --- a/packages/plugins/sdk/src/protocol.ts +++ b/packages/plugins/sdk/src/protocol.ts @@ -662,6 +662,9 @@ export interface PluginEnvironmentResumeLeaseParams extends PluginEnvironmentDri } export interface PluginEnvironmentReleaseLeaseParams extends PluginEnvironmentDriverBaseParams { + /** Explicit operator cancellation: terminate active work instead of waiting + * for command/sync activity to drain. Still requires a provider receipt. */ + cancelActiveWork?: boolean; providerLeaseId: string | null; leaseMetadata?: Record; } diff --git a/server/src/__tests__/environment-run-orchestrator.test.ts b/server/src/__tests__/environment-run-orchestrator.test.ts index 01bb0442a9..f3d1fc659c 100644 --- a/server/src/__tests__/environment-run-orchestrator.test.ts +++ b/server/src/__tests__/environment-run-orchestrator.test.ts @@ -255,6 +255,20 @@ describe("environmentRunOrchestrator — realizeForRun", () => { mockLogActivity.mockResolvedValue(undefined); }); + it.each([false, true])("only requests active-work cancellation for explicit Stop: %s", async (cancelActiveWork) => { + const releaseRunLeases = vi.fn().mockResolvedValue([]); + const runtime = makeMockRuntime({ releaseRunLeases }); + const orchestrator = environmentRunOrchestrator(mockDb, { environmentRuntime: runtime }); + await orchestrator.releaseForRun({ + heartbeatRunId: "run-1", companyId: "company-1", agentId: "agent-1", + providerResourceDisposition: "stop_and_retain", cancelActiveWork, + }); + expect(releaseRunLeases).toHaveBeenCalledWith( + "run-1", "released", expect.any(Function), "stop_and_retain", + ...(cancelActiveWork ? [true] : []), + ); + }); + it("happy path: returns lease, executionTarget, and remoteExecution on successful realization", async () => { const executionTarget = { kind: "local", environmentId: "env-1", leaseId: "lease-1" }; const remoteExecution = { kind: "local", environmentId: "env-1", leaseId: "lease-1" }; diff --git a/server/src/__tests__/environment-runtime.test.ts b/server/src/__tests__/environment-runtime.test.ts index 25e78070e5..9782336280 100644 --- a/server/src/__tests__/environment-runtime.test.ts +++ b/server/src/__tests__/environment-runtime.test.ts @@ -39,6 +39,7 @@ import * as sandboxProviderRuntime from "../services/sandbox-provider-runtime.ts import * as environmentsModule from "../services/environments.ts"; import { logger } from "../middleware/logger.ts"; import { environmentService } from "../services/environments.ts"; +import { remoteExecutionHasStopped } from "../services/remote-execution-termination.ts"; import { heartbeatService } from "../services/heartbeat.ts"; import { secretService } from "../services/secrets.ts"; import type { PluginWorkerManager } from "../services/plugin-worker-manager.ts"; @@ -463,6 +464,34 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { return { pluginId, companyId, agentId, environment, runId, executionWorkspaceId, reusableLease }; } + it.each(["stopped", "missing", "wrong-lease", "error"])( + "startup cancellation is run-scoped and needs a matching receipt: %s", + async (outcome) => { + const { pluginId, companyId, runId, reusableLease, environment } = await seedReusablePluginSandboxLease(); + const other = await environmentService(db).acquireLease({ + companyId, environmentId: environment.id, heartbeatRunId: null, + leasePolicy: "ephemeral", provider: "fake-plugin", providerLeaseId: "other-sandbox", + }); + const call = vi.fn(async () => { + if (outcome === "error") throw new Error("provider unavailable"); + if (outcome === "missing") return undefined; + return { providerLeaseId: outcome === "wrong-lease" ? "other-sandbox" : reusableLease.providerLeaseId, state: "stopped" }; + }); + const workerManager = { + isRunning: () => true, call, + getWorker: () => ({ supportedMethods: ["environmentReleaseLease"] }), + } as unknown as PluginWorkerManager; + const runtimeWithPlugin = environmentRuntimeService(db, { pluginWorkerManager: workerManager }); + await runtimeWithPlugin.releaseRunLeases(runId, "released", undefined, "stop_and_retain", true); + expect(call).toHaveBeenCalledWith(pluginId, "environmentReleaseLease", expect.objectContaining({ + companyId, providerLeaseId: reusableLease.providerLeaseId, cancelActiveWork: true, + }), expect.any(Number)); + expect(call).toHaveBeenCalledOnce(); + expect(await remoteExecutionHasStopped(db, companyId, runId)).toBe(outcome === "stopped"); + await expect(environmentService(db).getLeaseById(other.id)).resolves.toMatchObject({ status: "active" }); + }, + ); + it("retains a successful reusable sandbox lease without stopping the provider resource", async () => { const { pluginId, runId, reusableLease } = await seedReusablePluginSandboxLease(); const workerManager = { diff --git a/server/src/services/environment-run-orchestrator.ts b/server/src/services/environment-run-orchestrator.ts index 0c08deb714..97cca667e6 100644 --- a/server/src/services/environment-run-orchestrator.ts +++ b/server/src/services/environment-run-orchestrator.ts @@ -584,6 +584,8 @@ export function environmentRunOrchestrator( agentId: string; status?: Extract; failureReason?: string; + /** Explicit Stop during adapter startup; never used for ordinary cleanup. */ + cancelActiveWork?: boolean; /** Explicit paperclip_runner resource lifecycle. Omitted for legacy adapters. */ providerResourceDisposition?: ProviderResourceDisposition; nativeLifecycleTelemetry?: { @@ -606,6 +608,7 @@ export function environmentRunOrchestrator( status, (leaseId, error) => result.errors.push({ leaseId, error }), input.providerResourceDisposition, + ...(input.cancelActiveWork ? [true] as const : []), ); } catch (err) { result.errors.push({ leaseId: "*", error: err }); diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index c884c99da3..2deb410986 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -477,6 +477,8 @@ export interface EnvironmentDriverAcquireInput { } export interface EnvironmentDriverReleaseInput { + /** Explicit Stop may terminate in-flight setup rather than drain it. */ + cancelActiveWork?: boolean; environment: Environment; lease: EnvironmentLease; status: Extract; @@ -3032,9 +3034,11 @@ function createSandboxEnvironmentDriver( config: stripSandboxProviderEnvelope(config as SandboxEnvironmentConfig), providerLeaseId: input.lease.providerLeaseId, leaseMetadata: metadata, + ...(input.cancelActiveWork ? { cancelActiveWork: true } : {}), }, resolvePluginSandboxRpcTimeoutMs(stripSandboxProviderEnvelope(config as SandboxEnvironmentConfig))), ); termination = remoteTerminationReceipt(input.lease, receipt); + if (input.cancelActiveWork && !termination) cleanupStatus = "failed"; } catch { cleanupStatus = "failed"; } @@ -3699,6 +3703,7 @@ export function environmentRuntimeService( status: Extract = "released", onLeaseReleaseError?: (leaseId: string, error: unknown) => void, providerResourceDisposition?: ProviderResourceDisposition, + cancelActiveWork?: boolean, ): Promise { const leaseRows = await db .select() @@ -3781,6 +3786,7 @@ export function environmentRuntimeService( }) : driver ? await driver.releaseRunLease({ + ...(cancelActiveWork ? { cancelActiveWork: true } : {}), environment, lease: leaseSnapshot, // A stopped reusable provider resource must remain eligible diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 9c682d24a5..7bcb8e2d2a 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -23692,6 +23692,26 @@ export function heartbeatService( }, onDispatch: markDispatchStarted, signal: executionControl.controller.signal, + ...(executionTarget?.kind === "remote" && executionTarget.transport === "sandbox" ? { + stopRemoteStartup: async () => { + // Scope comes from the running host invocation, never agent + // config. Keep adapter ownership until setup has unwound. + if (!executionControl.controller.signal.aborted) { + throw new Error("Remote startup stop requires a cancelled run"); + } + const release = await envOrchestrator.releaseForRun({ + heartbeatRunId: run.id, + companyId: agent.companyId, + agentId: agent.id, + status: "released", + providerResourceDisposition: "stop_and_retain", + cancelActiveWork: true, + }); + if (release.errors.length || !await remoteExecutionHasStopped(db, agent.companyId, run.id)) { + throw new Error("Could not verify remote startup stopped"); + } + }, + } : {}), onCancellationReady: async () => { await registerAdapterExecutionControl(run.id, executionControl); const current = await getRun(run.id);