ZacharyZhang-NY/AgentCompany
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

fix(db): make task checkout atomic

Browse Source
This commit is contained in:
ZacharyZhang-NY committed 2026-04-22 13:38:54 -04:00
1 parent dffeccf3e2
commit 0785c9ba47
8 files changed
+329 -128
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/main/database/migrations/v53-v53.ts
+77
View File
@@ -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");
},
},
];
}
src/main/database/runs.ts
+64 -8
View File
@@ -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;
src/main/database/schema.ts
+2
View File
@@ -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
src/main/database/task-checkout.ts
+53
View File
@@ -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;
}
src/main/database/tasks-status.ts
+3 -74
View File
@@ -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 = ?`,