diff --git a/packages/storage/src/__tests__/scheduled-task-row-operations.test.ts b/packages/storage/src/__tests__/scheduled-task-row-operations.test.ts new file mode 100644 index 0000000000..8c94a0490e --- /dev/null +++ b/packages/storage/src/__tests__/scheduled-task-row-operations.test.ts @@ -0,0 +1,605 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { after, test, type TestContext } from 'node:test'; +import type { DatabaseSync, SQLInputValue } from 'node:sqlite'; +import { + acquireOperationalStateDatabase, + type OperationalStateDatabaseLease, +} from '../operational-state-store.js'; +import { + openInteractiveScheduledTaskStoreForWrite, + ScheduledTaskStoreError, + type InteractiveScheduledTaskStoreWriter, +} from '../scheduled-task-store.js'; +import { + resolveStorageRoot, + runWithStorageRootLease, + tryAcquireInteractiveRootOwner, +} from '../root-authority.js'; +import { + removeTrackedControlDirectories, + trackControlDirectory, +} from './fixtures/control-directory-hygiene.js'; + +after(removeTrackedControlDirectories); + +const NOW = 1_000_000; +const EXECUTION = { + sessionId: 'session-target', + turnId: 'turn-target', + runId: 'run-target', + userMessageId: 'message-target', +}; + +test('ScheduledTask point operations do not materialize or rewrite unrelated rows', async (t) => { + for (const unrelatedCount of [0, 32, 256]) { + await t.test(`${unrelatedCount} unrelated tasks and pending fires`, async (t) => { + await withStore(t, async ({ store, probe }) => { + // Setup goes through the actual owner and public store, outside measurement. + // Each unrelated task has a pending claim, so scanning either table fails. + for (let index = 0; index < unrelatedCount; index += 1) { + const other = await store.create(notifyInput(`Unrelated ${index}`), NOW); + await store.claimNow(other.id, NOW); + } + const target = await store.create(agentInput(), NOW); + const unchanged = probe.snapshotExcluding(target.id); + + const read = await probe.measure(() => store.get(target.id)); + assert.equal(read.value?.id, target.id); + assertPointCost(read.cost, { rows: 1, payloadRows: 1, changes: 0 }); + assert.ok( + read.cost.reads.every((read) => !read.sql.includes('workflow_scheduled_task_fires')), + ); + + const snoozed = await probe.measure(() => store.snooze(target.id, 1_000, NOW)); + assert.equal(snoozed.value.nextFireAt, target.nextFireAt! + 1_000); + assertPointCost(snoozed.cost, { rows: 2, payloadRows: 1, changes: 1 }); + + const paused = await probe.measure(() => store.pause(target.id, NOW + 1)); + assert.equal(paused.value.status, 'paused'); + assertPointCost(paused.cost, { rows: 2, payloadRows: 1, changes: 1 }); + await store.resume(target.id, NOW + 2); + + const claimed = await probe.measure(() => store.claimNow(target.id, NOW + 3)); + assert.equal(claimed.value.taskId, target.id); + assertPointCost(claimed.cost, { rows: 2, payloadRows: 1, changes: 1 }); + + const bound = await probe.measure(() => + store.bindFireExecution(claimed.value.id, EXECUTION), + ); + assert.deepEqual(bound.value.execution, EXECUTION); + assertPointCost(bound.cost, { rows: 1, payloadRows: 1, changes: 1 }); + + const settled = await probe.measure(() => + store.settleFire(claimed.value.id, { + id: 'run-receipt-target', + at: NOW + 4, + outcome: 'ok', + message: 'done', + sessionId: EXECUTION.sessionId, + runId: EXECUTION.runId, + }), + ); + assert.equal(settled.value.fireCount, 1); + assertPointCost(settled.cost, { rows: 2, payloadRows: 2, changes: 2 }); + assert.deepEqual(probe.snapshotExcluding(target.id), unchanged); + }); + }); + } +}); + +test('ScheduledTask polls without new due or expired tasks perform no DML', async (t) => { + await withStore(t, async ({ store, probe }) => { + for (let index = 0; index < 32; index += 1) { + await store.create(notifyInput(`Future ${index}`), NOW); + } + const due = await store.create( + { ...notifyInput('Already claimed'), schedule: { kind: 'once', runAt: NOW + 1 } }, + NOW, + ); + await store.claimNow(due.id, NOW); + const before = probe.snapshotExcluding(); + + for (let index = 0; index < 3; index += 1) { + const poll = await probe.measure(() => store.claimNextDue(NOW + 2)); + assert.deepEqual(poll.value, { claim: null, expired: [] }); + assert.equal(poll.cost.writes.length, 0, 'an unchanged poll must issue no DML'); + assert.equal(poll.cost.totalChanges, 0); + } + assert.deepEqual(probe.snapshotExcluding(), before); + }); +}); + +test('ScheduledTask expiry updates only newly expired tasks, including a pending task', async (t) => { + await withStore(t, async ({ store, probe }) => { + const expiringInput = { + ...notifyInput('Expired with pending fire'), + schedule: { kind: 'interval', everySeconds: 60, startAt: NOW + 1_000 }, + expiresAt: NOW + 2_000, + }; + const pending = await store.create(expiringInput, NOW); + const pendingClaim = await store.claimNow(pending.id, NOW); + const otherExpired = await store.create({ ...expiringInput, title: 'Other expired task' }, NOW); + const due = await store.create( + { ...notifyInput('Due'), schedule: { kind: 'once', runAt: NOW + 1_500 } }, + NOW, + ); + for (let index = 0; index < 32; index += 1) { + await store.create(notifyInput(`Unchanged ${index}`), NOW); + } + + const first = await probe.measure(() => store.claimNextDue(NOW + 2_000)); + assert.deepEqual( + first.value.expired.map((task) => task.id).sort(), + [pending.id, otherExpired.id].sort(), + ); + assert.equal(first.value.claim?.taskId, due.id); + assert.equal(first.cost.totalChanges, 3, 'two task updates and one new claim'); + assert.equal( + first.cost.writes.reduce((count, write) => count + write.changes, 0), + 3, + ); + assert.ok((await store.listPendingFires()).some((claim) => claim.id === pendingClaim.id)); + + const second = await probe.measure(() => store.claimNextDue(NOW + 2_000)); + assert.deepEqual(second.value, { claim: null, expired: [] }); + assert.equal(second.cost.writes.length, 0); + assert.equal(second.cost.totalChanges, 0); + + // Even though pause(expired) otherwise does nothing, a pending fire still + // takes precedence and must reject the mutation. + const rejectedPause = await probe.measure(() => + assert.rejects(() => store.pause(pending.id, NOW + 3_000), isOperationConflict), + ); + assertNoDml(rejectedPause.cost); + }); +}); + +test('ScheduledTask native delivery allows waiting cancellation but cannot undo admission', async (t) => { + await withStore(t, async ({ store, probe }) => { + const task = await store.create(notifyInput('Native notification'), NOW); + const waiting = await store.claimNow(task.id, NOW); + await store.setFireNativeState(waiting.id, 'waiting_for_provider'); + const cancelled = await probe.measure(() => store.cancelWaitingNativeFire(task.id)); + assert.equal(cancelled.value, true); + assert.equal(cancelled.cost.totalChanges, 1); + assert.deepEqual(await store.listPendingFires(), []); + assert.equal((await store.get(task.id))?.fireCount, 0); + + const invoking = await store.claimNow(task.id, NOW + 1); + await store.setFireNativeState(invoking.id, 'waiting_for_provider'); + await store.setFireNativeState(invoking.id, 'invoking'); + const admitted = probe.snapshotExcluding(); + const rejections = await probe.measure(async () => { + await assert.rejects(() => store.cancelWaitingNativeFire(task.id), isOperationConflict); + await assert.rejects( + () => store.setFireNativeState(invoking.id, 'waiting_for_provider'), + isOperationConflict, + ); + await store.setFireNativeState(invoking.id, 'invoking'); + }); + assertNoDml(rejections.cost); + assert.deepEqual(probe.snapshotExcluding(), admitted); + + // The same writer queue must remain usable after both rejected operations. + const settled = await store.settleFire(invoking.id, { + at: NOW + 2, + outcome: 'failed', + message: 'Delivery outcome was not observed.', + }); + assert.equal(settled.fireCount, 1); + assert.deepEqual(await store.listPendingFires(), []); + }); +}); + +test('ScheduledTask execution binding is idempotent and does not retain caller-owned objects', async (t) => { + await withStore(t, async ({ store, probe }) => { + const task = await store.create(agentInput(), NOW); + const claim = await store.claimNow(task.id, NOW); + const input = { ...EXECUTION }; + const bound = await store.bindFireExecution(claim.id, input); + const persisted = probe.snapshotExcluding(); + input.runId = 'mutated-input-run'; + assert.ok(bound.execution); + bound.execution.userMessageId = 'mutated-return-message'; + bound.task.title = 'Mutated returned task'; + assert.deepEqual(probe.snapshotExcluding(), persisted); + + const repeated = await probe.measure(() => store.bindFireExecution(claim.id, EXECUTION)); + assert.deepEqual(repeated.value.execution, EXECUTION); + assert.equal(repeated.value.task.title, task.title); + assertNoDml(repeated.cost); + const conflict = await probe.measure(() => + assert.rejects( + () => store.bindFireExecution(claim.id, { ...EXECUTION, runId: 'another-run' }), + isOperationConflict, + ), + ); + assertNoDml(conflict.cost); + assert.deepEqual(probe.snapshotExcluding(), persisted); + }); +}); + +test('ScheduledTask metadata updates keep schedule and expired-trigger semantics', async (t) => { + await withStore(t, async ({ store, probe }) => { + const task = await store.create(notifyInput('Original title'), NOW); + assert.equal(task.nextFireAt, NOW + 60_000); + const updated = await store.update(task.id, { title: 'New title' }, NOW + 61_000); + assert.equal(updated.title, 'New title'); + assert.equal(updated.nextFireAt, NOW + 120_000); + + const paused = await store.pause(task.id, NOW + 61_001); + const repeatedPause = await probe.measure(() => store.pause(task.id, NOW + 61_002)); + assert.deepEqual(repeatedPause.value, paused); + assertNoDml(repeatedPause.cost); + + const expiring = await store.create( + { + ...notifyInput('Expired trigger'), + schedule: { kind: 'interval', everySeconds: 60, startAt: NOW + 1_000 }, + expiresAt: NOW + 2_000, + }, + NOW, + ); + const before = probe.snapshotExcluding(); + const rejected = await probe.measure(() => + assert.rejects(() => store.claimNow(expiring.id, NOW + 2_000), isOperationConflict), + ); + assertNoDml(rejected.cost); + assert.deepEqual(probe.snapshotExcluding(), before); + assert.equal((await store.get(expiring.id))?.status, 'active'); + }); +}); + +test('ScheduledTask due discovery rejects a damaged task identity before changing another task', async (t) => { + await withStore(t, async ({ store, probe }) => { + const expiring = await store.create( + { + ...notifyInput('Expiring task'), + schedule: { kind: 'interval', everySeconds: 60, startAt: NOW + 1_000 }, + expiresAt: NOW + 2_000, + }, + NOW, + ); + const future = await store.create(notifyInput('Unrelated future task'), NOW); + // Both records came from the public API. This single-field corruption is a + // fault injection: the expiry write must not follow a damaged JSON identity. + probe.damageTaskIdentity(expiring.id, future.id); + const damaged = probe.snapshotExcluding(); + const rejected = await probe.measure(() => + assert.rejects(() => store.claimNextDue(NOW + 2_000), /Invalid scheduled task identity/), + ); + assertNoDml(rejected.cost); + assert.deepEqual(probe.snapshotExcluding(), damaged); + assert.equal((await store.get(future.id))?.status, 'active'); + }); +}); + +test('ScheduledTask settlement rolls back both rows and the queue accepts a retry', async (t) => { + await withStore(t, async ({ store, probe }) => { + const task = await store.create(agentInput(), NOW); + const claim = await store.claimNow(task.id, NOW); + await store.bindFireExecution(claim.id, EXECUTION); + const before = probe.snapshotExcluding(); + const run = { + id: 'rollback-receipt', + at: NOW + 1, + outcome: 'ok' as const, + message: 'durable settlement', + }; + + // Execute the task mutation normally, then fail before deleting the claim. + // This exercises SQLite rollback, not an early rejection in the store facade. + probe.failNextClaimDelete(); + await assert.rejects(() => store.settleFire(claim.id, run), /injected claim-delete failure/); + assert.equal(probe.failedAfterTaskWrite, true); + assert.deepEqual(probe.snapshotExcluding(), before); + + const retried = await store.settleFire(claim.id, run); + assert.equal(retried.fireCount, 1); + assert.equal(retried.runs.filter((item) => item.id === run.id).length, 1); + assert.equal((await store.listPendingFires()).length, 0); + }); +}); + +test('ScheduledTask execution identity survives closing and reacquiring the root owner', async (t) => { + await withStore(t, async (fixture) => { + const task = await fixture.store.create(agentInput(), NOW); + const claim = await fixture.store.claimNow(task.id, NOW); + await fixture.store.bindFireExecution(claim.id, EXECUTION); + const oldWriter = fixture.store; + + await fixture.reopen(); + await assert.rejects(() => oldWriter.get(task.id), /writer is closed/); + assert.notEqual(fixture.store, oldWriter); + assert.deepEqual((await fixture.store.listPendingFires())[0]?.execution, EXECUTION); + assert.equal((await fixture.store.listPendingFires())[0]?.id, claim.id); + await fixture.store.settleFire(claim.id, { + at: NOW + 1, + outcome: 'ok', + message: 'after reopening', + sessionId: EXECUTION.sessionId, + runId: EXECUTION.runId, + }); + + await fixture.reopen(); + assert.deepEqual(await fixture.store.listPendingFires(), []); + const settled = await fixture.store.get(task.id); + assert.equal(settled?.fireCount, 1); + assert.equal(settled?.runs[0]?.runId, EXECUTION.runId); + }); +}); + +function notifyInput(title: string) { + return { + title, + intentBody: '', + schedule: { kind: 'interval', everySeconds: 60, startAt: NOW + 60_000 }, + effect: { kind: 'notify', channel: 'local' }, + createdBy: { kind: 'user' }, + }; +} + +function agentInput() { + return { + ...notifyInput('Target task'), + intentBody: 'Perform the scheduled work.', + effect: { + kind: 'agent_run', + execution: { + cwd: '/workspace', + llmConnectionId: 'connection-default', + llmConnectionSlug: 'default', + model: 'test-model', + permissionMode: 'ask', + collaborationMode: 'agent', + orchestrationMode: 'default', + }, + }, + }; +} + +interface QueryRead { + sql: string; + parameters: SQLInputValue[]; + rows: number; + payloadRows: number; + payloadBytes: number; + plan: string[]; +} + +interface OperationCost { + reads: QueryRead[]; + writes: Array<{ sql: string; changes: number }>; + totalChanges: number; +} + +function isOperationConflict(error: unknown): boolean { + return error instanceof ScheduledTaskStoreError && error.code === 'operation_conflict'; +} + +function assertNoDml(cost: OperationCost): void { + assert.equal(cost.writes.length, 0); + assert.equal(cost.totalChanges, 0); +} + +function assertPointCost( + cost: OperationCost, + expected: { rows: number; payloadRows: number; changes: number }, +): void { + const returnedRows = cost.reads.reduce((count, read) => count + read.rows, 0); + const payloadRows = cost.reads.reduce((count, read) => count + read.payloadRows, 0); + const payloadBytes = cost.reads.reduce((count, read) => count + read.payloadBytes, 0); + assert.ok(returnedRows <= expected.rows, JSON.stringify({ returnedRows, cost })); + assert.ok(payloadRows <= expected.payloadRows, JSON.stringify({ payloadRows, cost })); + assert.ok(payloadBytes < 8 * 1024, JSON.stringify({ payloadBytes, cost })); + assert.equal(cost.totalChanges, expected.changes, JSON.stringify(cost)); + assert.equal( + cost.writes.reduce((count, write) => count + write.changes, 0), + expected.changes, + JSON.stringify(cost), + ); + assert.ok(cost.reads.length > 0, 'the probe must observe the actual point read'); + for (const read of cost.reads) { + assert.ok( + read.plan.some((line) => /SEARCH .*USING .*INDEX/u.test(line)), + JSON.stringify(read), + ); + assert.ok( + read.plan.every((line) => !/\bSCAN\b/u.test(line)), + JSON.stringify(read), + ); + } +} + +class SqlProbe { + readonly #prepare: DatabaseSync['prepare']; + #cost: OperationCost | undefined; + #failDelete = false; + #taskWritten = false; + failedAfterTaskWrite = false; + + constructor(t: TestContext, database: DatabaseSync) { + this.#prepare = database.prepare.bind(database); + t.mock.method(database, 'prepare', (sql: string) => { + const statement = this.#prepare(sql); + const relevant = /\bworkflow_scheduled_task(?:s|_fires)\b/u.test(sql); + if (!relevant) return statement; + return new Proxy(statement, { + get: (target, key) => { + const value: unknown = Reflect.get(target, key, target); + if (typeof value !== 'function') return value; + if (!['all', 'get', 'iterate', 'run'].includes(String(key))) return value.bind(target); + return (...parameters: SQLInputValue[]) => { + const isWrite = /^\s*(?:INSERT|UPDATE|DELETE)\b/iu.test(sql); + if ( + this.#failDelete && + /^\s*DELETE\s+FROM\s+workflow_scheduled_task_fires\b/iu.test(sql) + ) { + this.#failDelete = false; + this.failedAfterTaskWrite = this.#taskWritten; + throw new Error('injected claim-delete failure'); + } + const result: unknown = Reflect.apply(value, target, parameters); + if (isWrite && key === 'run') { + const changes = Number((result as { changes: number | bigint }).changes); + if (/\bworkflow_scheduled_tasks\b/u.test(sql) && changes > 0) { + this.#taskWritten = true; + } + this.#cost?.writes.push({ sql, changes }); + } else if (this.#cost && /^\s*SELECT\b/iu.test(sql)) { + const read: QueryRead = { + sql, + parameters, + rows: 0, + payloadRows: 0, + payloadBytes: 0, + plan: [], + }; + this.#cost.reads.push(read); + const record = (row: unknown) => { + if (row === undefined) return; + read.rows += 1; + const json = (row as { record_json?: unknown }).record_json; + if (typeof json === 'string') { + read.payloadRows += 1; + read.payloadBytes += Buffer.byteLength(json, 'utf8'); + } + }; + if (key === 'iterate') { + return (function* () { + for (const row of result as Iterable) { + record(row); + yield row; + } + })(); + } + if (key === 'all') { + for (const row of result as unknown[]) record(row); + } else if (key === 'get') record(result); + } + return result; + }; + }, + }); + }); + } + + async measure(operation: () => Promise): Promise<{ value: T; cost: OperationCost }> { + const before = this.#totalChanges(); + const cost: OperationCost = { reads: [], writes: [], totalChanges: 0 }; + this.#cost = cost; + try { + const value = await operation(); + cost.totalChanges = this.#totalChanges() - before; + for (const read of cost.reads) { + read.plan = this.#prepare(`EXPLAIN QUERY PLAN ${read.sql}`) + .all(...read.parameters) + .map((row) => String(row.detail)); + } + return { value, cost }; + } finally { + this.#cost = undefined; + } + } + + failNextClaimDelete(): void { + this.#taskWritten = false; + this.#failDelete = true; + this.failedAfterTaskWrite = false; + } + + damageTaskIdentity(taskId: string, replacementId: string): void { + const result = this.#prepare( + "UPDATE workflow_scheduled_tasks SET record_json = json_set(record_json, '$.id', ?) WHERE task_id = ?", + ).run(replacementId, taskId); + assert.equal(result.changes, 1); + } + + snapshotExcluding(taskId = ''): unknown { + return { + tasks: this.#prepare( + 'SELECT * FROM workflow_scheduled_tasks WHERE task_id <> ? ORDER BY task_id', + ).all(taskId), + claims: this.#prepare( + 'SELECT * FROM workflow_scheduled_task_fires WHERE task_id <> ? ORDER BY claim_id', + ).all(taskId), + }; + } + + #totalChanges(): number { + return Number(this.#prepare('SELECT total_changes() AS count').get()?.count); + } +} + +interface Fixture { + store: InteractiveScheduledTaskStoreWriter; + probe: SqlProbe; + reopen(): Promise; +} + +async function withStore(t: TestContext, run: (fixture: Fixture) => Promise): Promise { + const root = await mkdtemp(join(tmpdir(), 'maka-scheduled-task-rows-')); + const capability = trackControlDirectory( + await resolveStorageRoot({ path: root, kind: 'interactive' }), + ); + let owner = await tryAcquireInteractiveRootOwner(capability); + assert.ok(owner); + let lease: OperationalStateDatabaseLease = await runWithStorageRootLease( + owner.lease, + 'interactive', + 'write', + async (canonicalRoot) => acquireOperationalStateDatabase(canonicalRoot), + ); + const probe = new SqlProbe(t, lease.database); + let store = await openInteractiveScheduledTaskStoreForWrite(owner.lease); + const fixture: Fixture = { + store, + probe, + async reopen() { + store.close(); + lease.close(); + await owner!.close(); + owner = await tryAcquireInteractiveRootOwner(capability); + assert.ok(owner); + lease = await runWithStorageRootLease( + owner.lease, + 'interactive', + 'write', + async (canonicalRoot) => acquireOperationalStateDatabase(canonicalRoot), + ); + fixture.probe = new SqlProbe(t, lease.database); + store = await openInteractiveScheduledTaskStoreForWrite(owner.lease); + fixture.store = store; + }, + }; + try { + await run(fixture); + } finally { + store.close(); + lease.close(); + await owner?.close(); + await rm(root, { recursive: true, force: true }); + } +} diff --git a/packages/storage/src/scheduled-task-store.ts b/packages/storage/src/scheduled-task-store.ts index 9f28ec3dcc..8409ea104e 100644 --- a/packages/storage/src/scheduled-task-store.ts +++ b/packages/storage/src/scheduled-task-store.ts @@ -123,11 +123,6 @@ export interface InteractiveScheduledTaskStoreWriter extends ScheduledTaskStore readonly [writerBrand]: true; } -interface ScheduledTaskStoreState { - tasks: ScheduledTask[]; - claims: ScheduledTaskFireClaim[]; -} - export function authenticateInteractiveScheduledTaskStoreWriter( writer: InteractiveScheduledTaskStoreWriter, ): InteractiveScheduledTaskStoreWriter { @@ -242,11 +237,11 @@ class SqliteScheduledTaskStore implements ScheduledTaskStore { } async list(): Promise { - return (await this.read()).sort(compareScheduledTasksForList); + return this.readTasks().sort(compareScheduledTasksForList); } async get(id: string): Promise { - return (await this.read()).find((task) => task.id === id); + return this.readTask(id); } async create(input: unknown, now = Date.now()): Promise { @@ -271,103 +266,79 @@ class SqliteScheduledTaskStore implements ScheduledTaskStore { runs: [], lastError: null, }; - await this.mutate((state) => ({ ...state, tasks: [...state.tasks, task] })); + await this.enqueueWrite(() => { + this.#lease.database + .prepare(` + INSERT INTO workflow_scheduled_tasks(task_id, created_at, updated_at, record_json) + VALUES (?, ?, ?, ?) + `) + .run(task.id, task.createdAt, task.updatedAt, JSON.stringify(task)); + }); return task; } async update(id: string, patch: unknown, now = Date.now()): Promise { const normalized = normalizeUpdateScheduledTaskInput(patch, now); if (!normalized.ok) throw storeError('invalid_input', normalized.message); - let updated: ScheduledTask | undefined; - await this.mutate((state) => ({ - ...state, - tasks: state.tasks.map((task) => { - if (task.id !== id) return task; - assertNoPendingClaim(state.claims, id); - if (task.status === 'completed' || task.status === 'expired') { - throw storeError('operation_conflict', 'Cannot update a terminal scheduled task'); - } - const schedule = normalized.value.schedule ?? task.schedule; - const nextFireAt = task.status === 'active' ? computeRequiredNext(schedule, now) : null; - const effect = normalized.value.effect ?? task.effect; - const intentBody = normalized.value.intentBody ?? task.intent.body; - const expiresAt = Object.prototype.hasOwnProperty.call(normalized.value, 'expiresAt') - ? (normalized.value.expiresAt ?? null) - : task.expiresAt; - const maxFires = Object.prototype.hasOwnProperty.call(normalized.value, 'maxFires') - ? (normalized.value.maxFires ?? null) - : task.maxFires; - if (effect.kind !== 'notify' && !intentBody.trim()) { - throw storeError('invalid_input', 'Agent intent body is required'); - } - if (maxFires !== null && maxFires <= task.fireCount) { - throw storeError( - 'operation_conflict', - 'maxFires must be greater than the current fireCount', - ); - } - if (nextFireAt !== null && expiresAt !== null && nextFireAt >= expiresAt) { - throw storeError('invalid_input', 'Schedule must fire before expiresAt'); - } - updated = { - ...task, - ...(normalized.value.title !== undefined ? { title: normalized.value.title } : {}), - ...(normalized.value.intentBody !== undefined - ? { intent: { kind: 'text', body: normalized.value.intentBody } } - : {}), - schedule, - effect, - ...(Object.prototype.hasOwnProperty.call(normalized.value, 'maxFires') - ? { maxFires } - : {}), - expiresAt, - nextFireAt, - updatedAt: now, - }; - return updated; - }), - })); - if (!updated) throw storeError('not_found', `No such scheduled task: ${id}`); - return updated; + return this.updateTask(id, (task) => { + if (task.status === 'completed' || task.status === 'expired') { + throw storeError('operation_conflict', 'Cannot update a terminal scheduled task'); + } + const schedule = normalized.value.schedule ?? task.schedule; + const nextFireAt = task.status === 'active' ? computeRequiredNext(schedule, now) : null; + const effect = normalized.value.effect ?? task.effect; + const intentBody = normalized.value.intentBody ?? task.intent.body; + const expiresAt = Object.prototype.hasOwnProperty.call(normalized.value, 'expiresAt') + ? (normalized.value.expiresAt ?? null) + : task.expiresAt; + const maxFires = Object.prototype.hasOwnProperty.call(normalized.value, 'maxFires') + ? (normalized.value.maxFires ?? null) + : task.maxFires; + if (effect.kind !== 'notify' && !intentBody.trim()) { + throw storeError('invalid_input', 'Agent intent body is required'); + } + if (maxFires !== null && maxFires <= task.fireCount) { + throw storeError( + 'operation_conflict', + 'maxFires must be greater than the current fireCount', + ); + } + if (nextFireAt !== null && expiresAt !== null && nextFireAt >= expiresAt) { + throw storeError('invalid_input', 'Schedule must fire before expiresAt'); + } + return { + ...task, + ...(normalized.value.title !== undefined ? { title: normalized.value.title } : {}), + ...(normalized.value.intentBody !== undefined + ? { intent: { kind: 'text', body: normalized.value.intentBody } } + : {}), + schedule, + effect, + ...(Object.prototype.hasOwnProperty.call(normalized.value, 'maxFires') ? { maxFires } : {}), + expiresAt, + nextFireAt, + updatedAt: now, + }; + }); } async pause(id: string, now = Date.now()): Promise { - let updated: ScheduledTask | undefined; - await this.mutate((state) => ({ - ...state, - tasks: state.tasks.map((task) => { - if (task.id !== id) return task; - assertNoPendingClaim(state.claims, id); - updated = pauseScheduledTask(task, now); - return updated; - }), - })); - if (!updated) throw storeError('not_found', `No such scheduled task: ${id}`); - return updated; + return this.updateTask(id, (task) => pauseScheduledTask(task, now)); } async resume(id: string, now = Date.now()): Promise { - let updated: ScheduledTask | undefined; - await this.mutate((state) => ({ - ...state, - tasks: state.tasks.map((task) => { - if (task.id !== id) return task; - assertNoPendingClaim(state.claims, id); - const result = resumeScheduledTask(task, now); - if ('error' in result) throw storeError('operation_conflict', result.error); - if ( - result.nextFireAt !== null && - result.expiresAt !== null && - result.nextFireAt >= result.expiresAt - ) { - throw storeError('invalid_input', 'Schedule must fire before expiresAt'); - } - updated = result; - return updated; - }), - })); - if (!updated) throw storeError('not_found', `No such scheduled task: ${id}`); - return updated; + return this.updateTask(id, (task) => { + const result = resumeScheduledTask(task, now); + if ('error' in result) throw storeError('operation_conflict', result.error); + if ( + result.nextFireAt !== null && + result.expiresAt !== null && + result.nextFireAt >= result.expiresAt + ) { + throw storeError('invalid_input', 'Schedule must fire before expiresAt'); + } + return result; + }); } async snooze(id: string, delayMs: number, now = Date.now()): Promise { @@ -377,67 +348,51 @@ class SqliteScheduledTaskStore implements ScheduledTaskStore { 'Scheduled task snooze delay must be between 1 ms and 7 days', ); } - let updated: ScheduledTask | undefined; - await this.mutate((state) => ({ - ...state, - tasks: state.tasks.map((task) => { - if (task.id !== id) return task; - assertNoPendingClaim(state.claims, id); - if (task.status !== 'active' || task.nextFireAt === null) { - throw storeError('operation_conflict', 'Only active scheduled tasks can be snoozed'); - } - const nextFireAt = Math.max(now, task.nextFireAt) + Math.floor(delayMs); - if (task.expiresAt !== null && nextFireAt >= task.expiresAt) { - throw storeError('invalid_input', 'Snooze would move the task beyond expiresAt'); - } - updated = { ...task, nextFireAt, updatedAt: now }; - return updated; - }), - })); - if (!updated) throw storeError('not_found', `No such scheduled task: ${id}`); - return updated; + return this.updateTask(id, (task) => { + if (task.status !== 'active' || task.nextFireAt === null) { + throw storeError('operation_conflict', 'Only active scheduled tasks can be snoozed'); + } + const nextFireAt = Math.max(now, task.nextFireAt) + Math.floor(delayMs); + if (task.expiresAt !== null && nextFireAt >= task.expiresAt) { + throw storeError('invalid_input', 'Snooze would move the task beyond expiresAt'); + } + return { ...task, nextFireAt, updatedAt: now }; + }); } async clearRunHistory(id: string, now = Date.now()): Promise { - let updated: ScheduledTask | undefined; - await this.mutate((state) => ({ - ...state, - tasks: state.tasks.map((task) => { - if (task.id !== id) return task; - assertNoPendingClaim(state.claims, id); - updated = { ...task, runs: [], lastError: null, updatedAt: now }; - return updated; - }), + return this.updateTask(id, (task) => ({ + ...task, + runs: [], + lastError: null, + updatedAt: now, })); - if (!updated) throw storeError('not_found', `No such scheduled task: ${id}`); - return updated; } async remove(id: string): Promise { - let found = false; - await this.mutate((state) => { - assertNoPendingClaim(state.claims, id); - const next = state.tasks.filter((task) => { - if (task.id === id) { - found = true; - return false; - } - return true; - }); - return { ...state, tasks: next }; + await this.enqueueWrite(() => { + this.assertNoPendingClaim(id); + this.requireTask(id); + this.#lease.database + .prepare('DELETE FROM workflow_scheduled_tasks WHERE task_id = ?') + .run(id); }); - if (!found) throw storeError('not_found', `No such scheduled task: ${id}`); } async claimNextDue(now = Date.now()): Promise { - let claimed: ScheduledTaskFireClaim | undefined; - const expired: ScheduledTask[] = []; - await this.mutate((state) => { - const claimedTaskIds = new Set(state.claims.map((claim) => claim.taskId)); - const tasks = state.tasks.map((task) => { + return this.enqueueWrite(() => { + // Due discovery still traverses the catalog to return every newly expired + // task. Only claim keys are needed, and only changed rows are written. + const claimRows = this.#lease.database + .prepare('SELECT task_id FROM workflow_scheduled_task_fires') + .all() as Array<{ task_id: string }>; + const claimedTaskIds = new Set(claimRows.map((row) => row.task_id)); + const expired: ScheduledTask[] = []; + const tasks = this.readTasks().map((task) => { if (task.status === 'active' && task.expiresAt !== null && now >= task.expiresAt) { const next = { ...task, status: 'expired' as const, nextFireAt: null, updatedAt: now }; expired.push(next); + this.writeTask(next); return next; } return task; @@ -447,213 +402,252 @@ class SqliteScheduledTaskStore implements ScheduledTaskStore { .sort( (left, right) => left.nextFireAt! - right.nextFireAt! || left.id.localeCompare(right.id), )[0]; - if (!task) return { ...state, tasks }; - claimed = createClaim(task, task.nextFireAt!, now); - return { tasks, claims: [...state.claims, claimed] }; + const claim = task ? createClaim(task, task.nextFireAt!, now) : null; + if (claim) this.insertClaim(claim); + return { claim, expired }; }); - return { claim: claimed ?? null, expired }; } async claimNow(id: string, now = Date.now()): Promise { - let claimed: ScheduledTaskFireClaim | undefined; - await this.mutate((state) => { - const task = state.tasks.find((entry) => entry.id === id); - if (!task) throw storeError('not_found', `No such scheduled task: ${id}`); - assertNoPendingClaim(state.claims, id); + return this.enqueueWrite(() => { + const task = this.requireTask(id); + this.assertNoPendingClaim(id); if (task.status !== 'active') { throw storeError('operation_conflict', 'Only active tasks can be triggered now'); } if (task.expiresAt !== null && now >= task.expiresAt) { throw storeError('operation_conflict', 'Scheduled task has expired'); } - claimed = createClaim(task, now, now); - return { ...state, claims: [...state.claims, claimed] }; + const claim = createClaim(task, now, now); + this.insertClaim(claim); + return claim; }); - return claimed!; } async listPendingFires(): Promise { - return structuredClone((await this.readState()).claims); + const rows = this.#lease.database + .prepare(` + SELECT record_json + FROM workflow_scheduled_task_fires + ORDER BY claimed_at, claim_id + `) + .all() as Array<{ record_json?: unknown }>; + return rows.map((row, index) => decodeClaimRow(row, `row ${index + 1}`)); } async bindFireExecution( claimId: string, execution: ScheduledTaskFireExecution, ): Promise { - let updated: ScheduledTaskFireClaim | undefined; - await this.mutate((state) => ({ - ...state, - claims: state.claims.map((claim) => { - if (claim.id !== claimId) return claim; - if (claim.task.effect.kind === 'notify') { - throw storeError( - 'operation_conflict', - `Scheduled task fire ${claimId} is not an Agent execution`, - ); - } - if (claim.execution && !sameExecution(claim.execution, execution)) { + return this.updateClaim(claimId, (claim) => { + if (claim.task.effect.kind === 'notify') { + throw storeError( + 'operation_conflict', + `Scheduled task fire ${claimId} is not an Agent execution`, + ); + } + if (claim.execution) { + if (!sameExecution(claim.execution, execution)) { throw storeError( 'operation_conflict', `Scheduled task fire ${claimId} already has another execution`, ); } - updated = { ...claim, execution: { ...execution } }; - return updated; - }), - })); - if (!updated) { - throw storeError('not_found', `No such scheduled task fire claim: ${claimId}`); - } - return structuredClone(updated); + return claim; + } + return { ...claim, execution: { ...execution } }; + }); } async setFireNativeState( claimId: string, nativeState: ScheduledTaskNativeFireState, ): Promise { - let updated: ScheduledTaskFireClaim | undefined; - await this.mutate((state) => ({ - ...state, - claims: state.claims.map((claim) => { - if (claim.id !== claimId) return claim; - if (claim.task.effect.kind !== 'notify') { - throw storeError( - 'operation_conflict', - `Scheduled task fire ${claimId} is not a native effect`, - ); - } - if (claim.nativeState === 'invoking' && nativeState !== 'invoking') { - throw storeError( - 'operation_conflict', - `Scheduled task fire ${claimId} already crossed delivery admission`, - ); - } - updated = { ...claim, nativeState }; - return updated; - }), - })); - if (!updated) { - throw storeError('not_found', `No such scheduled task fire claim: ${claimId}`); - } - return structuredClone(updated); + return this.updateClaim(claimId, (claim) => { + if (claim.task.effect.kind !== 'notify') { + throw storeError( + 'operation_conflict', + `Scheduled task fire ${claimId} is not a native effect`, + ); + } + if (claim.nativeState === 'invoking' && nativeState !== 'invoking') { + throw storeError( + 'operation_conflict', + `Scheduled task fire ${claimId} already crossed delivery admission`, + ); + } + return claim.nativeState === nativeState ? claim : { ...claim, nativeState }; + }); } async cancelWaitingNativeFire(taskId: string): Promise { - let cancelled = false; - await this.mutate((state) => { - const claim = state.claims.find((entry) => entry.taskId === taskId); - if (!claim) return state; + return this.enqueueWrite(() => { + const row = this.#lease.database + .prepare('SELECT record_json FROM workflow_scheduled_task_fires WHERE task_id = ?') + .get(taskId) as { record_json?: unknown } | undefined; + if (!row) return false; + const claim = decodeClaimRow(row, `task ${taskId}`); if (claim.nativeState !== 'waiting_for_provider') { throw storeError('operation_conflict', 'Scheduled task has a fire in progress'); } - cancelled = true; - return { - ...state, - claims: state.claims.filter((entry) => entry.id !== claim.id), - }; + this.#lease.database + .prepare('DELETE FROM workflow_scheduled_task_fires WHERE task_id = ?') + .run(taskId); + return true; }); - return cancelled; } async settleFire( claimId: string, run: Omit & { id?: string }, ): Promise { - let updated: ScheduledTask | undefined; - await this.mutate((state) => { - const claim = state.claims.find((entry) => entry.id === claimId); + return this.enqueueWrite(() => { + const claim = this.readClaim(claimId); if (!claim) throw new Error(`No such scheduled task fire claim: ${claimId}`); - const tasks = state.tasks.map((task) => { - if (task.id !== claim.taskId) return task; - const record: ScheduledTaskRun = { - id: run.id ?? randomUUID(), - at: run.at, - outcome: run.outcome, - message: [...run.message].slice(0, SCHEDULED_TASK_RUN_MESSAGE_MAX_CHARS).join(''), - ...(run.sessionId ? { sessionId: run.sessionId } : {}), - ...(run.runId ? { runId: run.runId } : {}), - }; - updated = nextScheduledTaskStateAfterFire(task, record); - return updated; - }); - return { tasks, claims: state.claims.filter((entry) => entry.id !== claimId) }; + const task = this.readTask(claim.taskId); + if (!task) throw new Error(`No such scheduled task: claim ${claimId}`); + const record: ScheduledTaskRun = { + id: run.id ?? randomUUID(), + at: run.at, + outcome: run.outcome, + message: [...run.message].slice(0, SCHEDULED_TASK_RUN_MESSAGE_MAX_CHARS).join(''), + ...(run.sessionId ? { sessionId: run.sessionId } : {}), + ...(run.runId ? { runId: run.runId } : {}), + }; + const updated = nextScheduledTaskStateAfterFire(task, record); + this.writeTask(updated); + this.#lease.database + .prepare('DELETE FROM workflow_scheduled_task_fires WHERE claim_id = ?') + .run(claimId); + return updated; }); - if (!updated) throw new Error(`No such scheduled task: claim ${claimId}`); - return updated; - } - - private async read(): Promise { - return (await this.readState()).tasks; } - private async readState(): Promise { + private readTasks(): ScheduledTask[] { const rows = this.#lease.database .prepare(` - SELECT record_json + SELECT task_id, record_json FROM workflow_scheduled_tasks ORDER BY created_at, task_id `) - .all() as Array<{ record_json?: unknown }>; - const tasks = rows.map((row, index) => { - if (typeof row.record_json !== 'string') { - throw new Error(`Invalid scheduled task at row ${index + 1}`); - } - return decodePersistedScheduledTask( - markPersisted(JSON.parse(row.record_json)), - ); + .all() as Array<{ task_id: string; record_json?: unknown }>; + return rows.map((row, index) => decodeTaskRow(row, `row ${index + 1}`)); + } + + private readTask(id: string): ScheduledTask | undefined { + const row = this.#lease.database + .prepare('SELECT task_id, record_json FROM workflow_scheduled_tasks WHERE task_id = ?') + .get(id) as { task_id: string; record_json?: unknown } | undefined; + if (!row) return undefined; + return decodeTaskRow(row, `task ${id}`); + } + + private requireTask(id: string): ScheduledTask { + const task = this.readTask(id); + if (!task) throw storeError('not_found', `No such scheduled task: ${id}`); + return task; + } + + private readClaim(id: string): ScheduledTaskFireClaim | undefined { + const row = this.#lease.database + .prepare('SELECT record_json FROM workflow_scheduled_task_fires WHERE claim_id = ?') + .get(id) as { record_json?: unknown } | undefined; + if (!row) return undefined; + const claim = decodeClaimRow(row, `claim ${id}`); + if (claim.id !== id) throw new Error(`Invalid scheduled task fire claim identity: ${id}`); + return claim; + } + + private assertNoPendingClaim(taskId: string): void { + const pending = this.#lease.database + .prepare('SELECT 1 FROM workflow_scheduled_task_fires WHERE task_id = ?') + .get(taskId); + if (pending) { + throw storeError('operation_conflict', 'Scheduled task has a fire in progress'); + } + } + + private updateTask( + id: string, + update: (task: ScheduledTask) => ScheduledTask, + ): Promise { + return this.enqueueWrite(() => { + const task = this.requireTask(id); + this.assertNoPendingClaim(id); + const updated = update(task); + if (updated !== task) this.writeTask(updated); + return updated; }); - const claimRows = this.#lease.database - .prepare(` - SELECT record_json - FROM workflow_scheduled_task_fires - ORDER BY claimed_at, claim_id - `) - .all() as Array<{ record_json?: unknown }>; - const claims = claimRows.map((row, index) => { - if (typeof row.record_json !== 'string') { - throw new Error(`Invalid scheduled task fire claim at row ${index + 1}`); + } + + private updateClaim( + id: string, + update: (claim: ScheduledTaskFireClaim) => ScheduledTaskFireClaim, + ): Promise { + return this.enqueueWrite(() => { + const claim = this.readClaim(id); + if (!claim) throw storeError('not_found', `No such scheduled task fire claim: ${id}`); + const updated = update(claim); + if (updated !== claim) { + this.#lease.database + .prepare('UPDATE workflow_scheduled_task_fires SET record_json = ? WHERE claim_id = ?') + .run(JSON.stringify(updated), id); } - const claim = JSON.parse(row.record_json) as ScheduledTaskFireClaim; - return { - ...claim, - task: decodePersistedScheduledTask(markPersisted(claim.task)), - }; + return structuredClone(updated); }); - return { tasks, claims }; } - private async mutate( - fn: (state: ScheduledTaskStoreState) => ScheduledTaskStoreState, - ): Promise { - const run = async () => { - const current = await this.readState(); - this.write(fn(current)); - }; - const next = this.queue.then(run, run); - this.queue = next.catch(() => {}); - await next; + private writeTask(task: ScheduledTask): void { + this.#lease.database + .prepare( + 'UPDATE workflow_scheduled_tasks SET updated_at = ?, record_json = ? WHERE task_id = ?', + ) + .run(task.updatedAt, JSON.stringify(task), task.id); } - private write(state: ScheduledTaskStoreState): void { - this.#lease.transaction('write', () => { - this.#lease.database.prepare('DELETE FROM workflow_scheduled_tasks').run(); - this.#lease.database.prepare('DELETE FROM workflow_scheduled_task_fires').run(); - const insert = this.#lease.database.prepare(` - INSERT INTO workflow_scheduled_tasks(task_id, created_at, updated_at, record_json) - VALUES (?, ?, ?, ?) - `); - for (const task of state.tasks) { - insert.run(task.id, task.createdAt, task.updatedAt, JSON.stringify(task)); - } - const insertClaim = this.#lease.database.prepare(` + private insertClaim(claim: ScheduledTaskFireClaim): void { + this.#lease.database + .prepare(` INSERT INTO workflow_scheduled_task_fires(claim_id, task_id, claimed_at, record_json) VALUES (?, ?, ?, ?) - `); - for (const claim of state.claims) { - insertClaim.run(claim.id, claim.taskId, claim.claimedAt, JSON.stringify(claim)); - } - }); + `) + .run(claim.id, claim.taskId, claim.claimedAt, JSON.stringify(claim)); + } + + private enqueueWrite(operation: () => T): Promise { + const run = () => this.#lease.transaction('write', operation); + const next = this.queue.then(run, run); + this.queue = next.then( + () => {}, + () => {}, + ); + return next; + } +} + +function decodeTaskRow( + row: { task_id: string; record_json?: unknown }, + location: string, +): ScheduledTask { + if (typeof row.record_json !== 'string') { + throw new Error(`Invalid scheduled task at ${location}`); } + const task = decodePersistedScheduledTask( + markPersisted(JSON.parse(row.record_json)), + ); + if (task.id !== row.task_id) throw new Error(`Invalid scheduled task identity: ${row.task_id}`); + return task; +} + +function decodeClaimRow(row: { record_json?: unknown }, location: string): ScheduledTaskFireClaim { + if (typeof row.record_json !== 'string') { + throw new Error(`Invalid scheduled task fire claim at ${location}`); + } + const claim = JSON.parse(row.record_json) as ScheduledTaskFireClaim; + return { + ...claim, + task: decodePersistedScheduledTask(markPersisted(claim.task)), + }; } function computeRequiredNext(schedule: ScheduledTaskSchedule, now: number): number { @@ -678,12 +672,6 @@ function createClaim( }; } -function assertNoPendingClaim(claims: readonly ScheduledTaskFireClaim[], taskId: string): void { - if (claims.some((claim) => claim.taskId === taskId)) { - throw storeError('operation_conflict', 'Scheduled task has a fire in progress'); - } -} - function storeError(code: ScheduledTaskStoreErrorCode, message: string): ScheduledTaskStoreError { return new ScheduledTaskStoreError(code, message); }