From 10e8f42a87a6e210c51470d45d36061ffb1bcc31 Mon Sep 17 00:00:00 2001 From: Totoro Date: Mon, 7 Sep 2026 12:06:43 +0800 Subject: [PATCH] perf(storage): avoid scheduled-task catalog rewrites Use keyed task and fire-claim reads and row-level writes within the existing serialized SQLite transaction. Keep task settlement and claim removal atomic, and preserve execution identity and native admission rules. Keep due discovery as an explicit catalog scan, writing only newly expired tasks and a new claim when present. Unchanged polls perform no DML. Validate task row identities before expiry updates. Add regression coverage for indexed point operations as unrelated data grows, no-op writes, expiry, identity corruption, rollback and owner reopen. Fixes #4954 Refs: https://github.com/apache/maka/discussions/4876 Generated-by: Codex --- .../scheduled-task-row-operations.test.ts | 605 ++++++++++++++++++ packages/storage/src/scheduled-task-store.ts | 566 ++++++++-------- 2 files changed, 882 insertions(+), 289 deletions(-) create mode 100644 packages/storage/src/__tests__/scheduled-task-row-operations.test.ts 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); }