@pikku/kysely 0.13.31 → 0.13.34

package/dist/src/kysely-workflow-service.js

@@ -1,14 +1,52 @@
-import { PikkuWorkflowService, WorkflowStepFunctionMismatchError, } from '@pikku/core/workflow';
+import { PikkuWorkflowService, WorkflowStepFunctionMismatchError, WorkflowStepLeaseExpiredError, WorkflowStepSupersededError, leaseAttemptsExhausted, } from '@pikku/core/workflow';
 import { sql } from 'kysely';
+import { appNowMs, leaseUntil } from './kysely-lease-clock.js';
 import { KyselyWorkflowRunService } from './kysely-workflow-run-service.js';
 import { parseJson } from './kysely-json.js';
 import { requirePikkuSchema } from './schema/index.js';
 import { workflowSchema } from './schema/workflow.schema.js';
+/**
+ * What a step read returns, one list for the single read and the whole-run
+ * read alike, so the replay snapshot cannot drift from `getStepState`.
+ */
+const STEP_STATE_COLUMNS = [
+    'workflowStepId',
+    'status',
+    'rpcName',
+    'result',
+    'error',
+    'retries',
+    'retryDelay',
+    'fromStepName',
+    // `current_attempt` is bumped alongside every history insert, so it is
+    // already the count this used to aggregate — and reading it keeps the
+    // engine's hottest query off a table that grows for the life of the run.
+    'currentAttempt',
+    'leaseExpiresAt',
+    'createdAt',
+    'updatedAt',
+];
+const toStepState = (row) => ({
+    stepId: row.workflowStepId,
+    status: row.status,
+    rpcName: row.rpcName ?? null,
+    result: parseJson(row.result),
+    error: parseJson(row.error),
+    attemptCount: Number(row.currentAttempt ?? 1),
+    retries: row.retries != null ? Number(row.retries) : undefined,
+    retryDelay: row.retryDelay ?? undefined,
+    fromStepName: row.fromStepName ?? undefined,
+    leaseExpiresAt: row.leaseExpiresAt != null
+        ? new Date(Number(row.leaseExpiresAt))
+        : undefined,
+    createdAt: new Date(row.createdAt),
+    updatedAt: new Date(row.updatedAt),
+});
 export class KyselyWorkflowService extends PikkuWorkflowService {
     db;
     initialized = false;
     runService;
-    constructor(db, options = {}) {
+    constructor(db, options) {
         super(options);
         this.db = db;
         this.runService = new KyselyWorkflowRunService(db);
@@ -92,73 +130,25 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     async getStepState(runId, stepName) {
         const row = await this.db
             .selectFrom('workflowStep')
-            .select([
-            'workflowStepId',
-            'status',
-            'rpcName',
-            'result',
-            'error',
-            'retries',
-            'retryDelay',
-            'fromStepName',
-            // `current_attempt` is bumped alongside every history insert, so it is
-            // already the count this used to aggregate — and reading it keeps the
-            // engine's hottest query off a table that grows for the life of the run.
-            'currentAttempt',
-            'createdAt',
-            'updatedAt',
-        ])
+            .select(STEP_STATE_COLUMNS)
             .where('workflowRunId', '=', runId)
             .where('stepName', '=', stepName)
             .executeTakeFirst();
         if (!row) {
             throw new Error(`Step not found: runId=${runId}, stepName=${stepName}. Use insertStepState to create it.`);
         }
-        return {
-            stepId: row.workflowStepId,
-            status: row.status,
-            rpcName: row.rpcName ?? null,
-            result: parseJson(row.result),
-            error: parseJson(row.error),
-            attemptCount: Number(row.currentAttempt ?? 1),
-            retries: row.retries != null ? Number(row.retries) : undefined,
-            retryDelay: row.retryDelay ?? undefined,
-            fromStepName: row.fromStepName ?? undefined,
-            createdAt: new Date(row.createdAt),
-            updatedAt: new Date(row.updatedAt),
-        };
+        return toStepState(row);
     }
     async listStepStates(runId) {
         const rows = await this.db
             .selectFrom('workflowStep')
-            .select([
-            'workflowStepId',
-            'stepName',
-            'status',
-            'result',
-            'error',
-            'retries',
-            'retryDelay',
-            'fromStepName',
-            'currentAttempt',
-            'createdAt',
-            'updatedAt',
-        ])
+            .select([...STEP_STATE_COLUMNS, 'stepName'])
             .where('workflowRunId', '=', runId)
             .execute();
-        return rows.map((row) => ({
-            stepId: row.workflowStepId,
-            stepName: row.stepName,
-            status: row.status,
-            result: parseJson(row.result),
-            error: parseJson(row.error),
-            attemptCount: Number(row.currentAttempt ?? 1),
-            retries: row.retries != null ? Number(row.retries) : undefined,
-            retryDelay: row.retryDelay ?? undefined,
-            fromStepName: row.fromStepName ?? undefined,
-            createdAt: new Date(row.createdAt),
-            updatedAt: new Date(row.updatedAt),
-        }));
+        return rows.map((row) => ({ ...toStepState(row), stepName: row.stepName }));
+    }
+    async getRunSteps(runId) {
+        return this.runService.getRunSteps(runId);
     }
     async getRunHistory(runId) {
         return this.runService.getRunHistory(runId);
@@ -168,6 +158,24 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             runningAt: new Date(),
         });
     }
+    async refreshStepLease(stepId, leaseMs, attempt) {
+        let query = this.db
+            .updateTable('workflowStep')
+            .set({
+            leaseExpiresAt: leaseMs === null ? null : leaseUntil(this.nowMs(), leaseMs),
+            updatedAt: new Date(),
+        })
+            .where('workflowStepId', '=', stepId);
+        if (attempt !== undefined) {
+            query = query.where(sql `coalesce(current_attempt, 1)`, '=', attempt);
+        }
+        const renewed = await query.executeTakeFirst();
+        return Number(renewed?.numUpdatedRows ?? 0n) > 0;
+    }
+    /** Now, in epoch milliseconds, on the clock step leases are judged by. */
+    nowMs() {
+        return appNowMs();
+    }
     async setStepScheduledImpl(stepId) {
         await this.db
             .updateTable('workflowStep')
@@ -186,7 +194,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
      * over `created_at`, so a retry that lands in the same millisecond as the
      * attempt it replaces still resolves the newer row.
      */
-    async writeStepTransition(stepId, status, historyValues) {
+    async writeStepTransition(stepId, status, historyValues, attempt) {
         const now = new Date();
         const stepValues = { status, updatedAt: now };
         if ('result' in historyValues)
@@ -199,7 +207,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
         if (status === 'failed')
             stepValues.result = null;
         await this.db.transaction().execute(async (trx) => {
-            await trx
+            let step = trx
                 .updateTable('workflowStep')
                 .set({
                 ...stepValues,
@@ -209,8 +217,17 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
                 // MySQL would reject against the table being updated.
                 currentAttempt: sql `coalesce(current_attempt, 1)`,
             })
-                .where('workflowStepId', '=', stepId)
-                .execute();
+                .where('workflowStepId', '=', stepId);
+            // Fenced to the attempt that made the write, so a dispatch whose step
+            // was claimed again after its lease lapsed cannot record over the newer
+            // attempt. Throwing rolls the transaction back.
+            if (attempt !== undefined) {
+                step = step.where(sql `coalesce(current_attempt, 1)`, '=', attempt);
+            }
+            const moved = await step.executeTakeFirst();
+            if (attempt !== undefined && Number(moved?.numUpdatedRows ?? 0n) === 0) {
+                throw new WorkflowStepSupersededError(stepId, attempt);
+            }
             const written = await trx
                 .updateTable('workflowStepHistory')
                 .set({ status, ...historyValues })
@@ -247,7 +264,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             }
         });
     }
-    async insertHistoryRecord(stepId, status, attempt, result, error) {
+    async insertHistoryRecord(stepId, status, attempt, result, error, db = this.db) {
         const now = new Date();
         const values = {
             historyId: crypto.randomUUID(),
@@ -262,7 +279,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
         if (timestampField !== 'createdAt') {
             values[timestampField] = now;
         }
-        await this.db
+        await db
             .insertInto('workflowStepHistory')
             .values(values)
             .execute();
@@ -291,13 +308,13 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             .where('workflowStepId', '=', stepId)
             .execute();
     }
-    async setStepResultImpl(stepId, result) {
+    async setStepResultImpl(stepId, result, attempt) {
         await this.writeStepTransition(stepId, 'succeeded', {
             result: JSON.stringify(result),
             succeededAt: new Date(),
-        });
+        }, attempt);
     }
-    async setStepErrorImpl(stepId, error) {
+    async setStepErrorImpl(stepId, error, attempt) {
         const serializedError = {
             message: error.message,
             stack: error.stack,
@@ -306,7 +323,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
         await this.writeStepTransition(stepId, 'failed', {
             error: JSON.stringify(serializedError),
             failedAt: new Date(),
-        });
+        }, attempt);
     }
     async createRetryAttemptImpl(stepId, status) {
         // Read before writing: the new attempt number comes from the history this
@@ -329,6 +346,9 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             .where('workflowStepId', '=', stepId)
             .execute();
         await this.insertHistoryRecord(stepId, status, attempt);
+        return this.readStepAttempt(stepId);
+    }
+    async readStepAttempt(stepId) {
         const row = await this.db
             .selectFrom('workflowStep')
             .select([
@@ -360,22 +380,6 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             updatedAt: new Date(row.updatedAt),
         };
     }
-    async withRunLock(_id, fn) {
-        return fn();
-    }
-    /**
-     * A pass-through, and deliberately so: the one decision that must exclude —
-     * claiming a step to execute it — is made by `claimStepForExecution` below as
-     * a single conditional write, which needs no lock to be exclusive.
-     *
-     * What is left are the suspend and approval sections in the engine, where the
-     * lock is genuine mutual exclusion rather than a claim. A dialect with a real
-     * primitive overrides this to cover them — `kysely-postgres` with
-     * `pg_advisory_xact_lock`, `kysely-mysql` with `GET_LOCK`.
-     */
-    async withStepLock(_runId, _stepName, fn) {
-        return fn();
-    }
     /**
      * Claim the step with a status-guarded `UPDATE` and read the affected-row
      * count: the database decides the winner in one statement, so two dispatches
@@ -388,28 +392,56 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
      *
      * The winner then goes through the ordinary transition methods, so history
      * rows and the mirror see exactly what they saw before.
+     *
+     * A `running` step is taken back only once its lease has lapsed, and the
+     * lease is part of the same guarded `UPDATE` — checking it here and writing
+     * it afterwards would let two dispatches read one lapsed lease and both
+     * proceed. Whether it has lapsed is the database's call, on its own clock:
+     * a live lease simply matches no row.
      */
-    async claimStepForExecution(runId, stepName, rpcName) {
+    async claimStepForExecution(runId, stepName, rpcName, leaseMs) {
         const stepState = await this.getStepState(runId, stepName);
         // knowledge: decisions/security/a-step-runs-the-function-the-workflow-dispatched-it-with.md
         if (stepState.rpcName !== undefined &&
             stepState.rpcName !== (rpcName ?? null)) {
             throw new WorkflowStepFunctionMismatchError(runId, stepName);
         }
-        if (stepState.status === 'succeeded' || stepState.status === 'running') {
+        if (stepState.status === 'succeeded') {
             return null;
         }
-        if (stepState.status === 'failed') {
-            if (!(await this.claimStepStatus(stepState.stepId, ['failed']))) {
+        if (stepState.status === 'running') {
+            if (leaseAttemptsExhausted(stepState)) {
+                // Fail it only if its lease is still lapsed at the moment of writing:
+                // a worker that renewed since the read above still owns the step.
+                if (!(await this.claimStepStatus(stepState.stepId, ['running'], leaseMs))) {
+                    return null;
+                }
+                await this.setStepError(stepState.stepId, new WorkflowStepLeaseExpiredError(runId, stepName, stepState.attemptCount), stepState.attemptCount);
+                // `executeWorkflowStepInner` reads a null claim as "another dispatch
+                // owns it" and returns without resuming, so the run would stay
+                // `running` until a stalled sweep. Requeue the orchestrator here.
+                await this.resumeWorkflow(runId);
                 return null;
             }
-            return this.createRetryAttempt(stepState.stepId, 'running');
+        }
+        // A failed step that has spent every attempt is settled, however it failed;
+        // a redelivered message must not buy it another one.
+        if (stepState.status === 'failed' && leaseAttemptsExhausted(stepState)) {
+            return null;
+        }
+        if (stepState.status === 'failed' || stepState.status === 'running') {
+            const attempt = stepState.attemptCount + 1;
+            return this.mirrored(() => this.reclaimStepAttempt(stepState, attempt, leaseMs), async (mirror, newStep) => {
+                if (newStep) {
+                    await mirror.createRetryAttempt(stepState.stepId, {
+                        ...newStep,
+                        stepName,
+                    });
+                }
+            });
         }
         if (stepState.status === 'pending' || stepState.status === 'scheduled') {
-            if (!(await this.claimStepStatus(stepState.stepId, [
-                'pending',
-                'scheduled',
-            ]))) {
+            if (!(await this.claimStepStatus(stepState.stepId, ['pending', 'scheduled'], leaseMs))) {
                 return null;
             }
             await this.setStepRunning(stepState.stepId);
@@ -418,16 +450,62 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
         return stepState;
     }
     /**
-     * Move a step to `running` only if it is still in one of `from`, reporting
-     * whether this caller is the one that moved it.
+     * Take back a failed or lapsed step as a new attempt. The attempt number is
+     * advanced by the same guarded `UPDATE` that grants the lease, so the worker
+     * it is taken from is fenced out of its step at once — deferring that to
+     * `createRetryAttempt` would leave a window in which its result still lands,
+     * only to be cleared and the step run again.
+     */
+    async reclaimStepAttempt(stepState, attempt, leaseMs) {
+        const now = this.nowMs();
+        const claimed = await this.db.transaction().execute(async (trx) => {
+            let query = trx
+                .updateTable('workflowStep')
+                .set({
+                status: 'running',
+                result: null,
+                error: null,
+                currentAttempt: attempt,
+                leaseExpiresAt: leaseUntil(now, leaseMs),
+                updatedAt: new Date(),
+            })
+                .where('workflowStepId', '=', stepState.stepId)
+                .where('status', '=', stepState.status)
+                .where(sql `coalesce(current_attempt, 1)`, '=', stepState.attemptCount);
+            if (stepState.status === 'running') {
+                query = query.where('leaseExpiresAt', '<', now);
+            }
+            const moved = await query.executeTakeFirst();
+            if (Number(moved?.numUpdatedRows ?? 0n) === 0)
+                return false;
+            await this.insertHistoryRecord(stepState.stepId, 'running', attempt, undefined, undefined, trx);
+            return true;
+        });
+        return claimed ? this.readStepAttempt(stepState.stepId) : null;
+    }
+    /**
+     * Move a step to `running` under a fresh lease, only if it is still in one of
+     * `from`, reporting whether this caller is the one that moved it.
+     *
+     * Claiming back a `running` step additionally requires its lease to have
+     * lapsed. A null lease never matches, which is what keeps a step parked on a
+     * child run — running, with no worker on it by design — out of reach.
      */
-    async claimStepStatus(stepId, from) {
-        const claimed = await this.db
+    async claimStepStatus(stepId, from, leaseMs) {
+        const now = this.nowMs();
+        let query = this.db
             .updateTable('workflowStep')
-            .set({ status: 'running', updatedAt: new Date() })
+            .set({
+            status: 'running',
+            leaseExpiresAt: leaseUntil(now, leaseMs),
+            updatedAt: new Date(),
+        })
             .where('workflowStepId', '=', stepId)
-            .where('status', 'in', from)
-            .executeTakeFirst();
+            .where('status', 'in', from);
+        if (from.includes('running')) {
+            query = query.where('leaseExpiresAt', '<', now);
+        }
+        const claimed = await query.executeTakeFirst();
         return Number(claimed?.numUpdatedRows ?? 0n) > 0;
     }
     async getCompletedGraphState(runId) {
@@ -465,13 +543,16 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     async getStepInstances(runId) {
         const rows = await this.db
             .selectFrom('workflowStep')
-            .select(['stepName', 'status', 'fromStepName'])
+            .select(['stepName', 'status', 'fromStepName', 'leaseExpiresAt'])
             .where('workflowRunId', '=', runId)
             .execute();
         return rows.map((r) => ({
             stepName: r.stepName,
             status: r.status,
             fromStepName: r.fromStepName ?? undefined,
+            leaseExpiresAt: r.leaseExpiresAt != null
+                ? new Date(Number(r.leaseExpiresAt))
+                : undefined,
         }));
     }
     async getNodeResults(runId, nodeIds) {
@@ -537,7 +618,7 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
      * The undispatched-step query, shared by every dialect.
      *
      * It used to be deliberately NOT an override, because this class's
-     * `withStepLock` is a pass-through: a subclass inheriting it (kysely-sqlite)
+     * `withStepLock` was a pass-through: a subclass inheriting it (kysely-sqlite)
      * had no atomic claim, so the relay's redundant dispatches would have become
      * double executions, and only `kysely-postgres` and `kysely-mysql` opted in
      * on the strength of their real locks. `claimStepForExecution` no longer
@@ -571,7 +652,19 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             .where(({ not, exists, selectFrom }) => not(exists(selectFrom('workflowStep as s')
             .select('s.workflowStepId')
             .whereRef('s.workflowRunId', '=', 'r.workflowRunId')
-            .where('s.status', 'in', ['running', 'scheduled', 'suspended']))))
+            .where((eb) => eb.or([
+            eb('s.status', 'in', ['scheduled', 'suspended']),
+            // A `running` step counts as in flight only while someone
+            // still holds it. A lapsed lease is a dispatch that died, and
+            // a null one a step nothing was expected to hand back.
+            eb.and([
+                eb('s.status', '=', 'running'),
+                eb.or([
+                    eb('s.leaseExpiresAt', 'is', null),
+                    eb('s.leaseExpiresAt', '>', this.nowMs()),
+                ]),
+            ]),
+        ])))))
             // The run row only moves on a status change, so idleness has to hold for
             // the steps too — expressed as "no step touched since `before`" rather
             // than max(updated_at) so it stays one index-friendly NOT EXISTS.

package/dist/src/schema/agent.schema.js

@@ -13,89 +13,96 @@ export const agentSchema = {
     name: 'agent',
     ownedBy: ['agentStorage', 'agentRunState', 'agentRunService'],
     statements: [
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentThreads')
             .addColumn('id', 'varchar(36)', (col) => col.primaryKey())
             .addColumn('resourceId', 'varchar(255)', (col) => col.notNull())
-            .addColumn('title', 'text')
-            .addColumn('metadata', 'text')
+            .addColumn('title', ctx.text)
+            .addColumn('metadata', ctx.text)
             .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
             .addColumn('updatedAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull()),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_threads_resource')
             .on('agentThreads')
             .column('resourceId'),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentMessage')
             .addColumn('id', 'varchar(36)', (col) => col.primaryKey())
-            .addColumn('threadId', 'varchar(36)', (col) => col.notNull().references('agentThreads.id').onDelete('cascade'))
+            .addColumn('threadId', 'varchar(36)', (col) => ctx.references(col.notNull(), 'agentThreads.id'))
             .addColumn('role', 'varchar(50)', (col) => col.notNull())
-            .addColumn('content', 'text')
-            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull()),
-        (db) => db.schema
+            .addColumn('content', ctx.text)
+            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .$call(ctx.foreignKeys('agentMessage', { threadId: 'agentThreads.id' })),
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_message_thread')
             .on('agentMessage')
             .columns(['threadId', 'createdAt']),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentToolCall')
             .addColumn('id', 'varchar(36)', (col) => col.primaryKey())
-            .addColumn('threadId', 'varchar(36)', (col) => col.notNull().references('agentThreads.id').onDelete('cascade'))
-            .addColumn('messageId', 'varchar(36)', (col) => col.notNull().references('agentMessage.id').onDelete('cascade'))
+            .addColumn('threadId', 'varchar(36)', (col) => ctx.references(col.notNull(), 'agentThreads.id'))
+            .addColumn('messageId', 'varchar(36)', (col) => ctx.references(col.notNull(), 'agentMessage.id'))
             .addColumn('runId', 'varchar(36)')
             .addColumn('toolName', 'varchar(255)', (col) => col.notNull())
-            .addColumn('args', 'text', (col) => col.notNull())
-            .addColumn('result', 'text')
+            .addColumn('args', ctx.text, (col) => col.notNull())
+            .addColumn('result', ctx.text)
             .addColumn('approvalStatus', 'varchar(50)')
             .addColumn('approvalType', 'varchar(50)')
             .addColumn('agentRunId', 'varchar(36)')
             .addColumn('displayToolName', 'varchar(255)')
-            .addColumn('displayArgs', 'text')
-            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull()),
-        (db) => db.schema
+            .addColumn('displayArgs', ctx.text)
+            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .$call(ctx.foreignKeys('agentToolCall', {
+            threadId: 'agentThreads.id',
+            messageId: 'agentMessage.id',
+        })),
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_tool_call_thread')
             .on('agentToolCall')
             .column('threadId'),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_tool_call_message')
             .on('agentToolCall')
             .column('messageId'),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentWorkingMemory')
             .addColumn('id', 'varchar(255)', (col) => col.notNull())
             .addColumn('scope', 'varchar(50)', (col) => col.notNull())
-            .addColumn('data', 'text', (col) => col.notNull())
+            .addColumn('data', ctx.text, (col) => col.notNull())
             .addColumn('updatedAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
             .addPrimaryKeyConstraint('agent_working_memory_pk', ['id', 'scope']),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentRun')
             .addColumn('runId', 'varchar(36)', (col) => col.primaryKey())
             .addColumn('agentName', 'varchar(255)', (col) => col.notNull())
-            .addColumn('threadId', 'varchar(36)', (col) => col.notNull().references('agentThreads.id').onDelete('cascade'))
+            .addColumn('threadId', 'varchar(36)', (col) => ctx.references(col.notNull(), 'agentThreads.id'))
             .addColumn('resourceId', 'varchar(255)', (col) => col.notNull())
-            .addColumn('status', 'varchar(50)', (col) => col.notNull().defaultTo('running'))
-            .addColumn('errorMessage', 'text')
-            .addColumn('suspendReason', 'text')
-            .addColumn('missingRpcs', 'text')
-            .addColumn('pendingApprovals', 'text')
+            .addColumn('status', 'varchar(50)', (col) => col.notNull().defaultTo(ctx.defaultText('running')))
+            .addColumn('errorMessage', ctx.text)
+            .addColumn('suspendReason', ctx.text)
+            .addColumn('missingRpcs', ctx.text)
+            .addColumn('pendingApprovals', ctx.text)
             .addColumn('usageInputTokens', 'integer', (col) => col.notNull().defaultTo(0))
             .addColumn('usageOutputTokens', 'integer', (col) => col.notNull().defaultTo(0))
-            .addColumn('usageModel', 'varchar(255)', (col) => col.notNull().defaultTo(''))
+            .addColumn('usageModel', 'varchar(255)', (col) => col.notNull().defaultTo(ctx.defaultText('')))
             .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
-            .addColumn('updatedAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull()),
-        (db) => db.schema
+            .addColumn('updatedAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .$call(ctx.foreignKeys('agentRun', { threadId: 'agentThreads.id' })),
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_run_thread')
             .on('agentRun')
             .columns(['threadId', 'createdAt']),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('agentRunScore')
             .addColumn('id', 'varchar(36)', (col) => col.primaryKey())
-            .addColumn('runId', 'varchar(36)', (col) => col.notNull().references('agentRun.runId').onDelete('cascade'))
+            .addColumn('runId', 'varchar(36)', (col) => ctx.references(col.notNull(), 'agentRun.runId'))
             .addColumn('scorerName', 'varchar(255)', (col) => col.notNull())
             .addColumn('score', 'real', (col) => col.notNull())
-            .addColumn('reason', 'text')
-            .addColumn('metadata', 'text')
-            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull()),
-        (db) => db.schema
+            .addColumn('reason', ctx.text)
+            .addColumn('metadata', ctx.text)
+            .addColumn('createdAt', 'timestamp', (col) => col.defaultTo(sql `CURRENT_TIMESTAMP`).notNull())
+            .$call(ctx.foreignKeys('agentRunScore', { runId: 'agentRun.runId' })),
+        (db, _types, ctx) => db.schema
             .createIndex('idx_agent_run_score_run')
             .on('agentRunScore')
             .columns(['runId', 'createdAt']),

package/dist/src/schema/analytics.schema.js

@@ -17,30 +17,30 @@ export const analyticsSchema = {
     name: 'analytics',
     ownedBy: ['analyticsService'],
     statements: [
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createTable('pikkuAnalyticsEvents')
-            .addColumn('eventId', 'text', (col) => col.primaryKey())
-            .addColumn('name', 'text', (col) => col.notNull())
-            .addColumn('occurredAt', 'text', (col) => col.notNull())
-            .addColumn('source', 'text', (col) => col.notNull())
-            .addColumn('functionId', 'text')
-            .addColumn('wireType', 'text')
-            .addColumn('traceId', 'text')
-            .addColumn('userId', 'text')
-            .addColumn('orgId', 'text')
-            .addColumn('pikkuUserId', 'text')
+            .addColumn('eventId', ctx.key, (col) => col.primaryKey())
+            .addColumn('name', ctx.key, (col) => col.notNull())
+            .addColumn('occurredAt', ctx.key, (col) => col.notNull())
+            .addColumn('source', ctx.text, (col) => col.notNull())
+            .addColumn('functionId', ctx.text)
+            .addColumn('wireType', ctx.text)
+            .addColumn('traceId', ctx.text)
+            .addColumn('userId', ctx.text)
+            .addColumn('orgId', ctx.text)
+            .addColumn('pikkuUserId', ctx.text)
             // The device id for a visitor with no session. Its own column beside
             // `userId` and not folded into it: the whole point of the anonymous id
             // is the traffic that has no user to be attributed to yet.
-            .addColumn('anonymousId', 'text')
-            .addColumn('vendorIds', 'text')
-            .addColumn('consent', 'text')
-            .addColumn('props', 'text'),
-        (db) => db.schema
+            .addColumn('anonymousId', ctx.text)
+            .addColumn('vendorIds', ctx.text)
+            .addColumn('consent', ctx.text)
+            .addColumn('props', ctx.text),
+        (db, _types, ctx) => db.schema
             .createIndex('idx_pikku_analytics_events_occurred_at')
             .on('pikkuAnalyticsEvents')
             .column('occurredAt'),
-        (db) => db.schema
+        (db, _types, ctx) => db.schema
             .createIndex('idx_pikku_analytics_events_name')
             .on('pikkuAnalyticsEvents')
             .column('name'),