diff --git a/packages/paperclip-runner/src/index.ts b/packages/paperclip-runner/src/index.ts index 72e5aa5eb1e..665c5a4a0af 100644 --- a/packages/paperclip-runner/src/index.ts +++ b/packages/paperclip-runner/src/index.ts @@ -82,6 +82,7 @@ export { defaultCapabilityRunnerdBinary, readRunnerdArtifactBinding, drainRetainedRunnerdMaintenanceOperations, + coldAdmissionTimeoutMs, resolveSourceCodexHome, settleRetainedRunnerdSession, retainedRunnerdCleanupProofIsCurrent, diff --git a/packages/paperclip-runner/src/native-session-runtime.test.ts b/packages/paperclip-runner/src/native-session-runtime.test.ts index 1515b729ca3..e872c426e4b 100644 --- a/packages/paperclip-runner/src/native-session-runtime.test.ts +++ b/packages/paperclip-runner/src/native-session-runtime.test.ts @@ -2395,6 +2395,42 @@ describe("executeNativeSession recovery", () => { expect(openRun).not.toHaveBeenCalled(); }); + it.each(["synchronous failure", "expired deadline"].flatMap(scenario => [false, true].map(retained => ({ scenario, retained }))))("cleans up startup cancellation after $scenario (retained: $retained)", async ({ scenario, retained }) => { + const abort = new AbortController(); + const remove = vi.spyOn(abort.signal, "removeEventListener"); + const failure = new Error("synchronous backend failure"); + const openSession = vi.fn((): Promise => { throw failure; }); + const attachRun = vi.fn((): Promise => { throw failure; }); + const openRun = vi.fn(async () => undefined); + const capabilities = { resume: false, typedEvents: true, steering: false, interruption: true, structuredResult: true }; + const backend: NativeSessionBackend = { + async descriptor() { return { kind: "mock", name: `startup-${scenario}`, version: "1", capabilities }; }, + openSession, + }; + const existingSession: NativeSession = { + identity: () => ({ ...identity, runId: "prior-run" }), async capabilities() { return capabilities; }, attachRun, + async *events() {}, async startTurn() { throw new Error("unexpected turn"); }, async result() { return null; }, + async snapshot() { throw new Error("unexpected snapshot"); }, async close() {}, + }; + const port: ControlPlanePort = { + openRun, + async appendEvent() { throw new Error("unexpected event"); }, + async replayEvents() { return { events: [], highestContiguousSourceSeq: 0 }; }, + async completeRun() {}, + }; + await expect(executeNativeSession({ input, backend, controlPlane: port, + runnerInstanceId: "runner-recovery", controlPlaneInstanceId: "control-recovery", + ...(retained ? { existingSession } : {}), + startup: { deadlineAtMs: Date.now() + (scenario === "expired deadline" ? -1 : 60_000), signal: abort.signal }, + })).rejects.toThrow(scenario === "expired deadline" ? retained ? "native_execution_startup_expired" : "native session bootstrap timed out after 0ms" : failure.message); + expect(openSession).toHaveBeenCalledTimes(!retained && scenario !== "expired deadline" ? 1 : 0); + expect(attachRun).toHaveBeenCalledTimes(retained && scenario !== "expired deadline" ? 1 : 0); + expect(openRun).not.toHaveBeenCalled(); + if (scenario === "synchronous failure") expect(remove).toHaveBeenCalledWith("abort", expect.any(Function)); + abort.abort(new Error("after completion")); + await Promise.resolve(); + }); + it("retains late bootstrap cleanup through the next admission", async () => { vi.useFakeTimers(); let releaseClose = () => {}; diff --git a/packages/paperclip-runner/src/native-session-runtime.ts b/packages/paperclip-runner/src/native-session-runtime.ts index d950524c265..445d3e5f615 100644 --- a/packages/paperclip-runner/src/native-session-runtime.ts +++ b/packages/paperclip-runner/src/native-session-runtime.ts @@ -212,6 +212,7 @@ export interface ExecuteNativeSessionOptions { turnTimeoutMs?: number; /** Abort admission while waiting for prior cleanup in the same domain. */ signal?: AbortSignal; + startup?: { deadlineAtMs: number; signal: AbortSignal }; /** Internal test seam; production bounds checkpoint persistence to 30 seconds. */ checkpointTimeoutMs?: number; /** Internal test seam; production uses the fixed 120-second platform policy. */ @@ -392,12 +393,29 @@ async function runAbortableOperationWithin(input: { timeoutError?: () => Error; operation: (signal: AbortSignal) => Promise; onLateResolution?: (value: T) => Promise | void; + signal?: AbortSignal; }): Promise { + input.signal?.throwIfAborted(); + if (input.timeoutMs <= 0) throw input.timeoutError?.() ?? new Error(input.timeoutMessage); const operationAbort = new AbortController(); let timer: ReturnType | undefined; let timedOut = false; let timeoutError: Error | undefined; - const operation = input.operation(operationAbort.signal); + let removeAbort = () => {}; + const cancelled = input.signal ? new Promise((_resolve, reject) => { + const abort = () => { + timedOut = true; + timeoutError = input.signal!.reason instanceof Error ? input.signal!.reason : new Error("native session startup cancelled"); + reject(timeoutError); + operationAbort.abort(timeoutError); + }; + input.signal!.addEventListener("abort", abort, { once: true }); + removeAbort = () => input.signal!.removeEventListener("abort", abort); + }) : undefined; + const operation = Promise.resolve().then(() => { + input.signal?.throwIfAborted(); + return input.operation(operationAbort.signal); + }); // The timeout path deliberately stops awaiting an uncooperative adapter. // Keep its eventual settlement observed after mutation authority is revoked. void operation.catch(() => undefined); @@ -415,6 +433,7 @@ async function runAbortableOperationWithin(input: { try { const value = await Promise.race([ operation, + ...(cancelled ? [cancelled] : []), new Promise((_resolve, reject) => { timer = setTimeout(() => { const error = @@ -434,6 +453,7 @@ async function runAbortableOperationWithin(input: { return value; } finally { if (timer !== undefined) clearTimeout(timer); + removeAbort(); } } @@ -2044,12 +2064,23 @@ export async function executeNativeSession( // state). Prove the provider attachment before opening durable // control-plane state because ControlPlanePort has no rollback operation. try { - await options.existingSession.attachRun({ identity, + const attach = async (signal?: AbortSignal) => { + signal?.throwIfAborted(); + await options.existingSession!.attachRun!({ identity, ...(input.provider.kind === "acpx" && input.provider.agent === "copilot" ? { currentRunGrant: { runtimeContext: "runtimeContext" in input ? input.runtimeContext : null, instructions: nativeSystemInstructions(input) } } : {}), + }); + signal?.throwIfAborted(); + }; + if (options.startup) await runAbortableOperationWithin({ + timeoutMs: Math.max(0, options.startup.deadlineAtMs - Date.now()), + timeoutMessage: "native_execution_startup_expired", + signal: options.startup.signal, + operation: attach, }); + else await attach(); } catch (error) { // attachRun has no transactional guarantee: a provider may bind the new // run before reporting a later failure. Conservatively quarantine the @@ -2168,7 +2199,7 @@ export async function executeNativeSession( // A fresh provider session is part of admission. Prove that it exists // before opening durable run state because ControlPlanePort intentionally // exposes no rollback for an admitted run. - const bootstrapTimeoutMs = options.timeoutMs ?? 900_000; + const bootstrapTimeoutMs = options.startup ? Math.max(0, options.startup.deadlineAtMs - Date.now()) : options.timeoutMs ?? 900_000; const bootstrapInput = { identity, workingDirectory: input.workspace.cwd ?? undefined, @@ -2176,6 +2207,7 @@ export async function executeNativeSession( session = await runAbortableOperationWithin({ timeoutMs: bootstrapTimeoutMs, timeoutMessage: `native session bootstrap timed out after ${bootstrapTimeoutMs}ms`, + signal: options.startup?.signal, operation: (signal) => options.backend.openSession({ ...bootstrapInput, diff --git a/packages/plugins/sdk/src/environment-tasks.ts b/packages/plugins/sdk/src/environment-tasks.ts index 5154431ca0e..f2ded292d19 100644 --- a/packages/plugins/sdk/src/environment-tasks.ts +++ b/packages/plugins/sdk/src/environment-tasks.ts @@ -1,5 +1,6 @@ import { z } from "zod"; import type { PluginEnvironmentDriverBaseParams, PluginEnvironmentLease } from "./protocol.js"; +import type { PluginAgentRuntimeIdentity } from "./types.js"; export const MAX_ENVIRONMENT_TASK_FILE_BYTES = 10 * 1024 * 1024; const taskFilePath = z.string().min(1).max(2000).refine(value => @@ -135,6 +136,7 @@ export const environmentTaskResultSchema = z.discriminatedUnion("kind", [ runner: z.object({ version: z.string().min(1), digest: z.string().regex(/^sha256:[a-f0-9]{64}$/) }).strict(), warmTaskVersion: z.literal(1).optional(), launchPreparationVersion: z.literal(1).optional(), + executionContextVersion: z.literal(1).optional(), freshWorkspaceVersion: z.literal(1).optional(), projectAuthCheckVersion: z.literal(1).optional(), workspaceFileReadVersion: z.literal(1).optional(), @@ -146,6 +148,7 @@ export const environmentTaskResultSchema = z.discriminatedUnion("kind", [ exitCode: z.number().int().optional(), /** Provider observed the task and all descendants stopped; terminal phase alone is insufficient. */ executionStopped: z.boolean().optional(), + preparationStopped: z.boolean().optional(), /** The provider retains this task's recoverable session state for a later launch. */ sessionStateAvailable: z.boolean().optional(), }).strict(), @@ -168,6 +171,15 @@ export type PluginEnvironmentPreparedLaunch = z.infer; export type PluginEnvironmentTaskOperation = z.infer; export type PluginEnvironmentTaskResult = z.infer; +export interface PluginEnvironmentTaskExecutionContext { + id: string; + deadlineAtMs: number; + transport: { + instanceSecrets: PluginAgentRuntimeIdentity["instanceSecrets"]; + configuration: Record; + }; +} + export interface PluginEnvironmentTaskParams extends Omit { /** Null during cleanup after the environment is deleted. */ environmentId: string | null; @@ -182,6 +194,8 @@ export interface PluginEnvironmentTaskParams extends Omit { + directory = await mkdtemp(path.join(tmpdir(), "paperclip-startup-worker-")); + entrypointPath = path.join(directory, "worker.cjs"); + await writeFile(entrypointPath, ` +const readline = require("node:readline"); +const held = []; +const invocations = []; +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 message = JSON.parse(line); + if (message.method === "initialize") send(message.id, {ok:true,supportedMethods:["environmentTask","configChanged","health"]}); + else if (message.method === "environmentTask" || message.method === "environmentDestroyLease") { + held.push(message.id); + invocations.push(message.paperclipInvocation?.id); + if (message.params.config.hostWrite) process.stdout.write(JSON.stringify({jsonrpc:"2.0",id:"host-"+message.id,method:"companies.get",params:{companyId:"company"},paperclipInvocationId:message.paperclipInvocation.id})+"\\n"); + if (message.params.config.replyEarly) { held.splice(held.indexOf(message.id),1); send(message.id,{kind:"description",taskId:"task"}); } + } + else if (message.method === "health") send(message.id, {status:"ok",details:{pending:held.length}}); + else if (message.method === "configChanged") { + if (message.params.config.exit) process.exit(1); + if (message.params.config.hostAfter) process.stdout.write(JSON.stringify({jsonrpc:"2.0",id:"late-host",method:"companies.get",params:{companyId:"company"},paperclipInvocationId:invocations.at(-1)})+"\\n"); + if (message.params.config.release) for (const id of held.splice(0)) send(id, {kind:"description",taskId:"task"}); + send(message.id, {}); + } else if (message.method === "shutdown") {send(message.id, {});setImmediate(()=>process.exit(0));} +}); +`, { mode: 0o600 }); +}); +afterAll(async () => { if (directory) await rm(directory, { recursive: true, force: true }); }); + +async function startWorker(hostHandlers: Record = {}) { + const manager = createPluginWorkerManager(); + await manager.startWorker(manifest.id, { entrypointPath, manifest, config: {}, + instanceInfo: { instanceId: "instance", hostVersion: "1.0.0" }, apiVersion: 1, hostHandlers: hostHandlers as never, autoRestart: false }); + return manager; +} + +describe("Core startup responsibility across real worker IPC", () => { + it.each(["stop", "replace", "stopAll"])("keeps stopped startup work reachable through %s until verified retirement", async mode => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const manager = await startWorker({ "companies.get": async () => { entered(); await gate; return {}; } }); + const original = manager.getWorker(manifest.id)!, settled = vi.fn(); + try { + const request = manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 50, undefined, settled, "stopped-context"); + const rejected = expect(request).rejects.toThrow("timed out"); + await started; await rejected; + if (mode === "stop") await manager.stopWorker(manifest.id); + else if (mode === "replace") await original.stop(); + else await manager.stopAll(); + if (mode !== "stopAll") await manager.startWorker(manifest.id, { entrypointPath, manifest, config: {}, + instanceInfo: { instanceId: "replacement", hostVersion: "1.0.0" }, apiVersion: 1, hostHandlers: {} }); + let drained = false, retired = false; + const draining = manager.drainStartup!("stopped-context").then(() => { drained = true; }); + const retiring = manager.retireStartup!("stopped-context").then(() => { retired = true; }); + await new Promise(resolve => setImmediate(resolve)); + expect(drained).toBe(false); + expect(retired).toBe(false); + expect(settled).not.toHaveBeenCalled(); + release(); await Promise.all([draining, retiring]); + expect(settled).toHaveBeenCalledOnce(); + const oldDrain = vi.spyOn(original, "drainStartup"); + await manager.drainStartup!("stopped-context"); + expect(oldDrain).not.toHaveBeenCalled(); + } finally { release(); await manager.stopAll(); } + }); + + it("prunes a stopped handle after its last accepted startup host call settles", async () => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const manager = await startWorker({ "companies.get": async () => { entered(); await gate; return {}; } }); + const original = manager.getWorker(manifest.id)!, settled = vi.fn(); + try { + const request = manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true, replyEarly: true } }, 5_000, undefined, settled, "host-only-context"); + await started; await request; + expect(settled).toHaveBeenCalledOnce(); + await manager.stopWorker(manifest.id); + let drained = false; + const draining = manager.drainStartup!("host-only-context").then(() => { drained = true; }); + await new Promise(resolve => setImmediate(resolve)); + expect(drained).toBe(false); + release(); await draining; + const oldDrain = vi.spyOn(original, "drainStartup"); + await manager.drainStartup!("host-only-context"); + expect(oldDrain).not.toHaveBeenCalled(); + } finally { release(); await manager.stopAll(); } + }); + + it("keeps a stopped handle owned after bounded retirement fails and permits verified retry", async () => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const manager = await startWorker({ "companies.get": async () => { entered(); await gate; return {}; } }), settled = vi.fn(); + try { + const request = manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 50, undefined, settled, "stopped-retry"); + const rejected = expect(request).rejects.toThrow("timed out"); + await started; await rejected; + await manager.stopWorker(manifest.id); + vi.useFakeTimers(); + const failed = manager.retireStartup!("stopped-retry").catch(error => error.message); + await vi.advanceTimersByTimeAsync(15_000); + expect(await failed).toBe("Task cleanup is still pending"); + expect(settled).not.toHaveBeenCalled(); + vi.useRealTimers(); + let retired = false; + const retry = manager.retireStartup!("stopped-retry").then(() => { retired = true; }); + await new Promise(resolve => setImmediate(resolve)); + expect(retired).toBe(false); + release(); await retry; + expect(settled).toHaveBeenCalledOnce(); + } finally { vi.useRealTimers(); release(); await manager.stopAll(); } + }); + + it.each(["stop", "stopAll"])("preserves a replacement registered while %s is finishing", async mode => { + const manager = await startWorker(), original = manager.getWorker(manifest.id)!; + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), stopped = new Promise(resolve => { entered = resolve; }); + const stop = original.stop.bind(original); + let replacement: Awaited> | undefined; + vi.spyOn(original, "stop").mockImplementation(async () => { await stop(); entered(); await gate; }); + try { + const stopping = mode === "stop" ? manager.stopWorker(manifest.id) : manager.stopAll(); + await stopped; + replacement = await manager.startWorker(manifest.id, { entrypointPath, manifest, config: {}, + instanceInfo: { instanceId: "replacement", hostVersion: "1.0.0" }, apiVersion: 1, hostHandlers: {} }); + release(); await stopping; + expect(manager.getWorker(manifest.id)).toBe(replacement); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ status: "ok" }); + } finally { release(); await replacement?.stop(); await manager.stopAll(); } + }); + + it("refuses new calls while the same startup group retires and joins concurrent retirement", async () => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const write = vi.fn(async () => { entered(); await gate; return {}; }); + const handlers = createHostClientHandlers({ pluginId: manifest.id, capabilities: ["companies.read"], services: { companies: { get: write } } as never }); + const manager = await startWorker(handlers), firstSettled = vi.fn(), secondSettled = vi.fn(); + try { + const request = manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 5_000, undefined, firstSettled, "retiring-context"); + void request.catch(() => undefined); + await started; + let retired = false; + const retiring = manager.retireStartup!("retiring-context").then(() => { retired = true; }); + const joined = manager.retireStartup!("retiring-context"); + await expect(manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 100, undefined, secondSettled, "retiring-context")) + .rejects.toThrow("retiring"); + expect(write).toHaveBeenCalledOnce(); + expect(firstSettled).not.toHaveBeenCalled(); + expect(secondSettled).toHaveBeenCalledOnce(); + expect(retired).toBe(false); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ status: "ok" }); + release(); await Promise.all([retiring, joined]); + expect(firstSettled).toHaveBeenCalledOnce(); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + await request; + expect(firstSettled).toHaveBeenCalledOnce(); + } finally { release(); await manager.stopAll(); } + }); + + it("revokes a cleanup invocation created after the initial drain and waits for its accepted host handler", async () => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const write = vi.fn(async () => { entered(); await gate; return {}; }); + const handlers = createHostClientHandlers({ pluginId: manifest.id, capabilities: ["companies.read"], services: { companies: { get: write } } as never }); + const manager = await startWorker(handlers), settled = vi.fn(); + try { + await manager.drainStartup!("cleanup-context"); + const request = manager.call(manifest.id, "environmentDestroyLease", { companyId: "company", environmentId: "environment", issueId: "issue", driverKey: "tasks", + config: { hostWrite: true }, providerLeaseId: "task", leaseMetadata: { runnerTaskStartupCleanup: { contextId: "cleanup-context", requiresPreparationStop: true } } }, + 5_000, undefined, settled); + await started; + let retired = false; + const retiring = manager.retireStartup!("cleanup-context").then(() => { retired = true; }); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { hostAfter: true } }); + expect(write).toHaveBeenCalledOnce(); + expect(retired).toBe(false); + release(); await retiring; + expect(settled).toHaveBeenCalledOnce(); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + await request; + expect(settled).toHaveBeenCalledOnce(); + expect(write).toHaveBeenCalledOnce(); + } finally { release(); await manager.stopAll(); } + }); + it("bounds the cleanup wait without discarding an accepted host handler", async () => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const manager = await startWorker({ "companies.get": async () => { entered(); await gate; return {}; } }), settled = vi.fn(); + try { + const request = manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 100, undefined, settled, "bounded-context"); + const rejected = expect(request).rejects.toThrow("timed out"); + await started; await rejected; + vi.useFakeTimers(); + const waiting = manager.retireStartup!("bounded-context"); + const pending = expect(waiting).rejects.toThrow("cleanup is still pending"); + await vi.advanceTimersByTimeAsync(15_000); await pending; + expect(settled).not.toHaveBeenCalled(); + vi.useRealTimers(); + const retried = vi.fn(); + const retry = manager.call(manifest.id, "environmentTask", params, 5_000, undefined, retried, "bounded-context"); + void retry.catch(() => undefined); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ details: { pending: 2 } }); + expect(retried).not.toHaveBeenCalled(); + release(); + await manager.drainStartup!("bounded-context"); + await manager.retireStartup!("bounded-context"); + expect(settled).toHaveBeenCalledOnce(); + expect(retried).toHaveBeenCalledOnce(); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + await retry; + } finally { vi.useRealTimers(); release(); await manager.stopAll(); } + }); + it.each(["startup", "cleanup"])("drains accepted %s host writes before retiring only verified bookkeeping", async mode => { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }), started = new Promise(resolve => { entered = resolve; }); + const manager = await startWorker({ "companies.get": async () => { entered(); await gate; return {}; } }); + const verified = vi.fn(), unrelated = vi.fn(); + try { + const request = mode === "startup" + ? manager.call(manifest.id, "environmentTask", { ...params, config: { hostWrite: true } }, 100, undefined, verified, "verified-context") + : manager.call(manifest.id, "environmentDestroyLease", { companyId: "company", environmentId: "environment", issueId: "issue", driverKey: "tasks", + config: { hostWrite: true }, providerLeaseId: "task", leaseMetadata: { runnerTaskStartupCleanup: { contextId: "verified-context", requiresPreparationStop: true } } }, + 100, undefined, verified); + const rejected = expect(request).rejects.toThrow("timed out"); + await started; + await rejected; + await expect(manager.call(manifest.id, "environmentTask", params, 10, undefined, unrelated, "other-context")).rejects.toThrow("timed out"); + let drained = false; + const drain = manager.drainStartup!("verified-context").then(() => { drained = true; }); + await manager.call(manifest.id, "health", {}); + expect(drained).toBe(false); + expect(verified).not.toHaveBeenCalled(); + release(); await drain; + expect(verified).not.toHaveBeenCalled(); + await manager.retireStartup!("verified-context"); + expect(verified).toHaveBeenCalledOnce(); + expect(unrelated).not.toHaveBeenCalled(); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + expect(verified).toHaveBeenCalledOnce(); + expect(unrelated).toHaveBeenCalledOnce(); + } finally { release(); await manager.stopAll(); } + }); + + it("can retire a verified startup after worker death without declaring unrelated work settled", async () => { + const manager = await startWorker(), verified = vi.fn(), unrelated = vi.fn(); + try { + await expect(manager.call(manifest.id, "environmentTask", params, 10, undefined, verified, "verified-context")).rejects.toThrow("timed out"); + await expect(manager.call(manifest.id, "environmentTask", params, 10, undefined, unrelated, "other-context")).rejects.toThrow("timed out"); + await expect(manager.call(manifest.id, "configChanged", { companyId: "company", config: { exit: true } })).rejects.toThrow("exited"); + expect(verified).not.toHaveBeenCalled(); + await manager.drainStartup!("verified-context"); + await manager.retireStartup!("verified-context"); + expect(verified).toHaveBeenCalledOnce(); + expect(unrelated).not.toHaveBeenCalled(); + } finally { await manager.stopAll(); } + }); + it("keeps a timed-out caller pending until the actual late worker reply", async () => { + const manager = await startWorker(); + const context = new NativeExecutionContext(execution, 5_000); + const settled = vi.fn(context.trackWorker()); + try { + await expect(manager.call(manifest.id, "environmentTask", params, 10, undefined, settled)).rejects.toThrow("timed out"); + expect(context.pending).toBe(true); + expect(settled).not.toHaveBeenCalled(); + context.abort(new Error("Stop requested")); + let drained = false; + const drain = context.drain().then(result => { drained = true; return result; }); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ details: { pending: 1 } }); + expect(drained).toBe(false); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + expect(await drain).toBe(true); + expect(context.pending).toBe(false); + expect(settled).toHaveBeenCalledOnce(); + expect(() => context.assertOpen()).toThrow("Stop requested"); + } finally { context.finish(); await manager.stopAll(); } + }); + + it("does not treat the startup deadline as a worker completion receipt", async () => { + const manager = await startWorker(); + const context = new NativeExecutionContext(execution, 100); + const settled = vi.fn(context.trackWorker()); + try { + await expect(manager.call(manifest.id, "environmentTask", params, 10, undefined, settled)).rejects.toThrow("timed out"); + if (!context.signal.aborted) await new Promise(resolve => context.signal.addEventListener("abort", () => resolve(), { once: true })); + expect(await context.drain()).toBe(false); + expect(context.pending).toBe(true); + expect(settled).not.toHaveBeenCalled(); + expect(await manager.call(manifest.id, "health", {})).toMatchObject({ details: { pending: 1 } }); + await manager.call(manifest.id, "configChanged", { companyId: "company", config: { release: true } }); + expect(settled).toHaveBeenCalledOnce(); + expect(context.pending).toBe(false); + expect(await context.drain()).toBe(true); + expect(() => context.assertOpen()).toThrow("startup_expired"); + } finally { context.finish(); await manager.stopAll(); } + }); +}); diff --git a/server/src/services/native-runtime/native-execution-context.test.ts b/server/src/services/native-runtime/native-execution-context.test.ts new file mode 100644 index 00000000000..1ba080bb9da --- /dev/null +++ b/server/src/services/native-runtime/native-execution-context.test.ts @@ -0,0 +1,106 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { NativeExecutionInput } from "../../vendor/paperclip-runner/index.js"; +import { NativeExecutionContext } from "./native-execution-context.js"; + +const execution = { binding: { companyId: "company", agentId: "agent", runId: "run" } } as NativeExecutionInput; +const target = { pluginId: "plugin", leaseId: "lease", taskId: "task", environmentId: "environment" }; + +afterEach(() => vi.useRealTimers()); + +describe("Core execution startup context", () => { + it("retains the selected input and resolves one exact transport target per context", async () => { + const context = new NativeExecutionContext(execution, 1_000); + const resolve = vi.fn(async () => ({ instanceSecrets: { binding: "synthetic" }, configuration: {} })); + try { + expect(context.execution).toBe(execution); + context.assertScope(execution.binding); + expect(() => context.assertScope({ ...execution.binding, runId: "other" })).toThrow("scope_mismatch"); + const first = await context.projection(target, resolve); + expect(await context.projection(target, resolve)).toEqual(first); + expect(resolve).toHaveBeenCalledOnce(); + await expect(context.projection({ ...target, taskId: "other" }, resolve)).rejects.toThrow("target_changed"); + context.finish(); + await expect(context.projection(target, resolve)).rejects.toThrow("startup_completed"); + } finally { context.finish(); } + }); + + it("reuses an identical transport target regardless of property insertion order", async () => { + const context = new NativeExecutionContext(execution, 1_000); + const resolve = vi.fn(async () => ({ instanceSecrets: {}, configuration: {} })); + try { + const first = await context.projection(target, resolve); + const reordered = { environmentId: target.environmentId, taskId: target.taskId, + leaseId: target.leaseId, pluginId: target.pluginId }; + expect(await context.projection(reordered, resolve)).toEqual(first); + expect(resolve).toHaveBeenCalledOnce(); + } finally { context.finish(); } + }); + + it("finishes a pending lookup without aborting a successful backend or completing worker responsibility", async () => { + vi.useFakeTimers(); + const context = new NativeExecutionContext(execution, 1_000), settled = context.trackWorker(); + const remove = vi.spyOn(context.signal, "removeEventListener"); + let outcome = "pending"; + const waiting = context.projection(target, () => new Promise(() => {})) + .then(() => { outcome = "resolved"; }, error => { outcome = error.message; }); + try { + await vi.advanceTimersByTimeAsync(0); + context.finish(); + await vi.advanceTimersByTimeAsync(0); + expect(outcome).toBe("native_execution_startup_completed"); + expect(context.signal.aborted).toBe(false); + expect(context.pending).toBe(true); + expect(remove).toHaveBeenCalledWith("abort", expect.any(Function)); + settled(); + expect(await context.drain()).toBe(true); + } finally { settled(); context.abort(new Error("test cleanup")); context.finish(); await waiting; } + }); + + it("keeps pending worker responsibility after cancellation until actual settlement", async () => { + vi.useFakeTimers(); + const context = new NativeExecutionContext(execution, 1_000), settled = context.trackWorker(); + let drained = false; + const drain = context.drain().then(result => { drained = true; return result; }); + try { + context.abort(new Error("Stop requested")); + expect(() => context.assertOpen()).toThrow("Stop requested"); + await vi.advanceTimersByTimeAsync(100); + expect(drained).toBe(false); + expect(context.pending).toBe(true); + settled(); + expect(await drain).toBe(true); + expect(context.pending).toBe(false); + } finally { settled(); context.finish(); } + }); + + it("reports an unsettled worker at the deadline without treating expiry as completion", async () => { + vi.useFakeTimers(); + const context = new NativeExecutionContext(execution, 1_000), settled = context.trackWorker(); + const drain = context.drain(); + try { + await vi.advanceTimersByTimeAsync(1_000); + expect(await drain).toBe(false); + expect(context.pending).toBe(true); + expect(() => context.assertOpen()).toThrow("startup_expired"); + settled(); + expect(await context.drain()).toBe(true); + } finally { settled(); context.finish(); } + }); + + it("rejects a delayed transport lookup before it can send a worker request", async () => { + vi.useFakeTimers(); + const context = new NativeExecutionContext(execution, 100); + let resolve!: (value: { instanceSecrets: {}; configuration: {} }) => void; + const transport = new Promise<{ instanceSecrets: {}; configuration: {} }>(complete => { resolve = complete; }); + const send = vi.fn(); + const operation = context.projection(target, () => transport).then(send); + const rejected = expect(operation).rejects.toThrow("startup_expired"); + await vi.advanceTimersByTimeAsync(100); + await rejected; + resolve({ instanceSecrets: {}, configuration: {} }); + await Promise.resolve(); + expect(send).not.toHaveBeenCalled(); + expect(context.pending).toBe(false); + context.finish(); + }); +}); diff --git a/server/src/services/native-runtime/native-execution-context.ts b/server/src/services/native-runtime/native-execution-context.ts new file mode 100644 index 00000000000..e635061d2b6 --- /dev/null +++ b/server/src/services/native-runtime/native-execution-context.ts @@ -0,0 +1,110 @@ +import { randomUUID } from "node:crypto"; +import type { PluginEnvironmentTaskExecutionContext } from "@paperclipai/plugin-sdk"; +import type { NativeExecutionInput } from "../../vendor/paperclip-runner/index.js"; + +export class NativeExecutionContext { + readonly id = randomUUID(); + readonly deadlineAtMs: number; + readonly signal: AbortSignal; + readonly execution: NativeExecutionInput; + #abort = new AbortController(); + #timer: ReturnType; + #calls = new Set(); + #waiters = new Set<() => void>(); + #finishWaiters = new Set<() => void>(); + #transport?: Promise; + #target?: string; + #finished = false; + pluginSupported = false; + + constructor(execution: NativeExecutionInput, timeoutMs: number) { + this.execution = execution; + this.deadlineAtMs = Date.now() + timeoutMs; + this.signal = this.#abort.signal; + this.#timer = setTimeout(() => this.abort(new Error("native_execution_startup_expired")), timeoutMs); + this.#timer.unref?.(); + } + + get active(): boolean { return !this.#finished; } + get pending(): boolean { return this.#calls.size > 0; } + + assertOpen(): void { + if (this.#finished) throw new Error("native_execution_startup_completed"); + if (Date.now() >= this.deadlineAtMs) this.abort(new Error("native_execution_startup_expired")); + this.signal.throwIfAborted(); + } + + assertScope(scope: { companyId: string; agentId: string; runId: string }): void { + const binding = this.execution.binding; + if (scope.companyId !== binding.companyId || scope.agentId !== binding.agentId || scope.runId !== binding.runId) { + throw new Error("native_execution_context_scope_mismatch"); + } + this.assertOpen(); + } + + async wait(operation: () => Promise): Promise { + this.assertOpen(); + let remove = () => {}; + try { + const value = await Promise.race([ + Promise.resolve().then(() => { this.assertOpen(); return operation(); }), + new Promise((_resolve, reject) => { + const abort = () => reject(this.signal.reason); + const finish = () => reject(new Error("native_execution_startup_completed")); + this.#finishWaiters.add(finish); + this.signal.addEventListener("abort", abort, { once: true }); + remove = () => { + this.signal.removeEventListener("abort", abort); + this.#finishWaiters.delete(finish); + }; + }), + ]); + this.assertOpen(); + return value; + } finally { remove(); } + } + + async projection(target: { pluginId: string; leaseId: string; taskId: string; environmentId: string }, + resolve: () => Promise): Promise { + this.assertOpen(); + const key = JSON.stringify([target.pluginId, target.leaseId, target.taskId, target.environmentId]); + if (this.#target !== undefined && this.#target !== key) throw new Error("native_execution_context_target_changed"); + this.#target = key; + this.#transport ??= resolve(); + const transport = await this.wait(() => this.#transport!); + this.assertOpen(); + return { id: this.id, deadlineAtMs: this.deadlineAtMs, transport }; + } + + trackWorker(): () => void { + this.assertOpen(); + const call = {}; + this.#calls.add(call); + return () => { + this.#calls.delete(call); + if (!this.#calls.size) for (const waiter of this.#waiters) waiter(); + }; + } + + async drain(): Promise { + if (!this.pending) return true; + let timer: ReturnType | undefined, complete!: () => void; + try { + await Promise.race([ + new Promise(resolve => { complete = resolve; this.#waiters.add(complete); }), + new Promise(resolve => { timer = setTimeout(resolve, Math.max(0, this.deadlineAtMs - Date.now())); }), + ]); + return !this.pending; + } finally { clearTimeout(timer); this.#waiters.delete(complete); } + } + + abort(reason: Error): void { this.#abort.abort(reason); } + + finish(): void { + clearTimeout(this.#timer); + this.#finished = true; + this.#transport = undefined; + for (const finish of this.#finishWaiters) finish(); + this.#finishWaiters.clear(); + } +} diff --git a/server/src/services/native-runtime/native-execution-lease.test.ts b/server/src/services/native-runtime/native-execution-lease.test.ts index ff5174f9fe1..4159771e327 100644 --- a/server/src/services/native-runtime/native-execution-lease.test.ts +++ b/server/src/services/native-runtime/native-execution-lease.test.ts @@ -13,6 +13,20 @@ function fixture() { } afterEach(() => vi.useRealTimers()); describe("native execution renewal ownership", () => { + it("reports only the remaining confirmed ownership budget", async () => { + vi.useFakeTimers(); + const { input, returning } = fixture(); + let complete!: (value: { runId: string }[]) => void; + returning.mockImplementationOnce(() => new Promise(resolve => { complete = resolve; })); + const guard = startNativeSessionExecutionLeaseRenewal({ ...input, confirmedDeadline: performance.now() + 45_000 }); + expect(guard.remainingMs()).toBe(45_000); + await vi.advanceTimersByTimeAsync(40_000); + expect(guard.remainingMs()).toBe(5_000); + complete([{ runId: "run" }]); + await vi.advanceTimersByTimeAsync(0); + expect(guard.remainingMs()).toBe(40_000); + await guard.stop(); + }); it("keeps a successfully renewed lease through repeated slow connection acquisition", async () => { vi.useFakeTimers(); const { input, update, transaction } = fixture(); diff --git a/server/src/services/native-runtime/native-execution-lease.ts b/server/src/services/native-runtime/native-execution-lease.ts index ae0cd119277..dc28b1ad9fe 100644 --- a/server/src/services/native-runtime/native-execution-lease.ts +++ b/server/src/services/native-runtime/native-execution-lease.ts @@ -107,9 +107,10 @@ type RenewalInput = { db: Db; runId: string; companyId: string; issueId: string; leaseOwner: string; attempt: number; controller: NativeControllerIdentity; remote?: boolean; + confirmedDeadline?: number; onLost: () => Promise; }; -export interface NativeExecutionLeaseRenewal { stop(): Promise; } +export interface NativeExecutionLeaseRenewal { stop(): Promise; remainingMs(): number; } // Preparation and execution share one renewal timer. Each caller releases its // own reference; the final release drains any SQL already in flight. @@ -117,6 +118,7 @@ const renewals = new WeakMap; pending: Promise; busy: boolean; lost: Error | null; deadline: ReturnType; loss: Promise; + confirmedDeadline: number; }>>(); export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): NativeExecutionLeaseRenewal { @@ -154,6 +156,7 @@ export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): Na // or recreate a timer after the final caller has stopped. if (state!.lost || state!.references === 0) return; clearTimeout(state!.deadline); + state!.confirmedDeadline = confirmedDeadline; state!.deadline = deadline(confirmedDeadline - performance.now()); }) .catch((error: unknown) => { @@ -167,13 +170,14 @@ export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): Na }; const timer = setInterval(renew, input.remote ? REMOTE_CONTROLLER_LEASE_MS / 3 : NATIVE_SESSION_EXECUTION_LEASE_RENEW_INTERVAL_MS); timer.unref?.(); + const confirmedDeadline = input.confirmedDeadline ?? performance.now() + ttlMs; state = { input, references: 1, timer, pending: Promise.resolve(), busy: false, lost: null, - deadline: deadline(ttlMs), loss: Promise.resolve() }; + deadline: deadline(confirmedDeadline - performance.now()), confirmedDeadline, loss: Promise.resolve() }; owners.set(input.leaseOwner, state); } const owned = state; let stopped = false; - return { stop: async () => { + return { remainingMs: () => owned.lost ? 0 : Math.max(0, owned.confirmedDeadline - performance.now()), stop: async () => { if (!stopped) { stopped = true; if (--owned.references === 0) { diff --git a/server/src/services/native-runtime/native-session-executor.test.ts b/server/src/services/native-runtime/native-session-executor.test.ts index c14d4875ff3..67194b6a9f3 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -11060,7 +11060,7 @@ describe("native session bounded recovery", () => { return query; }); const stop = vi.fn(async () => undefined); - const renew = vi.spyOn(nativeLeaseRenewal, "startNativeSessionExecutionLeaseRenewal").mockReturnValue({ stop }); + const renew = vi.spyOn(nativeLeaseRenewal, "startNativeSessionExecutionLeaseRenewal").mockReturnValue({ stop, remainingMs: () => 60_000 }); const releasePreparation = vi.fn(async () => { expect(renew).toHaveBeenCalledOnce(); expect(stop).not.toHaveBeenCalled(); diff --git a/server/src/services/plugin-worker-manager.ts b/server/src/services/plugin-worker-manager.ts index b738e2e9aa9..363d15c065c 100644 --- a/server/src/services/plugin-worker-manager.ts +++ b/server/src/services/plugin-worker-manager.ts @@ -541,6 +541,7 @@ interface PendingRequest { } interface ActiveInvocation { + startupContextId?: string; providerSpanBudget: {remaining:number}; scope: PluginInvocationScope; timer?: ReturnType; @@ -680,6 +681,9 @@ interface ExecuteLogRoute { * with exponential backoff automatically when `autoRestart` is enabled. */ export interface PluginWorkerHandle { + hasStartupWork?(): boolean; + drainStartup?(contextId: string): Promise; + retireStartup?(contextId: string): Promise; prepareIdleSleep?(hold: { ownerId: string; expiresAt: number }): Promise<"none" | "present" | "unknown">; releaseIdleSleep?(): void; /** The plugin ID this worker serves. */ @@ -720,6 +724,7 @@ export interface PluginWorkerHandle { timeoutMs?: number, executeLogSink?: ExecuteLogSink, onWorkerSettled?: () => void, + startupContextId?: string, ): Promise; /** @@ -800,6 +805,8 @@ export interface WorkerDiagnostics { * for starting/stopping all workers and routing RPC calls. */ export interface PluginWorkerManager { + drainStartup?(contextId: string): Promise; + retireStartup?(contextId: string): Promise; /** Subscribe after a worker successfully applies live company configuration. */ onConfigApplied(listener: () => void): () => void; inspectIdleSleep?(hold: { ownerId: string; expiresAt: number }): Promise<{ @@ -856,6 +863,8 @@ export interface PluginWorkerManager { params: HostToWorkerMethods[M][0], timeoutMs?: number, executeLogSink?: ExecuteLogSink, + onWorkerSettled?: () => void, + startupContextId?: string, ): Promise; /** @@ -882,6 +891,7 @@ export interface PluginWorkerManager { export function createPluginWorkerHandle( pluginId: string, options: WorkerStartOptions, + onStartupWorkChanged?: () => void, ): PluginWorkerHandle { const log = logger.child({ service: "plugin-worker", pluginId }); const emitter = new EventEmitter(); @@ -904,6 +914,26 @@ export function createPluginWorkerHandle( // Unlike RPC promises, these survive timeouts. A late response is the // completion receipt. A worker crash cannot manufacture one. const unsettledCalls = new Map void>(); + const startupCalls = new Map void>>(); + const startupRetirements = new Map>(); + const startupHostHandlers = new Map>>(); + const revokeStartupInvocations = (contextId: string) => { + for (const [id, invocation] of activeInvocations) { + if (invocation.startupContextId !== contextId) continue; + if (invocation.timer) clearTimeout(invocation.timer); + activeInvocations.delete(id); + } + }; + const waitForStartupHandlers = async (contextId: string) => { + if (!startupHostHandlers.get(contextId)?.size) return; + let timer: ReturnType | undefined; + try { + await Promise.race([ + 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); } + }; const unsettledInspections = new Set<() => void>(); let activeHostHandlers = 0; let idleHold: { ownerId: string; expiresAt: number; ready: boolean } | null = null; @@ -2756,6 +2786,15 @@ export function createPluginWorkerHandle( async function handleWorkerRequest(request: JsonRpcRequest): Promise { const done = beginIdleTrackedWork(); activeHostHandlers++; + const invocationId = readNonEmptyString((request as { paperclipInvocationId?: unknown }).paperclipInvocationId); + const startupId = invocationId ? activeInvocations.get(invocationId)?.startupContextId : undefined; + let complete = () => {}, pending: Promise | undefined; + if (startupId) { + pending = new Promise(resolve => { complete = resolve; }); + let handlers = startupHostHandlers.get(startupId); + if (!handlers) { handlers = new Set(); startupHostHandlers.set(startupId, handlers); } + handlers.add(pending); + } try { // A queued notification may already be running in the worker while its // prepare request is in flight. Let its writes finish until that worker @@ -2767,7 +2806,12 @@ export function createPluginWorkerHandle( return; } await dispatchWorkerRequest(request); - } finally { activeHostHandlers--; done(); } + } finally { + if (pending) startupHostHandlers.get(startupId ?? "")?.delete(pending); + if (startupId && startupHostHandlers.get(startupId)?.size === 0) startupHostHandlers.delete(startupId); + complete(); activeHostHandlers--; done(); + if (startupId) onStartupWorkChanged?.(); + } } async function dispatchWorkerRequest(request: JsonRpcRequest): Promise { @@ -3378,8 +3422,9 @@ export function createPluginWorkerHandle( timeoutMs?: number, executeLogSink?: ExecuteLogSink, onWorkerSettled?: () => void, + startupContextId?: string, ): Promise { - const dispatched = measureRunPhase("execution.plugin_dispatch", () => callInternalRaw(method, params, timeoutMs, executeLogSink, onWorkerSettled)); + const dispatched = measureRunPhase("execution.plugin_dispatch", () => callInternalRaw(method, params, timeoutMs, executeLogSink, onWorkerSettled, startupContextId)); // Tracing wraps the RPC in a new promise. Guard the promise handed to the // caller as well, without swallowing its eventual rejection for that caller. void dispatched.catch(() => undefined); @@ -3392,8 +3437,14 @@ export function createPluginWorkerHandle( timeoutMs?: number, executeLogSink?: ExecuteLogSink, onWorkerSettled?: () => void, + startupContextId?: string, ): Promise { const rpcPromise = new Promise((resolve, reject) => { + if (startupContextId && startupRetirements.has(startupContextId)) { + onWorkerSettled?.(); + reject(new Error("Plugin startup group is retiring")); + return; + } const idleControl = method === "prepareIdleSleep" || method === "releaseIdleSleep"; if (!idleControl && method !== "shutdown" && idleSleepHeld()) { onWorkerSettled?.(); @@ -3420,14 +3471,23 @@ export function createPluginWorkerHandle( finishIdle(); unsettledInspections.delete(finish); unsettledCalls.delete(id); + startupCalls.get(startupContextId ?? "")?.delete(finish); + if (startupContextId && startupCalls.get(startupContextId)?.size === 0) startupCalls.delete(startupContextId); onWorkerSettled?.(); + if (startupContextId) onStartupWorkChanged?.(); }; unsettledCalls.set(id, finish); + if (startupContextId) { + let calls = startupCalls.get(startupContextId); + if (!calls) { calls = new Set(); startupCalls.set(startupContextId, calls); } + calls.add(finish); + } if (method === "environmentInspectLease") unsettledInspections.add(finish); } const timeout = resolveRpcCallTimeoutMs(timeoutMs, rpcTimeoutMs); const invocationScope = deriveInvocationScope(method, params); const invocation = invocationScope ? registerInvocation(invocationScope) : null; + if (invocation && startupContextId) activeInvocations.get(invocation.id)!.startupContextId = startupContextId; // Register the host-owned execute route only for an execute call that // carries a log sink. The company id comes from the host-derived // invocation scope, never from the worker. This binds the sink to the @@ -3574,6 +3634,26 @@ export function createPluginWorkerHandle( await stopInternal(); await startInternal(); }, + hasStartupWork() { return startupCalls.size > 0 || startupHostHandlers.size > 0 || startupRetirements.size > 0; }, + async drainStartup(contextId: string) { + revokeStartupInvocations(contextId); + await waitForStartupHandlers(contextId); + }, + async retireStartup(contextId: string) { + const current = startupRetirements.get(contextId); + if (current) return current; + const retiring = (async () => { + revokeStartupInvocations(contextId); + await waitForStartupHandlers(contextId); + for (const finish of [...startupCalls.get(contextId) ?? []]) finish(); + })(); + startupRetirements.set(contextId, retiring); + try { await retiring; } + finally { + if (startupRetirements.get(contextId) === retiring) startupRetirements.delete(contextId); + onStartupWorkChanged?.(); + } + }, call( method: M, @@ -3581,6 +3661,7 @@ export function createPluginWorkerHandle( timeoutMs?: number, executeLogSink?: ExecuteLogSink, onWorkerSettled?: () => void, + startupContextId?: string, ): Promise { if (status !== "running" && status !== "starting") { onWorkerSettled?.(); @@ -3590,7 +3671,7 @@ export function createPluginWorkerHandle( ), ); } - return callInternal(method, params, timeoutMs, executeLogSink, onWorkerSettled); + return callInternal(method, params, timeoutMs, executeLogSink, onWorkerSettled, startupContextId); }, openLoginPtySession(input: LoginPtyOpenInput) { @@ -3786,6 +3867,14 @@ export function createPluginWorkerManager( ): PluginWorkerManager { const log = logger.child({ service: "plugin-worker-manager" }); const workers = new Map(); + const retiredStartupWorkers = new Set(); + const retainStartupWorker = (handle: PluginWorkerHandle) => { + if (handle.hasStartupWork?.()) retiredStartupWorkers.add(handle); + }; + const pruneStartupWorker = (handle: PluginWorkerHandle) => { + if (!handle.hasStartupWork?.()) retiredStartupWorkers.delete(handle); + }; + const startupWorkers = () => new Set([...workers.values(), ...retiredStartupWorkers]); const configApplied = new Set<() => void>(); let activeInspections = 0; /** Per-plugin startup locks to prevent concurrent spawn races. */ @@ -3839,7 +3928,8 @@ export function createPluginWorkerManager( // caller already supplied its own (a test may inject its own). duplexRouteSlots, ...options, - }); + }, () => pruneStartupWorker(handle)); + if (existing) retainStartupWorker(existing); workers.set(pluginId, handle); // Subscribe to crash/ready events for live event forwarding @@ -3886,12 +3976,19 @@ export function createPluginWorkerManager( log.info({ pluginId }, "stopping plugin worker"); await trackIdleWork(handle.stop()); - workers.delete(pluginId); + retainStartupWorker(handle); + if (workers.get(pluginId) === handle) workers.delete(pluginId); }, getWorker(pluginId: string): PluginWorkerHandle | undefined { return workers.get(pluginId); }, + async retireStartup(contextId: string) { + await Promise.all([...startupWorkers()].map(worker => worker.retireStartup?.(contextId))); + }, + async drainStartup(contextId: string) { + await Promise.all([...startupWorkers()].map(worker => worker.drainStartup?.(contextId))); + }, isRunning(pluginId: string): boolean { const handle = workers.get(pluginId); @@ -3915,10 +4012,12 @@ export function createPluginWorkerManager( }, "error stopping worker during shutdown", ); + } finally { + retainStartupWorker(handle); + if (workers.get(handle.pluginId) === handle) workers.delete(handle.pluginId); } }); await Promise.all(promises); - workers.clear(); }, diagnostics(): WorkerDiagnostics[] { @@ -3931,23 +4030,34 @@ export function createPluginWorkerManager( params: HostToWorkerMethods[M][0], timeoutMs?: number, executeLogSink?: ExecuteLogSink, + onWorkerSettled?: () => void, + startupContextId?: string, ): Promise { const handle = workers.get(pluginId); + if (!startupContextId && ["environmentReleaseLease", "environmentDestroyLease"].includes(method)) { + const cleanup = (params as { leaseMetadata?: { runnerTaskStartupCleanup?: { contextId?: unknown; requiresPreparationStop?: unknown } } }).leaseMetadata?.runnerTaskStartupCleanup; + if (typeof cleanup?.requiresPreparationStop === "boolean" && typeof cleanup.contextId === "string" && cleanup.contextId.trim()) startupContextId = cleanup.contextId.trim(); + } if (!handle) { + onWorkerSettled?.(); return Promise.reject( new Error(`No worker registered for plugin "${pluginId}"`), ); } let finishInspection: (() => void) | undefined; if (method === "environmentInspectLease") { - if (activeInspections >= 2) return Promise.reject(new Error("Provider inspection concurrency limit reached")); + if (activeInspections >= 2) { + onWorkerSettled?.(); + return Promise.reject(new Error("Provider inspection concurrency limit reached")); + } activeInspections += 1; let finished = false; finishInspection = () => { if (!finished) { finished = true; activeInspections -= 1; } }; } // A caller timeout is not worker completion. Inspection capacity is // released only by its actual reply, send failure, or worker exit. - const result = handle.call(method, params, timeoutMs, executeLogSink, finishInspection); + const finish = onWorkerSettled ? () => { finishInspection?.(); onWorkerSettled(); } : finishInspection; + const result = handle.call(method, params, timeoutMs, executeLogSink, finish, startupContextId); if (method === "configChanged") void result.then(() => { for (const listener of configApplied) { try { listener(); } diff --git a/server/src/vendor/paperclip-runner/index.ts b/server/src/vendor/paperclip-runner/index.ts index c2122cb1c6e..2abc5ca4213 100644 --- a/server/src/vendor/paperclip-runner/index.ts +++ b/server/src/vendor/paperclip-runner/index.ts @@ -138,6 +138,7 @@ export const createRunnerdCodexTransport: ( export const defaultCapabilityRunnerdBinary = runner.defaultCapabilityRunnerdBinary; export const runnerBinaryTarget = runner.runnerBinaryTarget; +export const coldAdmissionTimeoutMs = runner.coldAdmissionTimeoutMs; export const executeNativeSession = runner.executeNativeSession; export const applyNativeSessionGoalControl = runner.applyNativeSessionGoalControl;