From 533b2a31811bc09b787df82a3dc541b998574463 Mon Sep 17 00:00:00 2001 From: "exe.dev user" Date: Thu, 29 Oct 2026 22:47:31 +0000 Subject: [PATCH] GEN-2716 preserve accepted run stop intent --- .../api/src/postgres/run-query-repository.ts | 6 +++- packages/api/src/postgres/run-repository.ts | 22 ++++++++---- packages/api/src/postgres/run-scheduler.ts | 7 +++- .../api/test/job-lifecycle.postgres.test.ts | 36 +++++++++++++++++++ .../api/test/postgres-run-repository.test.ts | 6 ++-- packages/api/test/run-finalization.test.ts | 1 + 6 files changed, 66 insertions(+), 12 deletions(-) diff --git a/packages/api/src/postgres/run-query-repository.ts b/packages/api/src/postgres/run-query-repository.ts index 8c7ec81a..e5f0f6d1 100644 --- a/packages/api/src/postgres/run-query-repository.ts +++ b/packages/api/src/postgres/run-query-repository.ts @@ -82,9 +82,13 @@ export class RunQueryRepository { async activity(runId: string) { return (await this.database.pool.query("SELECT action,detail,created_at FROM app.run_activity WHERE tenant_id=$1 AND run_id=$2 ORDER BY created_at,id", [this.tenantId, runId])).rows; } async stop(runId: string) { return this.database.transaction(async client => { - const result = await client.query("UPDATE app.job_runs SET state=CASE WHEN state='reserved' THEN 'stopped' ELSE 'stopping' END,next_poll_at=CASE WHEN state='reserved' THEN NULL ELSE now() END,updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('reserved','starting','running','blocked') RETURNING id,state", [this.tenantId, runId]); + const result = await client.query("UPDATE app.job_runs SET state=CASE WHEN state='reserved' THEN 'stopped' ELSE 'stopping' END,next_poll_at=CASE WHEN state='reserved' THEN NULL ELSE now() END,updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('reserved','starting','running','blocked','stopping') RETURNING id,state", [this.tenantId, runId]); const row = result.rows[0]; if (row?.state === "stopped") await client.query("UPDATE app.runs SET state='stopped',claim_released=true,vm_cleanup_complete=true,updated_at=now() WHERE tenant_id=$1 AND id=$2", [this.tenantId, runId]); + if (row?.state === "stopping") { + await client.query("UPDATE app.runs SET state='stopping',updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('starting','running','blocked','stopping')", [this.tenantId, runId]); + await client.query("INSERT INTO app.run_activity(tenant_id,id,run_id,action,detail) SELECT $1,$2,$3,'stop_requested','Stop request accepted' WHERE NOT EXISTS (SELECT 1 FROM app.run_activity WHERE tenant_id=$1 AND run_id=$3 AND action='stop_requested')", [this.tenantId, crypto.randomUUID(), runId]); + } if (row) await client.query("INSERT INTO app.wake_hints(key,not_before) VALUES ('global',now()) ON CONFLICT (key) DO UPDATE SET not_before=least(app.wake_hints.not_before,excluded.not_before),updated_at=now()"); return row ?? null; }); diff --git a/packages/api/src/postgres/run-repository.ts b/packages/api/src/postgres/run-repository.ts index fb2bb97d..3bc5e6a6 100644 --- a/packages/api/src/postgres/run-repository.ts +++ b/packages/api/src/postgres/run-repository.ts @@ -53,16 +53,18 @@ export class RunRepository { async markLaunched(run: PersistedRun, handle: RunHandle, destinationUrl: string, capabilities: readonly string[], traceSources?: TraceSources): Promise { await this.database.transaction(async client => { - await client.query("UPDATE app.job_runs SET state='running',started_at=COALESCE(started_at,now()),launch_lease_expires_at=NULL,next_poll_at=now()+interval '3 seconds',updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id]); - await client.query("UPDATE app.runs SET state='running',execution_handle=$3,execution_backend_kind=$4,execution_capabilities=$5,destination_url=$6,trace_sources=$7,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, handle, handle.backendKind, JSON.stringify(capabilities), destinationUrl, traceSources ? JSON.stringify(traceSources) : null]); + const result = await client.query<{ state: ExecutionState }>("UPDATE app.job_runs SET state=CASE WHEN state='stopping' THEN 'stopping' ELSE 'running' END,started_at=COALESCE(started_at,now()),launch_lease_expires_at=NULL,next_poll_at=now()+interval '3 seconds',updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('starting','stopping') RETURNING state", [run.tenantId, run.id]); + const state = result.rows[0]?.state; + if (!state) return; + await client.query("UPDATE app.runs SET state=$3,execution_handle=$4,execution_backend_kind=$5,execution_capabilities=$6,destination_url=$7,trace_sources=$8,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state, handle, handle.backendKind, JSON.stringify(capabilities), destinationUrl, traceSources ? JSON.stringify(traceSources) : null]); }); run.traceSources = traceSources ?? null; } async schedulePoll(run: PersistedRun, state: Extract, seconds = 3): Promise { await this.database.transaction(async client => { - await client.query("UPDATE app.job_runs SET state=$3,next_poll_at=now()+make_interval(secs=>$4),updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state, seconds]); - await client.query("UPDATE app.runs SET state=$3,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state]); + const result = await client.query<{ state: ExecutionState }>("UPDATE app.job_runs SET state=CASE WHEN state='stopping' THEN 'stopping' ELSE $3 END,next_poll_at=now()+make_interval(secs=>$4),updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('running','blocked','stopping') RETURNING state", [run.tenantId, run.id, state, seconds]); + if (result.rows[0]) await client.query("UPDATE app.runs SET state=$3,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, result.rows[0].state]); }); } @@ -74,10 +76,15 @@ export class RunRepository { return result.rows[0].execution_diagnostics; } + async recordStopAttempt(run: PersistedRun): Promise { + await this.database.pool.query("INSERT INTO app.run_activity(tenant_id,id,run_id,action,detail) VALUES ($1,$2,$3,'stop_attempt','Scheduler attempted agent termination')", [run.tenantId, crypto.randomUUID(), run.id]); + } + async terminal(run: PersistedRun, state: "succeeded" | "failed" | "stopped", artifactState: "stored" | "partial" | "failed", error?: string, cleanupComplete = false): Promise { await this.database.transaction(async client => { - await client.query("UPDATE app.job_runs SET state=$3,next_poll_at=NULL,launch_lease_expires_at=NULL,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state]); - await client.query("UPDATE app.runs SET state=$3,artifact_state=$4,artifact_error=$5,claim_released=true,vm_cleanup_complete=$6,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state, artifactState, error ? safeDiagnosticText(error) : null, cleanupComplete]); + const result = await client.query<{ state: ExecutionState }>("UPDATE app.job_runs SET state=CASE WHEN state='stopping' THEN 'stopped' ELSE $3 END,next_poll_at=NULL,launch_lease_expires_at=NULL,updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('starting','running','blocked','stopping') RETURNING state", [run.tenantId, run.id, state]); + if (!result.rows[0]) return; + await client.query("UPDATE app.runs SET state=$3,artifact_state=$4,artifact_error=$5,claim_released=true,vm_cleanup_complete=$6,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, result.rows[0].state, artifactState, error ? safeDiagnosticText(error) : null, cleanupComplete]); await client.query("DELETE FROM app.active_claims WHERE tenant_id=$1 AND run_id=$2", [run.tenantId, run.id]); // Automation runs before polling, so terminal events need another pass. await client.query("INSERT INTO app.wake_hints(key,not_before) VALUES ('global',now()) ON CONFLICT (key) DO UPDATE SET not_before=least(app.wake_hints.not_before,excluded.not_before),updated_at=now()"); @@ -88,7 +95,8 @@ export class RunRepository { await this.database.transaction(async client => { const attempt = (await client.query<{ vm_cleanup_attempt: number }>("UPDATE app.runs SET artifact_state=$3,artifact_error=$4,vm_cleanup_attempt=vm_cleanup_attempt+1,updated_at=now() WHERE tenant_id=$1 AND id=$2 RETURNING vm_cleanup_attempt", [run.tenantId, run.id, artifactState, safeDiagnosticText(error)])).rows[0]?.vm_cleanup_attempt ?? 1; const seconds = Math.min(60, 2 ** Math.min(attempt, 6)); - await client.query("UPDATE app.job_runs SET state='running',next_poll_at=now()+make_interval(secs=>$3),updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, seconds]); + const state = (await client.query<{ state: ExecutionState }>("UPDATE app.job_runs SET state=CASE WHEN state='stopping' THEN 'stopping' ELSE 'running' END,next_poll_at=now()+make_interval(secs=>$3),updated_at=now() WHERE tenant_id=$1 AND id=$2 AND state IN ('starting','running','blocked','stopping') RETURNING state", [run.tenantId, run.id, seconds])).rows[0]?.state; + if (state) await client.query("UPDATE app.runs SET state=$3,updated_at=now() WHERE tenant_id=$1 AND id=$2", [run.tenantId, run.id, state]); await client.query("INSERT INTO app.run_activity(tenant_id,id,run_id,action,detail) VALUES ($1,$2,$3,'finalization_retry',$4)", [run.tenantId, crypto.randomUUID(), run.id, `Attempt ${attempt}: ${safeDiagnosticText(error)}`]); }); } diff --git a/packages/api/src/postgres/run-scheduler.ts b/packages/api/src/postgres/run-scheduler.ts index 9b96394b..667b168c 100644 --- a/packages/api/src/postgres/run-scheduler.ts +++ b/packages/api/src/postgres/run-scheduler.ts @@ -47,8 +47,13 @@ export class RunScheduler { } private async poll(run: PersistedRun) { + if (!run.executionHandle && run.state === "stopping") return this.runs.schedulePoll(run, "stopping"); const backend = await this.backend(run); if (!run.executionHandle) throw new Error("Execution handle is missing"); - const observation = run.executionDiagnostics ?? (run.state === "stopping" && backend instanceof ExeVmBackend ? await backend.terminate(run.executionHandle) : await backend.inspect(run.executionHandle)); + const cancellationBackend = backend instanceof ExeVmBackend || backend instanceof AmpBackend; + if (run.state === "stopping" && cancellationBackend && !run.executionDiagnostics) await this.runs.recordStopAttempt(run); + const observation = run.executionDiagnostics ?? (run.state === "stopping" && cancellationBackend + ? backend instanceof ExeVmBackend ? await backend.terminate(run.executionHandle) : await backend.stop(run.executionHandle) + : await backend.inspect(run.executionHandle)); if (["running", "blocked", "stopping"].includes(observation.state)) { if (backend instanceof ExeVmBackend) await this.collectTraceChunk(run, backend).catch(() => { traceMetric("chunk_collection_failure", run.tenantId, run.id); }); return this.runs.schedulePoll(run, observation.state as "running" | "blocked" | "stopping"); diff --git a/packages/api/test/job-lifecycle.postgres.test.ts b/packages/api/test/job-lifecycle.postgres.test.ts index c98cfde2..3e6fa038 100644 --- a/packages/api/test/job-lifecycle.postgres.test.ts +++ b/packages/api/test/job-lifecycle.postgres.test.ts @@ -4,6 +4,7 @@ import { Database } from "../src/postgres/database"; import { PostgresJobRepository } from "../src/postgres/job-repository"; import { AutomationScheduler } from "../src/postgres/automation-scheduler"; import { RunRepository } from "../src/postgres/run-repository"; +import { RunQueryRepository } from "../src/postgres/run-query-repository"; import { InvocationService, type Job, type Trigger } from "../src/job-domain"; import { decrypt, encrypt } from "../src/crypto"; @@ -72,6 +73,41 @@ describe.skipIf(!process.env.AUTH_TEST_DATABASE_URL)("job lifecycle delivery (Po expect(await deliveries()).toHaveLength(1); }); + it("preserves one accepted stop across a stale poll and late launch completion", async () => { + const a = await tenant(), source = await job(a); + const create = (claimKey: string) => new InvocationService(repository(a), value => encrypt(value, key)).invoke(source.id, { + source: "manual", triggerId: source.triggers[0]!.id, claimKey, context: {}, + }); + + // Poll race: retain the pre-stop running snapshot, accept stop, then save + // the stale running observation as if backend inspection just completed. + const first = await create(crypto.randomUUID()); + let run = await runs.claimNext(); + await runs.markLaunched(run!, { backendKind: "exe-vm", id: "factorize-running" }, "https://vm", []); + const stalePoll = { ...run!, state: "running" as const, executionHandle: { backendKind: "exe-vm", id: "factorize-running" } }; + const stops = new RunQueryRepository(database, a); + expect(await stops.stop(first.run.id)).toMatchObject({ state: "stopping" }); + expect(await stops.stop(first.run.id)).toMatchObject({ state: "stopping" }); + await runs.schedulePoll(stalePoll, "running"); + await runs.retryFinalization(stalePoll, "partial", "temporary artifact issue"); + let states = await database.pool.query("SELECT jr.state job_state,r.state run_state FROM app.job_runs jr JOIN app.runs r USING(tenant_id,id) WHERE jr.tenant_id=$1 AND jr.id=$2", [a, first.run.id]); + expect(states.rows[0]).toEqual({ job_state: "stopping", run_state: "stopping" }); + await runs.terminal(stalePoll, "stopped", "stored", undefined, true); + + // Launch race: stop while the backend launch is paused, then let its + // completion persist the newly available handle. + const second = await create(crypto.randomUUID()); + run = await runs.claimNext(); + expect(run?.id).toBe(second.run.id); + expect(await stops.stop(second.run.id)).toMatchObject({ state: "stopping" }); + await runs.markLaunched(run!, { backendKind: "exe-vm", id: "factorize-late" }, "https://vm", []); + states = await database.pool.query("SELECT jr.state job_state,r.state run_state,r.execution_handle FROM app.job_runs jr JOIN app.runs r USING(tenant_id,id) WHERE jr.tenant_id=$1 AND jr.id=$2", [a, second.run.id]); + expect(states.rows[0]).toMatchObject({ job_state: "stopping", run_state: "stopping", execution_handle: { id: "factorize-late" } }); + await runs.terminal({ ...run!, id: second.run.id }, "succeeded", "stored", undefined, true); + states = await database.pool.query("SELECT jr.state job_state,r.state run_state FROM app.job_runs jr JOIN app.runs r USING(tenant_id,id) WHERE jr.tenant_id=$1 AND jr.id=$2", [a, second.run.id]); + expect(states.rows[0]).toEqual({ job_state: "stopped", run_state: "stopped" }); + }); + it("fires job-idle only for a completion after the trigger exists", async () => { const a = await tenant(), source = await job(a); await complete(a, source); diff --git a/packages/api/test/postgres-run-repository.test.ts b/packages/api/test/postgres-run-repository.test.ts index 6d1ea52f..ba085ec9 100644 --- a/packages/api/test/postgres-run-repository.test.ts +++ b/packages/api/test/postgres-run-repository.test.ts @@ -13,12 +13,12 @@ describe("RunRepository", () => { }); it("persists the source declaration at launch and reads it independently of current driver defaults", async () => { - const query = vi.fn(async (_sql: string, _values: unknown[]) => ({ rows: [], rowCount: 1 })); + const query = vi.fn(async (sql: string, _values: unknown[]) => ({ rows: sql.includes("UPDATE app.job_runs") ? [{ state: "running" }] : [], rowCount: 1 })); const repository = new RunRepository({ transaction: (work: any) => work({ query }), pool: { query } } as any); const sources = { primary: { kind: "execution_stream" as const, path: "/tmp/stream.jsonl", mediaType: "application/x-ndjson" as const, provider: "codex" as const, harnessVersion: "v1" } }; await repository.markLaunched({ tenantId: "tenant", id: "run" } as any, { backendKind: "exe-vm", id: "factorize-run" }, "https://vm", [], sources); - expect(query.mock.calls[1][0]).toContain("trace_sources=$7"); - expect(query.mock.calls[1][1][6]).toBe(JSON.stringify(sources)); + expect(query.mock.calls[1][0]).toContain("trace_sources=$8"); + expect(query.mock.calls[1][1][7]).toBe(JSON.stringify(sources)); await repository.dueForPoll(); expect(query.mock.calls.at(-1)![0]).toContain("x.trace_sources,x.execution_handle"); }); diff --git a/packages/api/test/run-finalization.test.ts b/packages/api/test/run-finalization.test.ts index dbed2953..38005b41 100644 --- a/packages/api/test/run-finalization.test.ts +++ b/packages/api/test/run-finalization.test.ts @@ -145,6 +145,7 @@ describe("durable run finalization", () => { f.run.state = "stopping"; f.run.traceSources = { primary: { kind: "execution_stream", path: "/tmp/codex-exec.jsonl", provider: "codex" } }; f.run.executionConfig.agentKind = "codex"; + vi.spyOn(RunRepository.prototype, "recordStopAttempt").mockResolvedValue(); const terminate = vi.spyOn(f.backend, "terminate").mockImplementation(async () => { f.events.push("terminate"); return { state: "stopped" }; }); vi.spyOn(f.scheduler, "finalizePrimary").mockResolvedValue([]); await f.scheduler.poll(f.run);