diff --git a/doc/plugins/PLUGIN_RUNNER_EXECUTION.md b/doc/plugins/PLUGIN_RUNNER_EXECUTION.md index 22400cd5ddb..1a871bf041f 100644 --- a/doc/plugins/PLUGIN_RUNNER_EXECUTION.md +++ b/doc/plugins/PLUGIN_RUNNER_EXECUTION.md @@ -350,7 +350,9 @@ call `onEnvironmentReleaseLease`. Core still calls it when the release did not complete a dispatched task, and the cleanup retry calls it after any failure. When `onEnvironmentResumeLease` returns no provider lease, Core sends `status` first. It sends `stop` only when that `status` does not report -`executionStopped: true`. A driver must release task-scoped resources in its +`executionStopped: true` and any required `preparationStopped: true` proof. +An old warm-start preparation marker does not force a repeated stop when both +proofs are present. A driver must release task-scoped resources in its `status` or `complete` handling, not only in `onEnvironmentReleaseLease`. Termination receipts cover all run leases that owned the retired physical task. diff --git a/server/src/__tests__/environment-task-runtime.test.ts b/server/src/__tests__/environment-task-runtime.test.ts index 8007d130478..349c335aaca 100644 --- a/server/src/__tests__/environment-task-runtime.test.ts +++ b/server/src/__tests__/environment-task-runtime.test.ts @@ -88,7 +88,7 @@ describe("plugin-provided Runner execution", () => { if (preparationStopped) { expect(result).toEqual({ value: "completed", error: null }); expect(operations).toEqual(["stop", "status", "complete"]); } else { expect(result.error?.message).toContain("termination is unconfirmed"); expect(fixture.sets).toContainEqual(expect.objectContaining({ status: "pending_cleanup", cleanupStatus: "failed" })); } }); - it("still stops a startup task first when resume already observed the stop", async () => { + it.each([true, false])("checks startup preparation proof after resume observed the stop (%s)", async preparationStopped => { const value = row(); Object.assign(value.lease.metadata, { runnerTaskDispatched: false, runnerTaskStartupCleanup: { contextId: "startup", requiresPreparationStop: true } }); @@ -96,13 +96,15 @@ describe("plugin-provided Runner execution", () => { fixture.db.select = ((...args: unknown[]) => (database(value).db.select as (...values: unknown[]) => unknown)(...args)) as Db["select"]; Object.assign(workers, { drainStartup: vi.fn(async () => undefined) }); const operations: string[] = []; + let stopped = false; workers.call.mockImplementation(async (...args: any[]) => { const operation = args[2].operation.kind; operations.push(operation); - return operation === "status" ? { kind: "status", taskId: "attempt-1", phase: "cancelled", executionStopped: true, preparationStopped: true } + if (operation === "stop") stopped = true; + return operation === "status" ? { kind: "status", taskId: "attempt-1", phase: "cancelled", executionStopped: true, preparationStopped: preparationStopped || stopped } : { kind: "accepted", taskId: "attempt-1" }; }); await expect(settleDispatchedEnvironmentTask(fixture.db, workers as never, { companyId: "company", id: leaseId }, { executionStopObserved: true })).resolves.toBe("completed"); - expect(operations).toEqual(["stop", "status", "complete"]); + expect(operations).toEqual(preparationStopped ? ["status", "complete"] : ["status", "stop", "status", "complete"]); }); it("negotiates the transient context without reading transport for an older plugin", async () => { const context = new NativeExecutionContext({ binding: { companyId: "company", agentId: "agent", runId: "run" } } as never, 1_000); diff --git a/server/src/__tests__/plugin-warm-tasks.integration.test.ts b/server/src/__tests__/plugin-warm-tasks.integration.test.ts index da9cc6e3ae5..88d4d8605ff 100644 --- a/server/src/__tests__/plugin-warm-tasks.integration.test.ts +++ b/server/src/__tests__/plugin-warm-tasks.integration.test.ts @@ -48,7 +48,7 @@ const support=external ? {supported:true} : await getEmbeddedPostgresTestSupport if(method==="environmentResumeLease") {if(outage) throw Error("unreachable"); if(resumeGate){if(++resumes===2)releaseResume!();await resumeGate;}return {providerLeaseId:stopped||unavailable?null:p.providerLeaseId,expiresAt:taskExpiries.get(p.providerLeaseId),metadata:{runId:p.runId}};} if(method==="environmentReleaseLease") {if(releaseFails) throw Error("release unavailable");return {providerLeaseId:p.providerLeaseId,state:"stopped"};} - if(p.operation.kind==="status") return {kind:"status",taskId:p.taskId,phase:stopped?"completed":"running",executionStopped:stopped,sessionStateAvailable:stopped}; + if(p.operation.kind==="status") return {kind:"status",taskId:p.taskId,phase:stopped?"completed":"running",executionStopped:stopped,sessionStateAvailable:stopped,preparationStopped:stopped}; if(p.operation.kind==="read_state") return {kind:"state",taskId:p.taskId,state:{...p.operation.identity,lifecycle:"ready"}}; if(p.operation.kind==="stop") stopped=true; if(p.operation.kind==="prepare") return {kind:"prepared",taskId:p.taskId,launch:{configurationId:"configuration",requirementsDigest:`sha256:${"a".repeat(64)}`,provider:{}}}; @@ -113,8 +113,9 @@ const support=external ? {supported:true} : await getEmbeddedPostgresTestSupport expect(f.calls).toContain("complete"); await expect(resolvePluginTaskResume(db,{companyId:f.companyId,leaseId:next.lease.id,priorIdentity:f.identity})).resolves.toBe(f.first.lease.providerLeaseId); }); - it("retires a stopped retained task with one status and complete before allocating its replacement",async()=>{ + it.each([false,true])("retires a stopped retained task with one status and complete after warm preparation (%s)",async prepared=>{ const f=await fixture();f.stopped();await f.nextRun(); + if(prepared) await db.update(environmentLeases).set({metadata:sql`${environmentLeases.metadata} || ${JSON.stringify({runnerTaskStartup:{contextId:"settled-startup",deadlineAtMs:1,runId:f.identity.runId}})}::jsonb`}).where(eq(environmentLeases.id,f.first.lease.id)); const calls=f.calls.length,dispatched=f.dispatched.length; const next=await f.acquire(); expect(f.calls.slice(calls)).toEqual(["environmentResumeLease","status","complete","environmentAcquireLease"]); diff --git a/server/src/services/environment-task-runtime.ts b/server/src/services/environment-task-runtime.ts index 5ad92fad86b..6c091152cdf 100644 --- a/server/src/services/environment-task-runtime.ts +++ b/server/src/services/environment-task-runtime.ts @@ -513,7 +513,7 @@ export async function settleDispatchedEnvironmentTask(db: Db, workers: PluginWor void draining?.catch(() => undefined); const confirmed = (state?: Awaited>) => state?.kind === "status" && state.executionStopped === true && (!(startup || startupCleanup?.requiresPreparationStop) || state.preparationStopped === true); - let state = options.executionStopObserved === true && !startup && !startupCleanup + let state = options.executionStopObserved === true ? await executeEnvironmentTask(db, workers, { ...scope, operation: { kind: "status" } }) : undefined; if (!confirmed(state)) { await executeEnvironmentTask(db, workers, { ...scope, operation: { kind: "stop" } });