diff --git a/src/daemon/dispatch.ts b/src/daemon/dispatch.ts index de91225..3f1b34a 100644 --- a/src/daemon/dispatch.ts +++ b/src/daemon/dispatch.ts @@ -45,6 +45,21 @@ export interface DispatchConfig { * the SAME tick and burn its whole failure budget in milliseconds. */ retryBaseMs: number; + /** + * Stable identity (host:port) of the daemon that owns the tasks it enqueues, + * scoping claim and reclaim to its own work. + * + * `claim_owner` is a BOOT id, which cannot tell "my own crashed run" apart + * from "another daemon's healthy claim" — so two daemons sharing one + * database reclaimed each other's in-flight tasks and auto-blocked work that + * had actually succeeded, and could claim across tenants. Port-scoped for + * the same reason `local-mode.md` port-scopes its token file: the port is + * what distinguishes two daemons on one machine, and it survives restarts. + * + * Undefined leaves tasks unowned (claimable by anyone) — the pre-upgrade + * path, and the default in tests that run a single dispatcher. + */ + daemonId?: string; } export const DEFAULT_DISPATCH_CONFIG: DispatchConfig = { @@ -246,6 +261,7 @@ export class Dispatcher { id, ...input, failureLimit: this.#config.failureLimit, + ownerDaemon: this.#config.daemonId, now: Date.now(), }); this.#host.audit( @@ -310,6 +326,7 @@ export class Dispatcher { createdBy: input.createdBy, groupId, groupOrdinal: i + 1, + ownerDaemon: this.#config.daemonId, now, }; }); @@ -409,6 +426,7 @@ export class Dispatcher { this.bootId, this.#config.leaseMs, Date.now(), + this.#config.daemonId, ); for (const task of reclaimed) { if (task.workerSessionId) this.#unwatch(task.workerSessionId, task.id); @@ -460,7 +478,12 @@ export class Dispatcher { // starve the rest of the queue). const touched: string[] = []; for (;;) { - const task = this.#store.dispatchClaimNext(this.bootId, Date.now(), touched); + const task = this.#store.dispatchClaimNext( + this.bootId, + Date.now(), + touched, + this.#config.daemonId, + ); if (!task) return; touched.push(task.id); if ( diff --git a/src/daemon/server.ts b/src/daemon/server.ts index 2b1400f..eeca58c 100644 --- a/src/daemon/server.ts +++ b/src/daemon/server.ts @@ -307,7 +307,15 @@ export class DaemonServer { this.#store, this.#transcriptStore, identityManager, rateLimiter, // Memory is wired post-construction via initMemory() — see start() undefined, - { config: config.fullConfig, compressionRegistry, hooks }, + { + config: config.fullConfig, + compressionRegistry, + hooks, + // Port-scoped, for the reason local mode port-scopes its token file: + // it is what distinguishes two daemons on one machine, and it survives + // restarts so a daemon still recognises its own crashed run. + daemonId: `${this.#config.host}:${this.#config.port}`, + }, ); console.log(`[codeoid] providers: ${this.#manager.providerIds().join(", ")}`); diff --git a/src/daemon/session-manager.ts b/src/daemon/session-manager.ts index b5fa542..69812a2 100644 --- a/src/daemon/session-manager.ts +++ b/src/daemon/session-manager.ts @@ -414,6 +414,13 @@ export class SessionManager { * Absent = no hooks; sessions pay zero overhead. */ hooks?: HookBus; + /** + * Stable identity of this daemon (host:port), scoping dispatch claims to + * its own tasks. Two daemons sharing one database otherwise reclaim each + * other's live claims and can claim across tenants — see + * `DispatchConfig.daemonId`. Absent in tests = unowned tasks. + */ + daemonId?: string; /** * Test-only: provider factory injected into every Session this manager * constructs, so manager-level integration tests (conductor injection, @@ -433,11 +440,10 @@ export class SessionManager { this.#providers = opts?.providers ?? createDefaultProviderRegistry(opts?.config); this.#hooks = opts?.hooks; this.#testProviderFactory = opts?._testProviderFactory; - this.#dispatcher = new Dispatcher( - store, - this.#makeDispatcherHost(), - opts?.config?.dispatch, - ); + this.#dispatcher = new Dispatcher(store, this.#makeDispatcherHost(), { + ...opts?.config?.dispatch, + daemonId: opts?.daemonId, + }); this.#pushService = new PushService( store, createPushTransport(opts?.config?.push, (token) => store.pruneDeadToken(token)), diff --git a/src/daemon/store.ts b/src/daemon/store.ts index 76e15aa..ea0e0eb 100644 --- a/src/daemon/store.ts +++ b/src/daemon/store.ts @@ -93,6 +93,7 @@ interface RawDispatchRow { failure_limit: number; claim_owner: string | null; claimed_at: number | null; + owner_daemon: string | null; not_before: number | null; worker_session_id: string | null; result_digest: string | null; @@ -490,6 +491,18 @@ export class Store { this.#addColumnIfMissing("dispatch_tasks", "group_id", "TEXT"); this.#addColumnIfMissing("dispatch_tasks", "group_ordinal", "INTEGER"); + // Owning daemon (host:port) — the executor scope for claim + reclaim. + // `claim_owner` is a BOOT id, which cannot distinguish "my own crashed + // run" from "another daemon's healthy claim". Two daemons sharing this + // database (a supported setup — local-mode.md port-scopes its token file + // for exactly that) therefore stole each other's in-flight tasks and + // auto-blocked them, and could claim across tenants. NULL = a pre-upgrade + // row, claimable by any daemon so nothing is stranded by the migration. + this.#addColumnIfMissing("dispatch_tasks", "owner_daemon", "TEXT"); + this.#db.exec( + `CREATE INDEX IF NOT EXISTS idx_dispatch_owner + ON dispatch_tasks(owner_daemon, status, created_at);`, + ); // AFTER the ALTERs, never inside the CREATE block above. On an existing // database `CREATE TABLE IF NOT EXISTS` is a no-op, so an index declared // there would run against a table that does not have the column yet and @@ -939,14 +952,20 @@ export class Store { /** 1-based position within that group. */ groupOrdinal?: number; now: number; + /** + * The daemon that may claim and execute this task (host:port). Omit only + * in tests that never run two dispatchers — a NULL owner is claimable by + * any daemon, which is the pre-upgrade compatibility path. + */ + ownerDaemon?: string; }): void { this.#db .prepare( `INSERT INTO dispatch_tasks (id, account_id, project_id, kind, shape, target_session, workdir, prompt, provider, model, failure_limit, created_by, group_id, - group_ordinal, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + group_ordinal, owner_daemon, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, ) .run( task.id, @@ -963,6 +982,7 @@ export class Store { task.createdBy, task.groupId ?? null, task.groupOrdinal ?? null, + task.ownerDaemon ?? null, task.now, task.now, ); @@ -981,6 +1001,7 @@ export class Store { bootId: string, now: number, excludeIds: readonly string[] = [], + ownerDaemon?: string, ): DispatchTaskRow | null { const row = this.#db .prepare( @@ -990,12 +1011,20 @@ export class Store { SELECT id FROM dispatch_tasks WHERE status = 'queued' AND (not_before IS NULL OR not_before <= ?) + AND (owner_daemon IS NULL OR owner_daemon IS ?) AND id NOT IN (SELECT value FROM json_each(?)) ORDER BY created_at LIMIT 1 ) RETURNING *`, ) - .get(bootId, now, now, now, JSON.stringify(excludeIds)) as RawDispatchRow | null; + .get( + bootId, + now, + now, + now, + ownerDaemon ?? null, + JSON.stringify(excludeIds), + ) as RawDispatchRow | null; return row ? rowToDispatchTask(row) : null; } @@ -1074,7 +1103,21 @@ export class Store { * stuck-loop escalation, in queue form. worker_session_id is preserved so * a re-claimed spawn can continue its (resumed) worker session. */ - dispatchReclaimStale(bootId: string, leaseMs: number, now: number): DispatchTaskRow[] { + /** + * Reclaim claims this daemon can prove are dead. + * + * The `owner_daemon` predicate is load-bearing: `claim_owner` is a BOOT id, + * so "not my boot" matches another *live* daemon's claim just as readily as + * my own crashed one. Scoping to this daemon's own tasks first is what makes + * the boot-id fast path (restart recovery on the first tick) correct rather + * than theft. A row with no owner is pre-upgrade and stays lease-governed. + */ + dispatchReclaimStale( + bootId: string, + leaseMs: number, + now: number, + ownerDaemon?: string, + ): DispatchTaskRow[] { const rows = this.#db .prepare( `UPDATE dispatch_tasks @@ -1083,10 +1126,24 @@ export class Store { claim_owner = NULL, claimed_at = NULL, error = COALESCE(error, 'reclaimed: stale claim'), updated_at = ? WHERE status IN ('claimed', 'running') - AND (claim_owner IS NOT ? OR claimed_at IS NULL OR claimed_at + ? < ?) + AND ( + -- Mine: a previous boot of THIS daemon crashed, or the lease ran out. + (owner_daemon IS ? AND (claim_owner IS NOT ? OR claimed_at IS NULL OR claimed_at + ? < ?)) + -- Unowned (pre-upgrade): lease expiry only — never on boot id, which + -- is what let another daemon steal a healthy claim. + OR (owner_daemon IS NULL AND (claimed_at IS NULL OR claimed_at + ? < ?)) + ) RETURNING *`, ) - .all(now, bootId, leaseMs, now) as RawDispatchRow[]; + .all( + now, + ownerDaemon ?? null, + bootId, + leaseMs, + now, + leaseMs, + now, + ) as RawDispatchRow[]; return rows.map(rowToDispatchTask); } diff --git a/src/tests/dispatch-daemon-scope.test.ts b/src/tests/dispatch-daemon-scope.test.ts new file mode 100644 index 0000000..353e433 --- /dev/null +++ b/src/tests/dispatch-daemon-scope.test.ts @@ -0,0 +1,136 @@ +/** + * Dispatch ownership scope — two daemons, one database. + * + * `claim_owner` is a daemon BOOT id, so the reclaim predicate's "not my boot" + * arm cannot tell *my own crashed run* apart from *another daemon's healthy + * claim*. Running a second daemon against the same `~/.codeoid/codeoid.db` is + * a supported setup — local-mode.md port-scopes its token file precisely so + * two daemons can coexist on one machine — and before `owner_daemon` existed + * the two dispatchers stole each other's in-flight tasks: + * + * - a healthy worker's task was reclaimed within a tick or two of a + * 10-MINUTE lease, hit `failure_limit`, and auto-BLOCKED, so the conductor + * was told the work had failed while the worker was reporting success; and + * - `dispatchClaimNext` had no ownership predicate at all, so one daemon + * could claim and execute another tenant's queued task. + * + * Both were observed end-to-end against a real daemon before this fix. + */ + +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import { Store } from "../daemon/store.js"; + +const LEASE_MS = 10 * 60_000; + +const DAEMON_A = "127.0.0.1:7400"; +const DAEMON_B = "127.0.0.1:7455"; +const BOOT_A = "boot-a"; +const BOOT_A2 = "boot-a-after-restart"; +const BOOT_B = "boot-b"; + +let dir: string; +let store: Store; + +function enqueue(ownerDaemon: string | undefined, now: number, accountId = "acc"): string { + const id = randomUUID(); + store.dispatchEnqueue({ + id, + accountId, + projectId: "proj", + kind: "spawn", + shape: "scout", + workdir: "/tmp", + prompt: "investigate", + failureLimit: 2, + createdBy: "conductor:test", + ownerDaemon, + now, + }); + return id; +} + +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), "dispatch-scope-")); + store = new Store(join(dir, "codeoid.db")); +}); + +afterEach(() => { + store.close(); + rmSync(dir, { recursive: true, force: true }); +}); + +describe("dispatch ownership across two daemons sharing a database", () => { + test("a second daemon does not reclaim a live claim (the auto-block bug)", () => { + const now = Date.now(); + const id = enqueue(DAEMON_A, now); + + expect(store.dispatchClaimNext(BOOT_A, now, [], DAEMON_A)?.id).toBe(id); + + // Daemon B's ordinary tick, one second in — 10 minutes of lease remain. + expect(store.dispatchReclaimStale(BOOT_B, LEASE_MS, now + 1_000, DAEMON_B)).toEqual([]); + + // A second tick must not accumulate attempts either: two of these used to + // be enough to reach failure_limit and auto-block healthy work. + expect(store.dispatchReclaimStale(BOOT_B, LEASE_MS, now + 6_000, DAEMON_B)).toEqual([]); + + const task = store.dispatchGet(id); + expect(task?.status).toBe("claimed"); + expect(task?.attempts).toBe(0); + expect(task?.error).toBeNull(); + }); + + test("a daemon still reclaims its OWN previous boot on the first tick", () => { + const now = Date.now(); + const id = enqueue(DAEMON_A, now); + store.dispatchClaimNext(BOOT_A, now, [], DAEMON_A); + + // Same daemon (same host:port), new boot id — a crash and restart. Fast + // recovery must survive the ownership scoping, well inside the lease. + const reclaimed = store.dispatchReclaimStale(BOOT_A2, LEASE_MS, now + 1_000, DAEMON_A); + expect(reclaimed.map((t) => t.id)).toEqual([id]); + expect(store.dispatchGet(id)?.status).toBe("queued"); + expect(store.dispatchGet(id)?.attempts).toBe(1); + }); + + test("a daemon cannot claim another daemon's queued task", () => { + const now = Date.now(); + const mine = enqueue(DAEMON_A, now, "acc-a"); + + // B has nothing of its own; the only queued row belongs to A. + expect(store.dispatchClaimNext(BOOT_B, now, [], DAEMON_B)).toBeNull(); + expect(store.dispatchGet(mine)?.status).toBe("queued"); + + // A still gets it. + expect(store.dispatchClaimNext(BOOT_A, now, [], DAEMON_A)?.id).toBe(mine); + }); + + test("an expired lease is still reclaimed by its owner", () => { + const now = Date.now(); + const id = enqueue(DAEMON_A, now); + store.dispatchClaimNext(BOOT_A, now, [], DAEMON_A); + + // Same boot, but the worker wedged and stopped renewing: lease expiry is + // the backstop and must keep working. + const reclaimed = store.dispatchReclaimStale(BOOT_A, LEASE_MS, now + LEASE_MS + 1, DAEMON_A); + expect(reclaimed.map((t) => t.id)).toEqual([id]); + }); + + test("pre-upgrade rows (no owner) stay claimable, and reclaim only on lease expiry", () => { + const now = Date.now(); + const id = enqueue(undefined, now); + + // Claimable by whoever is running — nothing is stranded by the migration. + expect(store.dispatchClaimNext(BOOT_A, now, [], DAEMON_A)?.id).toBe(id); + + // But an unowned row must NOT be reclaimed on the boot-id arm, or the + // original theft returns for exactly the rows that predate the fix. + expect(store.dispatchReclaimStale(BOOT_B, LEASE_MS, now + 1_000, DAEMON_B)).toEqual([]); + expect( + store.dispatchReclaimStale(BOOT_B, LEASE_MS, now + LEASE_MS + 1, DAEMON_B).map((t) => t.id), + ).toEqual([id]); + }); +});