refactor(database): move task runtime ownership to leases

This commit is contained in:
ZacharyZhang-NY committed 2026-04-24 00:18:51 -04:00
1 parent 1f4af89be0
commit f8eb57ad8c
21 files changed
+419 -162

No files matched your search

+18 -18
View File
@@ -184,12 +184,12 @@ function createAndStartRun(
* Verifies database consistency after crash recovery:
* - No runs in "running" status
* - No agents in "running" status
* - No orphaned checkout_run_ids pointing to non-existent runs
* - No orphaned task leases pointing to non-existent runs
*/
function verifyDbConsistency(db: AppDatabase): {
runningRuns: number;
runningAgents: number;
orphanedCheckouts: number;
orphanedTaskLeases: number;
} {
const runningRuns = (
(db as unknown as RawDb).db
@@ -203,17 +203,17 @@ function verifyDbConsistency(db: AppDatabase): {
.get() as { cnt: number }
).cnt;
const orphanedCheckouts = (
const orphanedTaskLeases = (
(db as unknown as RawDb).db
.prepare(
`SELECT count(*) as cnt FROM tasks t
LEFT JOIN runs r ON t.checkout_run_id = r.id
WHERE t.checkout_run_id IS NOT NULL AND r.id IS NULL`,
`SELECT count(*) as cnt FROM task_leases l
LEFT JOIN runs r ON l.run_id = r.id
WHERE r.id IS NULL`,
)
.get() as { cnt: number }
).cnt;
return { runningRuns, runningAgents, orphanedCheckouts };
return { runningRuns, runningAgents, orphanedTaskLeases };
}
// ---------------------------------------------------------------------------
@@ -292,20 +292,20 @@ describe("Process Recovery -- Worker Crash Detection", () => {
});
// Verify lock is held before crash
const taskBefore = (db as unknown as RawDb).db
const leaseBefore = (db as unknown as RawDb).db
.prepare(
"SELECT checkout_run_id FROM tasks WHERE id = ?",
"SELECT run_id FROM task_leases WHERE task_id = ?",
)
.get(taskId) as { checkout_run_id: string | null };
expect(taskBefore.checkout_run_id).toBe(runId);
.get(taskId) as { run_id: string | null };
expect(leaseBefore.run_id).toBe(runId);
simulateWorkerCrash(db);
// interruptStaleRuns clears active_run_id but does not directly clear checkout_run_id.
// The checkout_run_id is cleared by finishRun or releaseTaskLock.
// However, after crash, the task is set to "blocked" and a subsequent
// stale lock recovery pass handles the checkout_run_id.
// For this test, verify the task status is blocked (the recovery mechanism).
const leaseAfter = (db as unknown as RawDb).db
.prepare("SELECT run_id FROM task_leases WHERE task_id = ?")
.get(taskId);
expect(leaseAfter).toBeUndefined();
const task = db.core.entity.getTask(taskId);
expect(task.status).toBe("blocked");
});
@@ -483,7 +483,7 @@ describe("Process Recovery -- Context Preservation (D-02)", () => {
const consistency = verifyDbConsistency(db);
expect(consistency.runningRuns).toBe(0);
expect(consistency.runningAgents).toBe(0);
expect(consistency.orphanedCheckouts).toBe(0);
expect(consistency.orphanedTaskLeases).toBe(0);
});
});
@@ -579,6 +579,6 @@ describe("Process Recovery -- Multiple Concurrent Runs", () => {
const consistency = verifyDbConsistency(db);
expect(consistency.runningRuns).toBe(0);
expect(consistency.runningAgents).toBe(0);
expect(consistency.orphanedCheckouts).toBe(0);
expect(consistency.orphanedTaskLeases).toBe(0);
});
});
+47 -35
View File
@@ -24,6 +24,12 @@ type RawDb = {
};
};
type TaskLeaseRow = {
task_id: string;
run_id: string;
expires_at: string;
};
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
@@ -84,6 +90,14 @@ function createTestTask(db: AppDatabase, companyId: string, agentId: string | nu
});
}
function getTaskLease(db: AppDatabase, taskId: string): TaskLeaseRow | null {
return (
((db as unknown as RawDb).db
.prepare("SELECT task_id, run_id, expires_at FROM task_leases WHERE task_id = ?")
.get(taskId) as TaskLeaseRow | undefined) ?? null
);
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
@@ -157,11 +171,29 @@ describe("CAS checkout", () => {
const task = db.core.entity.getTask(taskId);
expect(task.activeRunId).toBeNull();
expect(task.status).toBe("todo");
const lock = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(lock.checkout_run_id).toBeNull();
expect(lock.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
});
it("failed checkout keeps a taskless run unbound", () => {
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 as unknown as RawDb).db.prepare("UPDATE tasks SET status = 'in_progress' WHERE id = ?").run(taskId);
const { id: runId } = db.core.runs.createRun({
companyId,
taskId: null,
agentId,
workspaceId,
connectorId: "openai-codex",
});
expect(db.core.tasks.checkoutTask(taskId, runId)).toBe(false);
expect(db.core.runs.getRun(runId).taskId).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
expect(db.core.runs.hasActiveRun({ companyId, taskId })).toBe(false);
});
it("createTaskRun atomically creates the run and claims the task row lock", () => {
@@ -183,10 +215,7 @@ describe("CAS checkout", () => {
const task = db.core.entity.getTask(taskId);
expect(task.status).toBe("in_progress");
expect(task.activeRunId).toBe(run?.id);
const lock = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null };
expect(lock.checkout_run_id).toBe(run?.id);
expect(getTaskLease(db, taskId)?.run_id).toBe(run?.id);
});
it("createTaskRun rolls back the run row when checkout loses the race", () => {
@@ -269,7 +298,7 @@ describe("CAS checkout", () => {
expect(db.core.tasks.checkoutTask(taskId, runId)).toBe(false);
});
it("lock release in finishRun clears checkout fields", () => {
it("lock release in finishRun clears the task lease", () => {
const companyId = createTestCompany(db);
createTestConnector(db);
const workspaceId = createTestWorkspace(db, companyId);
@@ -297,15 +326,10 @@ describe("CAS checkout", () => {
unattributedCost: false,
});
const row = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(row.checkout_run_id).toBeNull();
expect(row.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
});
it("lock release in cancelRunForDeletion clears checkout fields", () => {
it("lock release in cancelRunForDeletion clears the task lease", () => {
const companyId = createTestCompany(db);
createTestConnector(db);
const workspaceId = createTestWorkspace(db, companyId);
@@ -326,12 +350,7 @@ describe("CAS checkout", () => {
// Deleting the agent triggers cancelRunForDeletion for active runs
db.core.agents.deleteAgent(agentId, companyId);
const row = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(row.checkout_run_id).toBeNull();
expect(row.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
});
it("stale lock recovery resets tasks older than 15 minutes", () => {
@@ -351,10 +370,10 @@ describe("CAS checkout", () => {
db.core.tasks.checkoutTask(taskId, runId);
// Backdate execution_locked_at to 20 minutes ago
// Backdate the lease expiry to 20 minutes ago.
const backdated = new Date(Date.now() - 20 * 60 * 1000).toISOString();
(db as unknown as RawDb).db
.prepare("UPDATE tasks SET execution_locked_at = ? WHERE id = ?")
.prepare("UPDATE task_leases SET expires_at = ? WHERE task_id = ?")
.run(backdated, taskId);
const recovered = db.recoverStaleLocks();
@@ -366,12 +385,7 @@ describe("CAS checkout", () => {
const task = db.core.entity.getTask(taskId);
expect(task.status).toBe("todo");
const row = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(row.checkout_run_id).toBeNull();
expect(row.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
});
it("stale lock recovery ignores tasks within 15 minute window", () => {
@@ -404,10 +418,8 @@ describe("CAS checkout", () => {
it("failed dispatch (run never started) rolls task back to 'todo' so heartbeat can retry", () => {
// Regression: workspace-lock contention calls finishRun('failed') on a
// run that never reached 'running'. The old code cleared checkout_run_id
// but left task.status='in_progress' — the next checkoutTask pass then
// refused (WHERE status IN ('todo','backlog')), and the scheduler
// deadlocked while logging "lost checkout race" every 15 seconds.
// run that never reached 'running'. The lease must clear and the task
// status must return to 'todo' so the heartbeat retry path can reclaim it.
const companyId = createTestCompany(db);
createTestConnector(db);
const workspaceId = createTestWorkspace(db, companyId);
+6 -6
View File
@@ -31,9 +31,9 @@ function makeStaleRunDb(runningCount: number): StaleRunDb & { calls: string[] }
blockInterruptedTasks: vi.fn((_ts: string) => {
calls.push("blockInterruptedTasks");
}) as unknown as StaleRunDb["blockInterruptedTasks"],
clearActiveRunIds: vi.fn(() => {
calls.push("clearActiveRunIds");
}) as unknown as StaleRunDb["clearActiveRunIds"],
clearInterruptedTaskLeases: vi.fn(() => {
calls.push("clearInterruptedTaskLeases");
}) as unknown as StaleRunDb["clearInterruptedTaskLeases"],
resetAgentsToIdle: vi.fn((_ts: string) => {
calls.push("resetAgentsToIdle");
}) as unknown as StaleRunDb["resetAgentsToIdle"],
@@ -66,16 +66,16 @@ describe("Crash Recovery Detection", () => {
expect(db.calls).toContain("blockInterruptedTasks");
});
it("does NOT mark tasks as blocked when no runs were interrupted", () => {
it("leaves task blocking step idle when no runs were interrupted", () => {
const db = makeStaleRunDb(0);
interruptStaleRunsLogic(db, transaction, nowIso);
expect(db.calls).not.toContain("blockInterruptedTasks");
});
it("always clears active_run_id and resets agents to idle", () => {
it("always clears interrupted task leases and resets agents to idle", () => {
const db = makeStaleRunDb(0);
interruptStaleRunsLogic(db, transaction, nowIso);
expect(db.calls).toContain("clearActiveRunIds");
expect(db.calls).toContain("clearInterruptedTaskLeases");
expect(db.calls).toContain("resetAgentsToIdle");
});
+66 -9
View File
@@ -1,6 +1,7 @@
import { join } from "node:path";
import { tmpdir } from "node:os";
import { mkdtemp, rm } from "node:fs/promises";
import { DatabaseSync } from "node:sqlite";
import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
beforeAll(() => {
@@ -17,6 +18,8 @@ vi.mock("electron", () => ({
}));
import { AppDatabase } from "../../src/main/database";
import { installBaselineSchema } from "../../src/main/database/migrations/baseline-v56";
import { markBaselineSchema } from "../../src/main/database/migrations/legacy-to-v56";
describe("AppDatabase", () => {
let tempDir: string;
@@ -33,18 +36,25 @@ describe("AppDatabase", () => {
await rm(tempDir, { recursive: true, force: true });
});
it("keeps task runtime state on the tasks row", () => {
const shadowTables = db.queryAll<{ name: string }>(
it("keeps task runtime ownership in task_leases", () => {
const leaseTables = db.queryAll<{ name: string }>(
"select name from sqlite_master where type = 'table' and name = 'task_leases'",
);
const phaseColumns = db.queryAll<{ name: string }>("pragma table_info(tasks)")
.filter((column) => column.name === "phase");
const removedTaskColumns = new Set([
"phase",
"active_run_id",
"checkout_run_id",
"execution_locked_at",
"active_session_id",
]);
const taskColumns = db.queryAll<{ name: string }>("pragma table_info(tasks)")
.filter((column) => removedTaskColumns.has(column.name));
const phaseTriggers = db.queryAll<{ name: string }>(
"select name from sqlite_master where type = 'trigger' and name like 'trg_tasks_phase_from_status_%'",
);
expect(shadowTables).toEqual([]);
expect(phaseColumns).toEqual([]);
expect(leaseTables).toEqual([{ name: "task_leases" }]);
expect(taskColumns).toEqual([]);
expect(phaseTriggers).toEqual([]);
});
@@ -462,11 +472,14 @@ describe("AppDatabase", () => {
expect(db.core.agents.isAgentPendingApproval(pendingAgentId)).toBe(true);
});
it("installs the v56 baseline once on a fresh profile", () => {
it("installs the v56 baseline and v57 task leases once on a fresh profile", () => {
const migrations = db.queryAll<{ version: number; name: string }>(
"select version, name from schema_migrations order by version",
);
expect(migrations).toEqual([{ version: 56, name: "baseline_v56" }]);
expect(migrations).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
db.close();
db = new AppDatabase(tempDir);
@@ -475,7 +488,51 @@ describe("AppDatabase", () => {
const migrationsAfterReinit = db.queryAll<{ version: number; name: string }>(
"select version, name from schema_migrations order by version",
);
expect(migrationsAfterReinit).toEqual([{ version: 56, name: "baseline_v56" }]);
expect(migrationsAfterReinit).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
});
it("migrates v56 task ownership columns into task_leases", async () => {
const legacyDir = await mkdtemp(join(tmpdir(), "agentcompany-v56-lease-"));
const rawDb = new DatabaseSync(join(legacyDir, "profile.sqlite"));
const ts = "2026-04-24T00:00:00.000Z";
try {
installBaselineSchema(rawDb);
markBaselineSchema(rawDb);
rawDb
.prepare("insert into companies (id, name, description, status, created_at, updated_at) values (?, ?, ?, ?, ?, ?)")
.run("company-1", "Lease Co", "", "active", ts, ts);
rawDb
.prepare("insert into agents (id, company_id, name, role, title, status, reports_to, connector_id, workspace_id, capabilities, created_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
.run("agent-1", "company-1", "Worker", "IC", "Worker", "active", null, "openai-codex", null, "Executes", ts, ts);
rawDb
.prepare("insert into tasks (id, company_id, title, description, assignee_agent_id, status, priority, active_run_id, checkout_run_id, execution_locked_at, created_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
.run("task-1", "company-1", "Lease task", "", "agent-1", "in_progress", "medium", "run-1", "run-1", ts, ts, ts);
rawDb
.prepare("insert into runs (id, company_id, task_id, agent_id, connector_id, status, summary, created_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)")
.run("run-1", "company-1", "task-1", "agent-1", "openai-codex", "queued", "", ts, ts);
} finally {
rawDb.close();
}
const legacyAppDb = new AppDatabase(legacyDir);
try {
legacyAppDb.init();
expect(
legacyAppDb.queryAll<{ taskId: string; runId: string; agentId: string }>(
"select task_id as taskId, run_id as runId, agent_id as agentId from task_leases",
),
).toEqual([{ taskId: "task-1", runId: "run-1", agentId: "agent-1" }]);
const columns = legacyAppDb.queryAll<{ name: string }>("pragma table_info(tasks)").map((column) => column.name);
expect(columns).not.toContain("active_run_id");
expect(columns).not.toContain("checkout_run_id");
expect(columns).not.toContain("execution_locked_at");
} finally {
legacyAppDb.close();
await rm(legacyDir, { recursive: true, force: true });
}
});
it("acquires and releases workspace locks", () => {
+21 -17
View File
@@ -23,6 +23,12 @@ type RawDb = {
};
};
type TaskLeaseRow = {
task_id: string;
run_id: string;
expires_at: string;
};
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
@@ -83,6 +89,14 @@ function createTestTask(db: AppDatabase, companyId: string, workspaceId: string)
});
}
function getTaskLease(db: AppDatabase, taskId: string): TaskLeaseRow | null {
return (
((db as unknown as RawDb).db
.prepare("SELECT task_id, run_id, expires_at FROM task_leases WHERE task_id = ?")
.get(taskId) as TaskLeaseRow | undefined) ?? null
);
}
// ---------------------------------------------------------------------------
// Tests — Edge Case: Concurrent Task Checkout (D-05.1)
// ---------------------------------------------------------------------------
@@ -234,15 +248,10 @@ describe("Edge Case: Concurrent Task Checkout (D-05.1)", () => {
});
expect(db.core.tasks.checkoutTask(taskId, runIdA)).toBe(true);
// Release the lock (sets checkout_run_id = NULL)
// Release the lease while leaving the task lifecycle unchanged.
db.core.tasks.releaseTaskLock(taskId, runIdA);
// Verify lock fields are cleared
const row = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(row.checkout_run_id).toBeNull();
expect(row.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
// Task is still in_progress though, so a new checkout fails
// (CAS checks status IN ('todo', 'backlog'))
@@ -273,9 +282,9 @@ describe("Edge Case: Concurrent Task Checkout (D-05.1)", () => {
});
expect(db.core.tasks.checkoutTask(taskId, runIdA)).toBe(true);
// Backdate execution_locked_at to well past the stale timeout
// Backdate the lease expiry past the stale timeout.
(db as unknown as RawDb).db
.prepare("UPDATE tasks SET execution_locked_at = '2020-01-01T00:00:00Z' WHERE id = ?")
.prepare("UPDATE task_leases SET expires_at = '2020-01-01T00:00:00Z' WHERE task_id = ?")
.run(taskId);
const recovered = db.recoverStaleLocks();
@@ -283,12 +292,7 @@ describe("Edge Case: Concurrent Task Checkout (D-05.1)", () => {
expect(recovered[0].taskId).toBe(taskId);
expect(recovered[0].companyId).toBe(companyId);
// Checkout fields should be cleared
const row = (db as unknown as RawDb).db
.prepare("SELECT checkout_run_id, execution_locked_at FROM tasks WHERE id = ?")
.get(taskId) as { checkout_run_id: string | null; execution_locked_at: string | null };
expect(row.checkout_run_id).toBeNull();
expect(row.execution_locked_at).toBeNull();
expect(getTaskLease(db, taskId)).toBeNull();
});
it("after stale recovery, different agent can checkout the recovered task", () => {
@@ -309,9 +313,9 @@ describe("Edge Case: Concurrent Task Checkout (D-05.1)", () => {
});
expect(db.core.tasks.checkoutTask(taskId, runIdA)).toBe(true);
// Backdate to make stale
// Backdate the lease to make it stale.
(db as unknown as RawDb).db
.prepare("UPDATE tasks SET execution_locked_at = '2020-01-01T00:00:00Z' WHERE id = ?")
.prepare("UPDATE task_leases SET expires_at = '2020-01-01T00:00:00Z' WHERE task_id = ?")
.run(taskId);
// Recover stale locks (sets status back to "todo")
@@ -171,6 +171,7 @@ describe("legacy to v56 baseline migration", () => {
try {
expect(appDb.queryAll("select version, name from schema_migrations")).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
expect(appDb.queryAll<{ companyId: string }>("select company_id as companyId from secrets where id = 'secret-1'")[0])
.toEqual({ companyId: "company-1" });
@@ -202,6 +203,7 @@ describe("legacy to v56 baseline migration", () => {
try {
expect(appDb.queryAll("select version, name from schema_migrations")).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
expect(appDb.queryAll("select id from budget_policies where scope_type = 'agent' and scope_id = 'agent-1'"))
.toHaveLength(1);
@@ -360,6 +362,7 @@ describe("legacy to v56 baseline migration", () => {
.toEqual([{ agentId: "agent-1", adapterType: "openai-codex" }]);
expect(appDb.queryAll("select version, name from schema_migrations")).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
} finally {
appDb.close();
@@ -115,8 +115,13 @@ describe("legacy to v56 baseline migration", () => {
{ id: "task-cognitive", requiresFilesystem: 0 },
{ id: "task-inflight", requiresFilesystem: 1 },
]);
const migration = appDb.queryAll<{ version: number }>("select version from schema_migrations where version = 56")[0];
expect(migration?.version).toBe(56);
const migrations = appDb.queryAll<{ version: number; name: string }>(
"select version, name from schema_migrations where version in (56, 57) order by version",
);
expect(migrations).toEqual([
{ version: 56, name: "baseline_v56" },
{ version: 57, name: "task_leases_runtime_ownership" },
]);
} finally {
appDb.close();
}