diff --git a/src/main/database/migrations/v53-v53.ts b/src/main/database/migrations/v53-v53.ts new file mode 100644 index 00000000..83bf3131 --- /dev/null +++ b/src/main/database/migrations/v53-v53.ts @@ -0,0 +1,77 @@ +import log from "electron-log/main.js"; +import type { Migration, MigrationDeps } from "./types"; +import { nowIso, tableExists } from "./types"; + +export function createMigrationsV53V53(deps: MigrationDeps): Migration[] { + const { db } = deps; + return [ + { + version: 53, + name: "task_leases_restore_run_fk", + fn: () => { + if (!tableExists(db, "task_leases")) { + log.info("[migration v53] task_leases table missing; nothing to do"); + return; + } + + const danglingLeases = db.prepare( + `SELECT l.task_id AS taskId, l.run_id AS runId + FROM task_leases l + LEFT JOIN runs r ON r.id = l.run_id + WHERE r.id IS NULL`, + ).all() as Array<{ taskId: string; runId: string }>; + + if (danglingLeases.length > 0) { + const ts = nowIso(); + const repairTask = db.prepare( + `UPDATE tasks + SET active_run_id = CASE WHEN active_run_id = ? THEN NULL ELSE active_run_id END, + checkout_run_id = CASE WHEN checkout_run_id = ? THEN NULL ELSE checkout_run_id END, + execution_locked_at = CASE + WHEN active_run_id = ? OR checkout_run_id = ? THEN NULL + ELSE execution_locked_at + END, + updated_at = ? + WHERE id = ?`, + ); + for (const lease of danglingLeases) { + repairTask.run(lease.runId, lease.runId, lease.runId, lease.runId, ts, lease.taskId); + } + db.prepare( + `DELETE FROM task_leases + WHERE run_id NOT IN (SELECT id FROM runs)`, + ).run(); + log.warn( + `[migration v53] removed ${danglingLeases.length} dangling task lease(s) before restoring run_id FK`, + ); + } + + db.exec("PRAGMA defer_foreign_keys = ON"); + db.exec(` + CREATE TABLE task_leases_v53 ( + task_id TEXT PRIMARY KEY REFERENCES tasks(id) ON DELETE CASCADE, + run_id TEXT NOT NULL REFERENCES runs(id) ON DELETE CASCADE, + acquired_at TEXT NOT NULL, + expires_at TEXT NOT NULL, + checkout_run_id TEXT, + UNIQUE(task_id, run_id) + ) + `); + db.exec(` + INSERT INTO task_leases_v53 + (task_id, run_id, acquired_at, expires_at, checkout_run_id) + SELECT task_id, run_id, acquired_at, expires_at, checkout_run_id + FROM task_leases + `); + db.exec("DROP INDEX IF EXISTS idx_task_leases_expires_at"); + db.exec("DROP TABLE task_leases"); + db.exec("ALTER TABLE task_leases_v53 RENAME TO task_leases"); + db.exec( + "CREATE INDEX IF NOT EXISTS idx_task_leases_expires_at ON task_leases(expires_at)", + ); + + log.info("[migration v53] task_leases.run_id FK restored"); + }, + }, + ]; +} diff --git a/src/main/database/runs.ts b/src/main/database/runs.ts index c8a9c60b..03aa422e 100644 --- a/src/main/database/runs.ts +++ b/src/main/database/runs.ts @@ -7,6 +7,13 @@ import { type ReadAll, type Transaction, nowIso, toBoolean } from "./helpers"; import { buildRunLogMethods } from "./run-logs"; import { buildRunEventMethods } from "./run-events"; import { buildCostEntryMethods } from "./cost-entries"; +import { checkoutTaskInTransaction } from "./task-checkout"; + +class TaskCheckoutConflict extends Error { + constructor(readonly taskId: string) { + super(`Task checkout conflict: ${taskId}`); + } +} export function buildRunMethods( getDb: () => DatabaseSync, @@ -16,6 +23,33 @@ export function buildRunMethods( logsDir: string; }, ) { + function insertRunRow(input: { + id: string; + companyId: string; + taskId: string | null; + agentId: string; + workspaceId: string | null; + connectorId: ConnectorId; + logPath: string; + ts: string; + }): void { + getDb() + .prepare( + "insert into runs (id, company_id, task_id, agent_id, workspace_id, connector_id, status, summary, unattributed_cost, log_path, created_at, updated_at) values (?, ?, ?, ?, ?, ?, 'queued', '', 0, ?, ?, ?)", + ) + .run( + input.id, + input.companyId, + input.taskId, + input.agentId, + input.workspaceId, + input.connectorId, + input.logPath, + input.ts, + input.ts, + ); + } + return { ...buildRunLogMethods(getDb, readAll, transaction), ...buildRunEventMethods(getDb, readAll, transaction), @@ -51,19 +85,41 @@ export function buildRunMethods( const ts = nowIso(); const logPath = join(deps.logsDir, `${id}.log`); transaction(() => { - getDb() - .prepare( - "insert into runs (id, company_id, task_id, agent_id, workspace_id, connector_id, status, summary, unattributed_cost, log_path, created_at, updated_at) values (?, ?, ?, ?, ?, ?, 'queued', '', 0, ?, ?, ?)", - ) - .run(id, input.companyId, input.taskId, input.agentId, input.workspaceId, input.connectorId, logPath, ts, ts); - if (input.taskId !== null) { - getDb().prepare("update tasks set active_run_id = ?, updated_at = ? where id = ?").run(id, ts, input.taskId); - } + insertRunRow({ ...input, id, logPath, ts }); }); writeFileSync(logPath, "", "utf8"); return { id, logPath }; }, + createTaskRun(input: { + id?: string; + companyId: string; + taskId: string; + agentId: string; + workspaceId: string | null; + connectorId: ConnectorId; + }): { id: string; logPath: string } | null { + const id = input.id ?? randomUUID(); + const ts = nowIso(); + const logPath = join(deps.logsDir, `${id}.log`); + try { + transaction(() => { + insertRunRow({ ...input, id, logPath, ts }); + const checkoutOk = checkoutTaskInTransaction(getDb(), input.taskId, id); + if (!checkoutOk) { + throw new TaskCheckoutConflict(input.taskId); + } + }); + } catch (error) { + if (error instanceof TaskCheckoutConflict) { + return null; + } + throw error; + } + writeFileSync(logPath, "", "utf8"); + return { id, logPath }; + }, + markRunStatus(runId: string, status: RunRecord["status"], message: string | null) { const ts = nowIso(); const startedAt = status === "running" ? ts : null; diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index 2c8f6178..a2991d0e 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -18,6 +18,7 @@ import { createMigrationsV49V49 } from "./migrations/v49-v49"; import { createMigrationsV50V50 } from "./migrations/v50-v50"; import { createMigrationsV51V51 } from "./migrations/v51-v51"; import { createMigrationsV52V52 } from "./migrations/v52-v52"; +import { createMigrationsV53V53 } from "./migrations/v53-v53"; import { columnExists, nowIso, @@ -216,6 +217,7 @@ export function runMigrations( ...createMigrationsV50V50(deps), ...createMigrationsV51V51(deps), ...createMigrationsV52V52(deps), + ...createMigrationsV53V53(deps), ]; // Pre-migration snapshot. If any migration is pending we take a full diff --git a/src/main/database/task-checkout.ts b/src/main/database/task-checkout.ts new file mode 100644 index 00000000..489d2530 --- /dev/null +++ b/src/main/database/task-checkout.ts @@ -0,0 +1,53 @@ +import type { DatabaseSync } from "node:sqlite"; + +export const TASK_LEASE_TTL_MS = 15 * 60 * 1000; + +function isoFromMs(ms: number): string { + return new Date(ms).toISOString(); +} + +function assertRunExists(db: DatabaseSync, runId: string): void { + const row = db + .prepare("SELECT 1 AS present FROM runs WHERE id = ?") + .get(runId) as { present: number } | undefined; + if (!row) { + throw new Error(`Run not found for task checkout: ${runId}`); + } +} + +export function checkoutTaskInTransaction( + db: DatabaseSync, + taskId: string, + runId: string, +): boolean { + assertRunExists(db, runId); + const nowMs = Date.now(); + const ts = isoFromMs(nowMs); + const result = db + .prepare( + // eslint-disable-next-line no-restricted-syntax -- B2 Phase 1 grandfathered: checkout CAS is the single sanctioned status='in_progress' lock acquisition. The same transaction writes active_run_id and task_leases. + `UPDATE tasks + SET checkout_run_id = ?, + active_run_id = ?, + execution_locked_at = ?, + status = 'in_progress', + updated_at = ? + WHERE id = ? + AND checkout_run_id IS NULL + AND status IN ('todo', 'backlog')`, + ) + .run(runId, runId, ts, ts, taskId); + + if (result.changes === 0) { + return false; + } + + const expiresAt = isoFromMs(nowMs + TASK_LEASE_TTL_MS); + db.prepare( + `INSERT OR REPLACE INTO task_leases + (task_id, run_id, acquired_at, expires_at, checkout_run_id) + VALUES (?, ?, ?, ?, ?)`, + ).run(taskId, runId, ts, expiresAt, runId); + + return true; +} diff --git a/src/main/database/tasks-status.ts b/src/main/database/tasks-status.ts index 4ca6fcd3..e6a79a18 100644 --- a/src/main/database/tasks-status.ts +++ b/src/main/database/tasks-status.ts @@ -7,38 +7,7 @@ import { taskPriorityOrderSql, } from "./helpers"; import { chooseBestTaskAssignee } from "./task-assignment"; - -/** - * Default lease TTL used when the dual-write path seeds `task_leases` - * from the legacy `checkout_run_id` / `execution_locked_at` columns. - * Matches the stale-lock timeout in database-base so the sweep logic - * stays in sync across the shadow tables. - */ -const DUAL_WRITE_LEASE_TTL_MS = 15 * 60 * 1000; - -/** - * Mirror a successful legacy checkout onto the new `task_leases` row. - * - * Called from `checkoutTask` immediately after the legacy CAS succeeds. - * `INSERT OR REPLACE` is safe here because the preceding SQL CAS - * guarded on `checkout_run_id IS NULL` — any live lease row at this - * point refers to the same run (replay under retry) or is itself stale - * (recoverStaleLocks has not caught up yet). The new expires_at resets - * the stale-lock sweep deadline either way. - */ -function writeCheckoutLease( - db: DatabaseSync, - taskId: string, - runId: string, - ts: string, -): void { - const expiresAt = new Date(Date.now() + DUAL_WRITE_LEASE_TTL_MS).toISOString(); - db.prepare( - `INSERT OR REPLACE INTO task_leases - (task_id, run_id, acquired_at, expires_at, checkout_run_id) - VALUES (?, ?, ?, ?, ?)`, - ).run(taskId, runId, ts, expiresAt, runId); -} +import { TASK_LEASE_TTL_MS, checkoutTaskInTransaction } from "./task-checkout"; export function buildTaskStatusMethods( getDb: () => DatabaseSync, @@ -222,47 +191,7 @@ export function buildTaskStatusMethods( }, checkoutTask(taskId: string, runId: string): boolean { - const ts = nowIso(); - // Primary path is `status IN ('todo','backlog')`. The extra - // `in_progress AND active_run_id IS NULL` branch recovers from a - // specific orphan: finishRun has cleared active_run_id and - // checkout_run_id for the previous run but the status transition - // to done/in_review never fired (e.g. the run succeeded without a - // deliverable and automation parked it at in_progress, or a - // crashed run left the row behind). Without the reclaim, heartbeat - // picks the task (activeRunId=null passes the candidate filter) - // and then loops on `checkout_race_lost` forever because status - // never flips back to todo on its own. - // - // The `active_run_id IS NULL` guard preserves the original CAS - // contract for the busy case: a task whose active_run_id still - // points at a live run is refused here, exactly as before. - const result = getDb() - .prepare( - // eslint-disable-next-line no-restricted-syntax -- B2 Phase 1 grandfathered: checkoutTask CAS writes status='in_progress' as part of the atomic lock acquisition. The status→phase trigger mirrors to the new phase column and writeCheckoutLease below writes the shadow lease in the same call. - `UPDATE tasks - SET checkout_run_id = ?, execution_locked_at = ?, status = 'in_progress', updated_at = ? - WHERE id = ? - AND checkout_run_id IS NULL - AND ( - status IN ('todo', 'backlog') - OR (status = 'in_progress' AND active_run_id IS NULL) - )`, - ) - .run(runId, ts, ts, taskId); - if (result.changes > 0) { - // B2 Phase 1 step 3/5 — dual-write the new lease shadow. The - // status → phase trigger from v50 already mirrored the phase - // column by this point; the lease table still needs an explicit - // write. Same-run re-acquires overwrite in place, different-run - // acquires would return null (and should not happen here - // because the SQL CAS above already gated on `checkout_run_id - // IS NULL`). Wrapped in a best-effort try to keep the new path - // non-blocking for step 3 — the observability pass in step 4/5 - // surfaces any divergence. - writeCheckoutLease(getDb(), taskId, runId, ts); - } - return result.changes > 0; + return transaction(() => checkoutTaskInTransaction(getDb(), taskId, runId)); }, releaseTaskLock(taskId: string, runId: string): void { @@ -291,7 +220,7 @@ export function buildTaskStatusMethods( // Shadow renewal: push `expires_at` forward in lockstep with // `execution_locked_at` so the stale-lock sweep over the new // table stays consistent with the legacy sweep. - const expiresAt = new Date(Date.now() + DUAL_WRITE_LEASE_TTL_MS).toISOString(); + const expiresAt = new Date(Date.now() + TASK_LEASE_TTL_MS).toISOString(); getDb() .prepare( `UPDATE task_leases SET expires_at = ? WHERE task_id = ? AND run_id = ?`, diff --git a/src/main/orchestration/heartbeat.ts b/src/main/orchestration/heartbeat.ts index dbcb4133..a581215b 100644 --- a/src/main/orchestration/heartbeat.ts +++ b/src/main/orchestration/heartbeat.ts @@ -1,4 +1,3 @@ -import { randomUUID } from "node:crypto"; import type { AppDatabase } from "../database"; import type { AgentRecord, ConnectorRecord } from "@shared/types"; import type { OrchestrationDependencies } from "./types"; @@ -165,22 +164,9 @@ export function createHeartbeat( const candidateTasks = assignedTasks; const candidateTask = candidateTasks.find((entry) => !entry.activeRunId) ?? null; - const runId = randomUUID(); let task: ReturnType | null = null; if (candidateTask) { - const checkoutSuccess = db.tasks.checkoutTask(candidateTask.id, runId); - if (!checkoutSuccess) { - db.activityMod.addActivity({ - companyId, - actor: "system", - action: "execution.checkout_rejected", - entityType: "task", - entityId: candidateTask.id, - detail: `Agent ${agent.name} lost checkout race on task "${candidateTask.title}"`, - }); - return SKIPPED("checkout_race_lost"); - } task = db.entityCrud.getTask(candidateTask.id); } else { const hasActiveAssignedRun = candidateTasks.some( @@ -209,16 +195,6 @@ export function createHeartbeat( workspaceResolution = resolveWorkspace(task, agent, companyId); } catch (error) { if (error instanceof WorkspaceConfigurationError) { - // A coding-role task landed here with no bound workspace. We - // refuse to silently pick a shared workspace (the old - // company_fallback silent-corruption path) — the operator - // needs to bind a workspace explicitly. Release the lease - // the earlier checkoutTask just claimed so the task re-enters - // the pool at the next heartbeat instead of stalling as - // `in_progress` with no run. - if (task) { - db.tasks.releaseTaskLock(task.id, runId); - } db.activityMod.addActivity({ companyId, actor: "system", @@ -237,14 +213,35 @@ export function createHeartbeat( `[heartbeat] Workspace resolution for ${agent.name}: task.workspaceId=${task?.workspaceId ?? "null"}, agent.workspaceId=${agent.workspaceId ?? "null"}, projectId=${task?.projectId ?? "null"}, source=${workspaceResolution.source}, mode=${workspaceResolution.lockMode}, resolved=${workspace?.name ?? "NONE"}, localPath=${workspace?.localPath ?? "NONE"}`, ); - const run = db.runs.createRun({ - id: runId, - companyId, - taskId: task?.id ?? null, - agentId: agent.id, - workspaceId: workspaceResolution.lockMode === "exclusive" ? (workspace?.id ?? null) : null, - connectorId: connector.id, - }); + const run = task + ? db.runs.createTaskRun({ + companyId, + taskId: task.id, + agentId: agent.id, + workspaceId: workspaceResolution.lockMode === "exclusive" ? (workspace?.id ?? null) : null, + connectorId: connector.id, + }) + : db.runs.createRun({ + companyId, + taskId: null, + agentId: agent.id, + workspaceId: workspaceResolution.lockMode === "exclusive" ? (workspace?.id ?? null) : null, + connectorId: connector.id, + }); + if (!run) { + if (!task) { + throw new Error("Taskless heartbeat run creation returned null."); + } + db.activityMod.addActivity({ + companyId, + actor: "system", + action: "execution.checkout_rejected", + entityType: "task", + entityId: task.id, + detail: `Agent ${agent.name} lost checkout race on task "${task.title}"`, + }); + return SKIPPED("checkout_race_lost"); + } const companyBasic = db.company.getCompanyBasic(companyId); dispatchRunToWorker({ diff --git a/src/main/orchestration/task-runner.ts b/src/main/orchestration/task-runner.ts index cd10056a..9750f374 100644 --- a/src/main/orchestration/task-runner.ts +++ b/src/main/orchestration/task-runner.ts @@ -1,4 +1,3 @@ -import { randomUUID } from "node:crypto"; import type { AgentRecord, ConnectorRecord } from "@shared/types"; import type { DesktopResult } from "@shared/contracts"; import type { OrchestrationDependencies } from "./types"; @@ -87,22 +86,11 @@ export function createTaskRunner( ); } - // CAS checkout — atomic claim (EXEC-01, EXEC-02) - const runId = randomUUID(); - const checkoutSuccess = db.tasks.checkoutTask(task.id, runId); - if (!checkoutSuccess) { - return fail("CHECKOUT_RACE_LOST", "Another run claimed this task. Refresh and try again."); - } - let workspaceResolution: ReturnType; try { workspaceResolution = resolveWorkspace(task, agent, task.companyId); } catch (error) { if (error instanceof WorkspaceConfigurationError) { - // Resolver refused to guess a workspace. Release the lease - // we just claimed so the task row does not stall at - // `in_progress`; surface the error to the caller. - db.tasks.releaseTaskLock(task.id, runId); return fail("WORKSPACE_UNRESOLVED", error.message); } throw error; @@ -112,14 +100,16 @@ export function createTaskRunner( return fail("WORKSPACE_COMPANY_MISMATCH", "The selected workspace does not belong to the task company."); } - const run = db.runs.createRun({ - id: runId, + const run = db.runs.createTaskRun({ companyId: task.companyId, taskId: task.id, agentId: agent.id, workspaceId: workspaceResolution.lockMode === "exclusive" ? (workspace?.id ?? null) : null, connectorId: connector.id, }); + if (!run) { + return fail("CHECKOUT_RACE_LOST", "Another run claimed this task. Refresh and try again."); + } const companyBasic = db.company.getCompanyBasic(task.companyId); dispatchRunToWorker({ diff --git a/tests/unit/atomic-checkout.test.ts b/tests/unit/atomic-checkout.test.ts index a6062053..6882f3fa 100644 --- a/tests/unit/atomic-checkout.test.ts +++ b/tests/unit/atomic-checkout.test.ts @@ -1,5 +1,6 @@ import { join } from "node:path"; import { tmpdir } from "node:os"; +import { randomUUID } from "node:crypto"; import { mkdtemp, rm } from "node:fs/promises"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; @@ -124,6 +125,102 @@ describe("CAS checkout", () => { expect(task.status).toBe("in_progress"); }); + it("createRun does not claim a task before checkout", () => { + const companyId = createTestCompany(db); + createTestConnector(db); + const workspaceId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, workspaceId); + const taskId = createTestTask(db, companyId, null, workspaceId); + + db.runs.createRun({ + companyId, + taskId, + agentId, + workspaceId, + connectorId: "openai-codex", + }); + + const task = db.entityCrud.getTask(taskId); + expect(task.activeRunId).toBeNull(); + expect(task.status).toBe("todo"); + }); + + it("checkout rejects a missing run without mutating task state", () => { + const companyId = createTestCompany(db); + createTestConnector(db); + const workspaceId = createTestWorkspace(db, companyId); + const taskId = createTestTask(db, companyId, null, workspaceId); + const missingRunId = randomUUID(); + + expect(() => db.tasks.checkoutTask(taskId, missingRunId)).toThrow(/Run not found/); + + const task = db.entityCrud.getTask(taskId); + expect(task.activeRunId).toBeNull(); + expect(task.status).toBe("todo"); + const lease = (db as unknown as RawDb).db + .prepare("SELECT run_id FROM task_leases WHERE task_id = ?") + .get(taskId); + expect(lease).toBeUndefined(); + }); + + it("createTaskRun atomically creates the run and claims the task lease", () => { + const companyId = createTestCompany(db); + createTestConnector(db); + const workspaceId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, workspaceId); + const taskId = createTestTask(db, companyId, null, workspaceId); + + const run = db.runs.createTaskRun({ + companyId, + taskId, + agentId, + workspaceId, + connectorId: "openai-codex", + }); + + expect(run).not.toBeNull(); + const task = db.entityCrud.getTask(taskId); + expect(task.status).toBe("in_progress"); + expect(task.activeRunId).toBe(run?.id); + const lease = (db as unknown as RawDb).db + .prepare("SELECT run_id FROM task_leases WHERE task_id = ?") + .get(taskId) as { run_id: string } | undefined; + expect(lease?.run_id).toBe(run?.id); + }); + + it("createTaskRun rolls back the run row when checkout loses the race", () => { + const companyId = createTestCompany(db); + createTestConnector(db); + const workspaceId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, workspaceId); + const taskId = createTestTask(db, companyId, null, workspaceId); + + const firstRun = db.runs.createTaskRun({ + companyId, + taskId, + agentId, + workspaceId, + connectorId: "openai-codex", + }); + expect(firstRun).not.toBeNull(); + + const losingRunId = randomUUID(); + const secondRun = db.runs.createTaskRun({ + id: losingRunId, + companyId, + taskId, + agentId, + workspaceId, + connectorId: "openai-codex", + }); + + expect(secondRun).toBeNull(); + const stored = (db as unknown as RawDb).db + .prepare("SELECT id FROM runs WHERE id = ?") + .get(losingRunId); + expect(stored).toBeUndefined(); + }); + it("CAS checkout fails for second claimer", () => { const companyId = createTestCompany(db); createTestConnector(db);