GEN-2716 preserve accepted run stop intent

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<any>("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<any>("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;
     });

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<void> {
     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<ExecutionState, "running" | "blocked" | "stopping">, seconds = 3): Promise<void> {
     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<void> {
+    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<void> {
     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)}`]);
     });
   }

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");

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);

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");
   });

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);