From a2c03c316eb71d8c534e831dbff8dbe80eb5d9c7 Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Sun, 16 Aug 2026 21:38:30 -0700 Subject: [PATCH] Add transactional safety to dispatch and message operations Wrap dispatch failures and batched message updates with SAVEPOINTs to ensure atomicity and idempotency: - Dispatch failures now check status guards and roll back if the related task update fails, preventing partial state corruption - Message batches (across multiple 500-id chunks) roll back entirely if any batch fails, avoiding partial mutations - Add tests verifying idempotency and atomicity under failure conditions --- .../dispatch-context/dispatch-completion.ts | 75 +++++++++------ .../db/messages/message-inbox.ts | 79 ++++++++++------ .../dispatch-failure-idempotency.test.ts | 60 ++++++++++++ .../message-batch-atomicity.test.ts | 94 +++++++++++++++++++ 4 files changed, 247 insertions(+), 61 deletions(-) create mode 100644 src/main/runtime/orchestration/dispatch-failure-idempotency.test.ts create mode 100644 src/main/runtime/orchestration/message-batch-atomicity.test.ts diff --git a/src/main/runtime/orchestration/db/dispatch-context/dispatch-completion.ts b/src/main/runtime/orchestration/db/dispatch-context/dispatch-completion.ts index e4de5f7280b..fa3006cd5b6 100644 --- a/src/main/runtime/orchestration/db/dispatch-context/dispatch-completion.ts +++ b/src/main/runtime/orchestration/db/dispatch-context/dispatch-completion.ts @@ -2,6 +2,8 @@ import type { TaskStatus, DispatchStatus, DispatchContextRow } from '../../types import { DISPATCH_CIRCUIT_BREAK_FAILURES } from './dispatch-circuit-breaker' import type { OrchestrationDb } from '../orchestration-db' +const FAIL_DISPATCH_SAVEPOINT = 'fail_dispatch' + export function completeDispatch(this: OrchestrationDb, ctxId: string): void { this.db .prepare( @@ -65,37 +67,50 @@ export function failDispatch( ctxId: string, error: string ): DispatchContextRow | undefined { - const ctx = this.db.prepare('SELECT * FROM dispatch_contexts WHERE id = ?').get(ctxId) as - | DispatchContextRow - | undefined - if (!ctx) { - return undefined + this.db.exec(`SAVEPOINT ${FAIL_DISPATCH_SAVEPOINT}`) + try { + const ctx = this.db.prepare('SELECT * FROM dispatch_contexts WHERE id = ?').get(ctxId) as + | DispatchContextRow + | undefined + if (!ctx || (ctx.status !== 'pending' && ctx.status !== 'dispatched')) { + this.db.exec(`RELEASE ${FAIL_DISPATCH_SAVEPOINT}`) + return ctx + } + + const newFailureCount = ctx.failure_count + 1 + const newStatus: DispatchStatus = + newFailureCount >= DISPATCH_CIRCUIT_BREAK_FAILURES ? 'circuit_broken' : 'failed' + const result = this.db + .prepare( + `UPDATE dispatch_contexts + SET status = ?, failure_count = ?, last_failure = ?, + completed_at = COALESCE(completed_at, datetime('now')), + capability_revoked_at = COALESCE(capability_revoked_at, datetime('now')) + WHERE id = ? AND status IN ('pending', 'dispatched')` + ) + .run(newStatus, newFailureCount, error, ctxId) + if (result.changes !== 1) { + this.db.exec(`RELEASE ${FAIL_DISPATCH_SAVEPOINT}`) + return this.db.prepare('SELECT * FROM dispatch_contexts WHERE id = ?').get(ctxId) as + | DispatchContextRow + | undefined + } + + // Why: back to 'ready' not 'pending' — 'pending' would strand it since promoteReadyTasks only runs when a dep completes. + const taskStatus: TaskStatus = newStatus === 'circuit_broken' ? 'failed' : 'ready' + // Why: the status guard keeps a late failure from reopening a task that already completed or was retried elsewhere. + this.db + .prepare("UPDATE tasks SET status = ? WHERE id = ? AND status IN ('dispatched', 'blocked')") + .run(taskStatus, ctx.task_id) + this.db.exec(`RELEASE ${FAIL_DISPATCH_SAVEPOINT}`) + return this.db.prepare('SELECT * FROM dispatch_contexts WHERE id = ?').get(ctxId) as + | DispatchContextRow + | undefined + } catch (cause) { + this.db.exec(`ROLLBACK TO ${FAIL_DISPATCH_SAVEPOINT}`) + this.db.exec(`RELEASE ${FAIL_DISPATCH_SAVEPOINT}`) + throw cause } - - const newFailureCount = ctx.failure_count + 1 - const newStatus: DispatchStatus = - newFailureCount >= DISPATCH_CIRCUIT_BREAK_FAILURES ? 'circuit_broken' : 'failed' - - this.db - .prepare( - `UPDATE dispatch_contexts - SET status = ?, failure_count = ?, last_failure = ?, - completed_at = COALESCE(completed_at, datetime('now')), - capability_revoked_at = COALESCE(capability_revoked_at, datetime('now')) - WHERE id = ?` - ) - .run(newStatus, newFailureCount, error, ctxId) - - // Why: back to 'ready' not 'pending' — 'pending' would strand it since promoteReadyTasks only runs when a dep completes. - const taskStatus: TaskStatus = newStatus === 'circuit_broken' ? 'failed' : 'ready' - // Why: the status guard keeps a late failure from reopening a task that already completed or was retried elsewhere. - this.db - .prepare("UPDATE tasks SET status = ? WHERE id = ? AND status IN ('dispatched', 'blocked')") - .run(taskStatus, ctx.task_id) - - return this.db.prepare('SELECT * FROM dispatch_contexts WHERE id = ?').get(ctxId) as - | DispatchContextRow - | undefined } export type DispatchCompletionMethods = { diff --git a/src/main/runtime/orchestration/db/messages/message-inbox.ts b/src/main/runtime/orchestration/db/messages/message-inbox.ts index d014a7991f6..e5d03eea771 100644 --- a/src/main/runtime/orchestration/db/messages/message-inbox.ts +++ b/src/main/runtime/orchestration/db/messages/message-inbox.ts @@ -4,6 +4,30 @@ import { addLifecycleRejectionMarker } from '../lifecycle-rejection-marker' import type { OrchestrationDb } from '../orchestration-db' const MESSAGE_ID_UPDATE_BATCH_SIZE = 500 +const MESSAGE_MUTATION_SAVEPOINT = 'message_id_mutation' + +function runBatchedMessageMutation( + db: OrchestrationDb, + ids: string[], + sqlForPlaceholders: (placeholders: string) => string +): void { + if (ids.length === 0) { + return + } + db.db.exec(`SAVEPOINT ${MESSAGE_MUTATION_SAVEPOINT}`) + try { + for (let offset = 0; offset < ids.length; offset += MESSAGE_ID_UPDATE_BATCH_SIZE) { + const batch = ids.slice(offset, offset + MESSAGE_ID_UPDATE_BATCH_SIZE) + const placeholders = batch.map(() => '?').join(',') + db.db.prepare(sqlForPlaceholders(placeholders)).run(...batch) + } + db.db.exec(`RELEASE ${MESSAGE_MUTATION_SAVEPOINT}`) + } catch (error) { + db.db.exec(`ROLLBACK TO ${MESSAGE_MUTATION_SAVEPOINT}`) + db.db.exec(`RELEASE ${MESSAGE_MUTATION_SAVEPOINT}`) + throw error + } +} export function getUnreadMessages( this: OrchestrationDb, @@ -123,35 +147,31 @@ export function getMessageById(this: OrchestrationDb, id: string): MessageRow | } export function markAsRead(this: OrchestrationDb, ids: string[]): void { - for (let offset = 0; offset < ids.length; offset += MESSAGE_ID_UPDATE_BATCH_SIZE) { - const batch = ids.slice(offset, offset + MESSAGE_ID_UPDATE_BATCH_SIZE) - const placeholders = batch.map(() => '?').join(',') - this.db.prepare(`UPDATE messages SET read = 1 WHERE id IN (${placeholders})`).run(...batch) - } + runBatchedMessageMutation( + this, + ids, + (placeholders) => `UPDATE messages SET read = 1 WHERE id IN (${placeholders})` + ) } // Why: use datetime('now') so delivered_at matches the space-format UTC shape of the table's other timestamps for correct ordering (§3.2). export function markAsDelivered(this: OrchestrationDb, ids: string[]): void { - for (let offset = 0; offset < ids.length; offset += MESSAGE_ID_UPDATE_BATCH_SIZE) { - const batch = ids.slice(offset, offset + MESSAGE_ID_UPDATE_BATCH_SIZE) - const placeholders = batch.map(() => '?').join(',') - this.db - .prepare(`UPDATE messages SET delivered_at = datetime('now') WHERE id IN (${placeholders})`) - .run(...batch) - } + runBatchedMessageMutation( + this, + ids, + (placeholders) => + `UPDATE messages SET delivered_at = datetime('now') WHERE id IN (${placeholders})` + ) } export function markAsUndelivered(this: OrchestrationDb, ids: string[]): void { - for (let offset = 0; offset < ids.length; offset += MESSAGE_ID_UPDATE_BATCH_SIZE) { - const batch = ids.slice(offset, offset + MESSAGE_ID_UPDATE_BATCH_SIZE) - const placeholders = batch.map(() => '?').join(',') - this.db - .prepare( - `UPDATE messages SET delivered_at = NULL - WHERE read = 0 AND id IN (${placeholders})` - ) - .run(...batch) - } + runBatchedMessageMutation( + this, + ids, + (placeholders) => + `UPDATE messages SET delivered_at = NULL + WHERE read = 0 AND id IN (${placeholders})` + ) } export function areUnreadMessages(this: OrchestrationDb, toHandle: string, ids: string[]): boolean { @@ -173,15 +193,12 @@ export function areUnreadMessages(this: OrchestrationDb, toHandle: string, ids: // Why: superseded lifecycle messages stay in history but must not be consumed or injected after their dispatch finished. export function markAsReadAndDelivered(this: OrchestrationDb, ids: string[]): void { - for (let offset = 0; offset < ids.length; offset += MESSAGE_ID_UPDATE_BATCH_SIZE) { - const batch = ids.slice(offset, offset + MESSAGE_ID_UPDATE_BATCH_SIZE) - const placeholders = batch.map(() => '?').join(',') - this.db - .prepare( - `UPDATE messages SET read = 1, delivered_at = COALESCE(delivered_at, datetime('now')) WHERE id IN (${placeholders})` - ) - .run(...batch) - } + runBatchedMessageMutation( + this, + ids, + (placeholders) => + `UPDATE messages SET read = 1, delivered_at = COALESCE(delivered_at, datetime('now')) WHERE id IN (${placeholders})` + ) } export function getInbox(this: OrchestrationDb, limit = 20): MessageRow[] { diff --git a/src/main/runtime/orchestration/dispatch-failure-idempotency.test.ts b/src/main/runtime/orchestration/dispatch-failure-idempotency.test.ts new file mode 100644 index 00000000000..9b526c480cc --- /dev/null +++ b/src/main/runtime/orchestration/dispatch-failure-idempotency.test.ts @@ -0,0 +1,60 @@ +import { describe, expect, it } from 'vitest' +import type Database from '../../sqlite/sync-database' +import { OrchestrationDb } from './db' + +describe('dispatch failure idempotency', () => { + it('counts an active dispatch failure only once', () => { + const db = new OrchestrationDb(':memory:') + const task = db.createTask({ spec: 'work' }) + const dispatch = db.createDispatchContext(task.id, 'term_worker') + + expect(db.failDispatch(dispatch.id, 'exit')?.failure_count).toBe(1) + const duplicate = db.failDispatch(dispatch.id, 'duplicate escalation') + + expect(duplicate?.failure_count).toBe(1) + expect(duplicate?.status).toBe('failed') + db.close() + }) + + it('does not overwrite a completed dispatch', () => { + const db = new OrchestrationDb(':memory:') + const task = db.createTask({ spec: 'work' }) + const dispatch = db.createDispatchContext(task.id, 'term_worker') + db.completeDispatch(dispatch.id) + + const lateFailure = db.failDispatch(dispatch.id, 'late exit') + + expect(lateFailure?.failure_count).toBe(0) + expect(lateFailure?.status).toBe('completed') + db.close() + }) + + it('rolls back the dispatch when the task update fails', () => { + const db = new OrchestrationDb(':memory:') + const sqlite = (db as unknown as { db: Database.Database }).db + const task = db.createTask({ spec: 'work' }) + const dispatch = db.createDispatchContext(task.id, 'term_worker') + sqlite.exec(` + CREATE TRIGGER reject_task_failure_update + BEFORE UPDATE ON tasks WHEN OLD.id = '${task.id}' + BEGIN + SELECT RAISE(ABORT, 'blocked'); + END; + `) + + expect(() => db.failDispatch(dispatch.id, 'exit')).toThrow('blocked') + expect(db.getDispatchContextById(dispatch.id)).toMatchObject({ + failure_count: 0, + status: 'dispatched' + }) + expect(db.getTask(task.id)?.status).toBe('dispatched') + + sqlite.exec('DROP TRIGGER reject_task_failure_update') + expect(db.failDispatch(dispatch.id, 'retry')).toMatchObject({ + failure_count: 1, + status: 'failed' + }) + expect(db.getTask(task.id)?.status).toBe('ready') + db.close() + }) +}) diff --git a/src/main/runtime/orchestration/message-batch-atomicity.test.ts b/src/main/runtime/orchestration/message-batch-atomicity.test.ts new file mode 100644 index 00000000000..a9bf6237f55 --- /dev/null +++ b/src/main/runtime/orchestration/message-batch-atomicity.test.ts @@ -0,0 +1,94 @@ +import { afterEach, describe, expect, it } from 'vitest' +import type Database from '../../sqlite/sync-database' +import { OrchestrationDb } from './db' + +type MessageMutation = + | 'markAsRead' + | 'markAsDelivered' + | 'markAsUndelivered' + | 'markAsReadAndDelivered' + +const messageIds = Array.from( + { length: 501 }, + (_, index) => `m${index.toString().padStart(3, '0')}` +) + +function seedMessages(sqlite: Database.Database): void { + sqlite.exec(` + WITH RECURSIVE ids(value) AS ( + VALUES(0) + UNION ALL SELECT value + 1 FROM ids WHERE value < 500 + ) + INSERT INTO messages (id, from_handle, to_handle, subject) + SELECT printf('m%03d', value), 'sender', 'recipient', 'subject' FROM ids; + `) +} + +function rejectLastMessageUpdate(sqlite: Database.Database): void { + sqlite.exec(` + CREATE TRIGGER reject_last_message_update + BEFORE UPDATE ON messages WHEN OLD.id = 'm500' + BEGIN + SELECT RAISE(ABORT, 'blocked'); + END; + `) +} + +describe('message batch atomicity', () => { + let db: OrchestrationDb | undefined + + afterEach(() => db?.close()) + + it.each<{ + method: MessageMutation + setup?: string + changedCountSql: string + }>([ + { method: 'markAsRead', changedCountSql: 'SELECT COUNT(*) count FROM messages WHERE read = 1' }, + { + method: 'markAsDelivered', + changedCountSql: 'SELECT COUNT(*) count FROM messages WHERE delivered_at IS NOT NULL' + }, + { + method: 'markAsUndelivered', + setup: "UPDATE messages SET delivered_at = datetime('now')", + changedCountSql: 'SELECT COUNT(*) count FROM messages WHERE delivered_at IS NULL' + }, + { + method: 'markAsReadAndDelivered', + changedCountSql: + 'SELECT COUNT(*) count FROM messages WHERE read = 1 OR delivered_at IS NOT NULL' + } + ])('rolls back $method when a later batch fails', ({ method, setup, changedCountSql }) => { + db = new OrchestrationDb(':memory:') + const sqlite = (db as unknown as { db: Database.Database }).db + seedMessages(sqlite) + if (setup) { + sqlite.exec(setup) + } + rejectLastMessageUpdate(sqlite) + + expect(() => db?.[method](messageIds)).toThrow('blocked') + + const changed = sqlite.prepare(changedCountSql).get() as { count: number } + expect(changed.count).toBe(0) + }) + + it('preserves an outer transaction when an inner batch rolls back', () => { + db = new OrchestrationDb(':memory:') + const sqlite = (db as unknown as { db: Database.Database }).db + seedMessages(sqlite) + rejectLastMessageUpdate(sqlite) + sqlite.exec('BEGIN IMMEDIATE') + sqlite.prepare("UPDATE messages SET subject = 'outer change' WHERE id = 'm000'").run() + + expect(() => db?.markAsRead(messageIds)).toThrow('blocked') + sqlite.exec('COMMIT') + + const first = sqlite.prepare("SELECT subject, read FROM messages WHERE id = 'm000'").get() as { + subject: string + read: number + } + expect(first).toEqual({ subject: 'outer change', read: 0 }) + }) +})