From 172a0c3bfbb935ab73854b03249a1eb99a8f86e3 Mon Sep 17 00:00:00 2001 From: Fabian Scherer Date: Tue, 14 Jul 2026 16:22:23 +0200 Subject: [PATCH] Fence task budget adjustments --- docs/quay-spec.md | 2 +- packages/cli/src/core/task_budget.ts | 39 +++++++++++--- packages/cli/tests/task/task_budget.test.ts | 60 +++++++++++++++++++++ 3 files changed, 93 insertions(+), 8 deletions(-) diff --git a/docs/quay-spec.md b/docs/quay-spec.md index 3292186..17c09d0 100644 --- a/docs/quay-spec.md +++ b/docs/quay-spec.md @@ -928,7 +928,7 @@ CREATE TABLE events ( event_id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL REFERENCES tasks(task_id), attempt_id INTEGER REFERENCES attempts(attempt_id), - event_type TEXT NOT NULL, -- spawned / pr_opened / ci_failed / ci_passed / merged / closed / changes_requested / conflict / blocker_ingested / malformed_signal_ingested / no_progress / crashed / spawn_failed / worker_auth_invalid / stale_detected / stale_killed / wall_clock_exceeded / wall_clock_killed / slack_reply_ingested / claimed / claim_expired / orchestrator_loop_parked / non_budget_loop_parked / brief_submitted / human_escalated / budget_exhausted / cancelled / worktree_error / tick_error + event_type TEXT NOT NULL, -- spawned / pr_opened / ci_failed / ci_passed / merged / closed / changes_requested / conflict / blocker_ingested / malformed_signal_ingested / no_progress / crashed / spawn_failed / worker_auth_invalid / stale_detected / stale_killed / wall_clock_exceeded / wall_clock_killed / slack_reply_ingested / claimed / claim_expired / orchestrator_loop_parked / non_budget_loop_parked / brief_submitted / human_escalated / budget_exhausted / task_budget_adjusted / cancelled / worktree_error / tick_error from_state TEXT, to_state TEXT, payload_artifact_id INTEGER REFERENCES artifacts(artifact_id), diff --git a/packages/cli/src/core/task_budget.ts b/packages/cli/src/core/task_budget.ts index 8ef0524..f034071 100644 --- a/packages/cli/src/core/task_budget.ts +++ b/packages/cli/src/core/task_budget.ts @@ -179,18 +179,24 @@ function adjustUnderLock( previous_budget_exhausted: task.budget_exhausted === 1, budget_exhausted: nextBudgetExhausted === 1, }; - const event = deps.db.transaction(() => { - deps.db + const txResult = deps.db.transaction(() => { + const upd = deps.db .query( `UPDATE tasks SET retry_budget = ?, budget_exhausted = ?, updated_at = ? - WHERE task_id = ?`, + WHERE task_id = ? + AND state = ? + AND cancel_requested_at IS NULL`, ) - .run(nextBudget, nextBudgetExhausted, now, task.task_id); + .run(nextBudget, nextBudgetExhausted, now, task.task_id, task.state); + const changes = (upd as { changes?: number }).changes ?? 0; + if (changes === 0) { + return { ok: false as const }; + } - return deps.db + const event = deps.db .query<{ event_id: number }, [string, string, string, string, string]>( `INSERT INTO events ( task_id, event_type, from_state, to_state, occurred_at, event_data @@ -204,8 +210,27 @@ function adjustUnderLock( now, JSON.stringify(eventData), ); + return { ok: true as const, event }; })(); - if (!event) throw new Error("task_budget_adjusted event insert returned no row"); + if (!txResult.ok) { + const current = loadTask(deps.db, task.task_id); + return { + ok: false, + error: { + code: "unsafe_state", + message: + `task ${task.task_id} changed before budget adjustment; retry after re-checking its current state`, + details: { + task_id: task.task_id, + observed_state: task.state, + current_state: current?.state ?? null, + }, + }, + }; + } + if (!txResult.event) { + throw new Error("task_budget_adjusted event insert returned no row"); + } return { ok: true, @@ -219,7 +244,7 @@ function adjustUnderLock( budget_exhausted: nextBudgetExhausted === 1, reason, forced: input.force === true, - event_id: event.event_id, + event_id: txResult.event.event_id, }, }; } diff --git a/packages/cli/tests/task/task_budget.test.ts b/packages/cli/tests/task/task_budget.test.ts index 1dc53f0..0583edc 100644 --- a/packages/cli/tests/task/task_budget.test.ts +++ b/packages/cli/tests/task/task_budget.test.ts @@ -131,6 +131,66 @@ test("adjust_task_budget raises budget, clears exhaustion, and records audit rea }); }); +test("adjust_task_budget refuses when state changes before fenced update", async () => { + const { h, built, taskId } = setupTask("awaiting-next-brief"); + let mutated = false; + const racingClock = { + now: () => h.clock.now(), + nowISO: () => { + if (!mutated) { + mutated = true; + h.db + .query(`UPDATE tasks SET state = 'queued' WHERE task_id = ?`) + .run(taskId); + } + return h.clock.nowISO(); + }, + }; + + const result = await adjust_task_budget( + { + db: h.db, + clock: racingClock, + supervisorLock: built.deps.supervisorLock, + }, + { + taskId, + by: 1, + reason: "operator recovery after substrate failure", + }, + ); + + expect(result).toMatchObject({ + ok: false, + error: { + code: "unsafe_state", + details: { + task_id: taskId, + observed_state: "awaiting-next-brief", + current_state: "queued", + }, + }, + }); + const task = h.db + .query<{ state: string; retry_budget: number; budget_exhausted: number }, [string]>( + `SELECT state, retry_budget, budget_exhausted FROM tasks WHERE task_id = ?`, + ) + .get(taskId); + expect(task).toEqual({ + state: "queued", + retry_budget: 5, + budget_exhausted: 1, + }); + const eventCount = h.db + .query<{ count: number }, [string]>( + `SELECT count(*) AS count + FROM events + WHERE task_id = ? AND event_type = 'task_budget_adjusted'`, + ) + .get(taskId); + expect(eventCount).toEqual({ count: 0 }); +}); + test("task increase-budget CLI lets blocker_resolved resume after budget is raised", async () => { const { h, built, taskId } = setupTask();