@pikku/kysely 0.13.19 → 0.13.21

package/src/kysely-virtual-user-run-store.test.ts

@@ -6,6 +6,7 @@ import { CamelCasePlugin, Kysely, SqliteDialect } from 'kysely'
 import { KyselyVirtualUserRunStore } from './kysely-virtual-user-run-store.js'
 import { SerializePlugin } from './serialize-plugin.js'
 import type { KyselyPikkuDB } from './kysely-tables.js'
+import { applyPikkuSchemas, virtualUserSchema } from './schema/index.js'
 
 // SerializePlugin is here because most projects in this package run with it, and
 // it is the harder case: it deserialises JSON columns the store has already
@@ -16,6 +17,12 @@ const newDb = () =>
     plugins: [new CamelCasePlugin(), new SerializePlugin()],
   })
 
+const migratedDb = async () => {
+  const db = newDb()
+  await applyPikkuSchemas(db, [virtualUserSchema])
+  return db
+}
+
 const TALLY = {
   steps: 12,
   calls: 30,
@@ -32,19 +39,49 @@ const FINDING = {
   step: 7,
 } as any
 
+const INTENT = {
+  id: 'intent_1',
+  sourceId: 'exportADoc',
+  title: 'Export a doc',
+  status: 'done',
+  steps: [0, 1],
+  suspensions: 0,
+} as any
+
+const step = (index: number, over: Record<string, unknown> = {}) =>
+  ({
+    index,
+    intentId: 'intent_1',
+    action: { kind: 'call', rpcName: 'listDocs', input: {} },
+    status: 200,
+    ok: true,
+    response: '{"docs":[]}',
+    tokensIn: 40,
+    tokensOut: 12,
+    ...over,
+  }) as any
+
 describe('KyselyVirtualUserRunStore', () => {
   let db: Kysely<KyselyPikkuDB>
   let store: KyselyVirtualUserRunStore
 
-  beforeEach(() => {
-    db = newDb()
+  beforeEach(async () => {
+    db = await migratedDb()
     store = new KyselyVirtualUserRunStore(db)
   })
 
-  // The runtime never needs this table, so it arrives with the feature that
-  // fills it rather than in every database — which means the first write has to
-  // create it.
-  test('creates its table on first use', async () => {
+  // The runtime never creates this table, so a project that skipped the
+  // migration has to be told which command writes it rather than quietly
+  // getting one at first use.
+  test('refuses to start against a database that never migrated it', async () => {
+    const store = new KyselyVirtualUserRunStore(newDb())
+    await assert.rejects(
+      store.start({ persona: 'susan', disposition: 'realistic', seed: 42 }),
+      /pikku db generate/
+    )
+  })
+
+  test('writes against a migrated table', async () => {
     const runId = await store.start({
       persona: 'susan',
       disposition: 'realistic',
@@ -99,6 +136,8 @@ describe('KyselyVirtualUserRunStore', () => {
       tally: TALLY,
       memory: { docId: 'doc_1', orgId: 'org_9' },
       stoppedBy: 'budget-steps',
+      intents: [INTENT],
+      steps: [step(0), step(1)],
     })
 
     const run = await store.get(runId)
@@ -167,14 +206,17 @@ describe('KyselyVirtualUserRunStore', () => {
     assert.equal(runs[0]?.persona, 'susan')
   })
 
-  test('completing an unknown run does not throw', async () => {
+  test('completing an unknown run does not throw, and leaves no steps behind', async () => {
     await store.init()
     await store.complete('nope', {
       findings: [],
       tally: TALLY,
       memory: {},
       stoppedBy: 'exhausted',
+      intents: [INTENT],
+      steps: [step(0)],
     })
+    assert.deepEqual(await store.steps('nope'), [])
   })
 
   // Not every project installs SerializePlugin, and a bare SQLite driver cannot
@@ -184,6 +226,7 @@ describe('KyselyVirtualUserRunStore', () => {
       dialect: new SqliteDialect({ database: new Database(':memory:') }),
       plugins: [new CamelCasePlugin()],
     })
+    await applyPikkuSchemas(bare, [virtualUserSchema])
     const bareStore = new KyselyVirtualUserRunStore(bare)
 
     const runId = await bareStore.start({
@@ -198,6 +241,8 @@ describe('KyselyVirtualUserRunStore', () => {
       tally: TALLY,
       memory: { docId: 'doc_1' },
       stoppedBy: 'exhausted',
+      intents: [INTENT],
+      steps: [step(0)],
     })
 
     const run = await bareStore.get(runId)
@@ -208,5 +253,117 @@ describe('KyselyVirtualUserRunStore', () => {
     assert.equal(run?.findings[0]?.detail, FINDING.detail)
     assert.ok(run?.createdAt instanceof Date)
     assert.ok(run?.finishedAt instanceof Date)
+    assert.deepEqual(await bareStore.steps(runId), [step(0)])
+  })
+
+  test('the transcript comes back in the order it happened', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 1,
+    })
+
+    await store.complete(runId, {
+      findings: [],
+      tally: TALLY,
+      memory: {},
+      stoppedBy: 'exhausted',
+      intents: [INTENT],
+      steps: [step(2), step(0), step(1)],
+    })
+
+    const steps = await store.steps(runId)
+    assert.deepEqual(
+      steps.map((one) => one.index),
+      [0, 1, 2]
+    )
+    assert.deepEqual(steps[0], step(0))
+  })
+
+  // The turn the model got wrong is the one worth reading, and it carries none
+  // of the fields a successful call does.
+  test('a turn with no response and no status survives the round trip', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'careless',
+      seed: 1,
+    })
+
+    const invalid = step(0, {
+      intentId: undefined,
+      action: { kind: 'invalid', detail: 'no such rpc: listDoc' },
+      status: undefined,
+      ok: undefined,
+      response: undefined,
+      findingKinds: ['invalid-action'],
+    })
+
+    await store.complete(runId, {
+      findings: [],
+      tally: TALLY,
+      memory: {},
+      stoppedBy: 'exhausted',
+      intents: [],
+      steps: [invalid],
+    })
+
+    const [stored] = await store.steps(runId)
+    assert.deepEqual(stored, {
+      index: 0,
+      action: { kind: 'invalid', detail: 'no such rpc: listDoc' },
+      findingKinds: ['invalid-action'],
+      tokensIn: 40,
+      tokensOut: 12,
+    })
+  })
+
+  // Ten columns times a 500-step budget is five thousand bound variables, and a
+  // bare sqlite driver takes 999 — so the long run is exactly the one that must
+  // not be lost.
+  test('a run longer than one insert can hold is stored whole', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'adversarial',
+      seed: 1,
+    })
+
+    const steps = Array.from({ length: 240 }, (_, index) => step(index))
+    await store.complete(runId, {
+      findings: [],
+      tally: TALLY,
+      memory: {},
+      stoppedBy: 'budget-steps',
+      intents: [],
+      steps,
+    })
+
+    assert.equal((await store.steps(runId, { limit: 500 })).length, 240)
+    assert.deepEqual(
+      (await store.steps(runId, { limit: 2, offset: 100 })).map(
+        (one) => one.index
+      ),
+      [100, 101]
+    )
+  })
+
+  test('intents ride on the run, not on the transcript', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 1,
+    })
+
+    await store.complete(runId, {
+      findings: [],
+      tally: TALLY,
+      memory: {},
+      stoppedBy: 'exhausted',
+      intents: [INTENT],
+      steps: [],
+    })
+
+    const [listed] = await store.list({ persona: 'susan' })
+    assert.deepEqual(listed?.intents, [INTENT])
+    assert.deepEqual(await store.steps(runId), [])
   })
 })

package/src/kysely-virtual-user-run-store.ts

@@ -1,15 +1,22 @@
 import type { Kysely, Selectable } from 'kysely'
 import { randomUUID } from 'node:crypto'
-import type { VirtualUserDisposition } from '@pikku/core/virtual-user'
+import type {
+  StepRecord,
+  VirtualUserDisposition,
+} from '@pikku/core/virtual-user'
 import type {
   VirtualUserRunOutcome,
   VirtualUserRunRecord,
   VirtualUserRunStart,
   VirtualUserRunStore,
 } from '@pikku/core/virtual-user'
-import type { KyselyPikkuDB, VirtualUserRunTable } from './kysely-tables.js'
+import type {
+  KyselyPikkuDB,
+  VirtualUserRunStepTable,
+  VirtualUserRunTable,
+} from './kysely-tables.js'
 import { parseJson } from './kysely-json.js'
-import { ensurePikkuSchema } from './schema/index.js'
+import { requirePikkuSchema } from './schema/index.js'
 import { virtualUserSchema } from './schema/virtual-user.schema.js'
 
 /**
@@ -31,6 +38,16 @@ import { virtualUserSchema } from './schema/virtual-user.schema.js'
  * ids. This is a privileged store; the scaffolded reads are scope-gated, and
  * exposing these rows more widely publishes your own exploits.
  */
+/**
+ * How many step rows go in one insert.
+ *
+ * Not a performance dial. A bare SQLite driver binds at most 999 variables per
+ * statement, and ten columns times a 500-step budget is five thousand — so an
+ * un-chunked insert does not fail on a long run, it fails on the *interesting*
+ * long run, which is the one nobody wants to lose.
+ */
+const STEP_INSERT_CHUNK = 50
+
 export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
   private initialized = false
 
@@ -43,7 +60,7 @@ export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
    */
   public async init(): Promise<void> {
     if (this.initialized) return
-    await ensurePikkuSchema(this.db, virtualUserSchema)
+    await requirePikkuSchema(this.db, virtualUserSchema)
     this.initialized = true
   }
 
@@ -71,7 +88,7 @@ export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
 
   async complete(runId: string, outcome: VirtualUserRunOutcome): Promise<void> {
     await this.init()
-    await this.db
+    const updated = await this.db
       .updateTable('virtualUserRun')
       .set({
         status: 'completed',
@@ -80,11 +97,16 @@ export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
         // Overwritten rather than merged: the engine's memory already carries
         // what it was given, plus what it learned on the way.
         memory: JSON.stringify(outcome.memory),
+        intents: JSON.stringify(outcome.intents),
         stoppedBy: outcome.stoppedBy,
         finishedAt: new Date().toISOString(),
       })
       .where('runId', '=', runId)
-      .execute()
+      .executeTakeFirst()
+    // Nothing owns steps written against a run that is not there, and nothing
+    // would ever read or reap them — there is no foreign key to refuse them.
+    if (updated.numUpdatedRows === 0n) return
+    await this.writeSteps(runId, outcome.steps)
   }
 
   async fail(runId: string, error: string): Promise<void> {
@@ -128,6 +150,69 @@ export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
     return rows.map((row) => this.toRecord(row))
   }
 
+  async steps(
+    runId: string,
+    options?: { limit?: number; offset?: number }
+  ): Promise<StepRecord[]> {
+    await this.init()
+    const rows = await this.db
+      .selectFrom('virtualUserRunStep')
+      .selectAll()
+      .where('runId', '=', runId)
+      .orderBy('stepIndex', 'asc')
+      .limit(options?.limit ?? 1000)
+      .offset(options?.offset ?? 0)
+      .execute()
+    return rows.map((row) => this.toStep(row))
+  }
+
+  private async writeSteps(
+    runId: string,
+    steps: readonly StepRecord[]
+  ): Promise<void> {
+    for (let at = 0; at < steps.length; at += STEP_INSERT_CHUNK) {
+      const chunk = steps.slice(at, at + STEP_INSERT_CHUNK)
+      await this.db
+        .insertInto('virtualUserRunStep')
+        .values(
+          chunk.map((step) => ({
+            runId,
+            stepIndex: step.index,
+            intentId: step.intentId ?? null,
+            action: JSON.stringify(step.action),
+            status: step.status ?? null,
+            ok: step.ok === undefined ? null : step.ok ? 1 : 0,
+            response:
+              step.response === undefined
+                ? null
+                : JSON.stringify(step.response),
+            findingKinds: step.findingKinds
+              ? JSON.stringify(step.findingKinds)
+              : null,
+            tokensIn: step.tokensIn,
+            tokensOut: step.tokensOut,
+          }))
+        )
+        .execute()
+    }
+  }
+
+  private toStep(row: Selectable<VirtualUserRunStepTable>): StepRecord {
+    return {
+      index: Number(row.stepIndex),
+      ...(row.intentId ? { intentId: row.intentId } : {}),
+      action: parseJson(row.action),
+      ...(row.status === null ? {} : { status: Number(row.status) }),
+      ...(row.ok === null ? {} : { ok: Boolean(row.ok) }),
+      ...(row.response === null ? {} : { response: parseJson(row.response) }),
+      ...(row.findingKinds
+        ? { findingKinds: parseJson(row.findingKinds) }
+        : {}),
+      tokensIn: Number(row.tokensIn),
+      tokensOut: Number(row.tokensOut),
+    }
+  }
+
   private toRecord(row: Selectable<VirtualUserRunTable>): VirtualUserRunRecord {
     return {
       runId: row.runId,
@@ -138,6 +223,7 @@ export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
       goals: parseJson(row.goals) ?? [],
       memory: parseJson(row.memory) ?? {},
       findings: parseJson(row.findings) ?? [],
+      intents: parseJson(row.intents) ?? [],
       tally: row.tally ? (parseJson(row.tally) ?? null) : null,
       stoppedBy: row.stoppedBy,
       error: row.error,

package/src/kysely-virtual-user-schedule-store.test.ts

@@ -0,0 +1,187 @@
+import { describe, test, beforeEach } from 'node:test'
+import assert from 'node:assert/strict'
+import Database from 'better-sqlite3'
+import { CamelCasePlugin, Kysely, SqliteDialect } from 'kysely'
+
+import { KyselyVirtualUserScheduleStore } from './kysely-virtual-user-schedule-store.js'
+import { SerializePlugin } from './serialize-plugin.js'
+import { applyPikkuSchemas, virtualUserScheduleSchema } from './schema/index.js'
+import type { KyselyPikkuDB } from './kysely-tables.js'
+
+// With the plugin, because most projects in this package run with it and it is
+// the harder case: it deserialises columns the store has already stringified.
+const withPlugin = () =>
+  new Kysely<KyselyPikkuDB>({
+    dialect: new SqliteDialect({ database: new Database(':memory:') }),
+    plugins: [new CamelCasePlugin(), new SerializePlugin()],
+  })
+
+const bare = () =>
+  new Kysely<KyselyPikkuDB>({
+    dialect: new SqliteDialect({ database: new Database(':memory:') }),
+    plugins: [new CamelCasePlugin()],
+  })
+
+const HOUR = 60 * 60 * 1000
+
+describe('KyselyVirtualUserScheduleStore', () => {
+  let db: Kysely<KyselyPikkuDB>
+  let store: KyselyVirtualUserScheduleStore
+
+  beforeEach(async () => {
+    db = withPlugin()
+    await applyPikkuSchemas(db, [virtualUserScheduleSchema])
+    store = new KyselyVirtualUserScheduleStore(db)
+  })
+
+  test('a persona given a cadence is not running until someone says so', async () => {
+    const written = await store.set({ persona: 'guest' })
+    assert.equal(written.enabled, false)
+    assert.deepEqual(await store.due(new Date(Date.now() + 1000)), [])
+    assert.equal((await store.get('guest'))?.enabled, false)
+  })
+
+  test('a schedule written without an interval gets one, sparse by default', async () => {
+    const written = await store.set({ persona: 'guest' })
+    assert.ok(written.minIntervalMs >= HOUR)
+    assert.ok(written.maxIntervalMs >= written.minIntervalMs)
+    const stored = await store.get('guest')
+    assert.equal(stored?.minIntervalMs, written.minIntervalMs)
+    assert.equal(stored?.maxIntervalMs, written.maxIntervalMs)
+  })
+
+  test('setting one field leaves the rest of the schedule alone', async () => {
+    await store.set({
+      persona: 'guest',
+      enabled: true,
+      disposition: 'adversarial',
+      goals: ['book a room'],
+      budget: { steps: 40 },
+      minIntervalMs: HOUR,
+      maxIntervalMs: 2 * HOUR,
+    })
+    await store.set({ persona: 'guest', enabled: false })
+
+    const stored = await store.get('guest')
+    assert.equal(stored?.enabled, false)
+    assert.equal(stored?.disposition, 'adversarial')
+    assert.deepEqual(stored?.goals, ['book a room'])
+    assert.deepEqual(stored?.budget, { steps: 40 })
+    assert.equal(stored?.minIntervalMs, HOUR)
+  })
+
+  test('only enabled personas whose time has come come back as due', async () => {
+    const now = new Date()
+    await store.set({
+      persona: 'due',
+      enabled: true,
+      nextRunAt: new Date(now.getTime() - HOUR),
+    })
+    await store.set({
+      persona: 'later',
+      enabled: true,
+      nextRunAt: new Date(now.getTime() + HOUR),
+    })
+    await store.set({
+      persona: 'off',
+      enabled: false,
+      nextRunAt: new Date(now.getTime() - HOUR),
+    })
+
+    assert.deepEqual(
+      (await store.due(now)).map((row) => row.persona),
+      ['due']
+    )
+  })
+
+  test('claiming pushes the persona out of the due list', async () => {
+    const now = new Date()
+    const due = new Date(now.getTime() - HOUR)
+    await store.set({ persona: 'guest', enabled: true, nextRunAt: due })
+    const acquired = await store.claim('guest', {
+      from: due,
+      nextRunAt: new Date(now.getTime() + HOUR),
+      runId: 'run-1',
+      at: now,
+    })
+
+    assert.equal(acquired, true)
+
+    assert.deepEqual(await store.due(now), [])
+    const stored = await store.get('guest')
+    assert.equal(stored?.lastRunId, 'run-1')
+    assert.equal(stored?.lastRunAt?.getTime(), now.getTime())
+  })
+
+  test('the claim written before a run starts does not blank the last run', async () => {
+    const now = new Date()
+    const written = await store.set({ persona: 'guest', enabled: true })
+    await store.claim('guest', {
+      from: written.nextRunAt,
+      nextRunAt: now,
+      runId: 'run-1',
+      at: now,
+    })
+    await store.claim('guest', {
+      from: now,
+      nextRunAt: new Date(now.getTime() + HOUR),
+      runId: null,
+      at: now,
+    })
+
+    assert.equal((await store.get('guest'))?.lastRunId, 'run-1')
+  })
+
+  test('only one of two ticks holding the same due row can claim it', async () => {
+    const now = new Date()
+    const due = new Date(now.getTime() - HOUR)
+    await store.set({ persona: 'guest', enabled: true, nextRunAt: due })
+
+    const [first] = await store.due(now)
+    const [second] = await store.due(now)
+    assert.equal(first?.nextRunAt.getTime(), second?.nextRunAt.getTime())
+
+    const won = await store.claim('guest', {
+      from: first!.nextRunAt,
+      nextRunAt: new Date(now.getTime() + HOUR),
+      runId: 'run-1',
+      at: now,
+    })
+    const lost = await store.claim('guest', {
+      from: second!.nextRunAt,
+      nextRunAt: new Date(now.getTime() + 2 * HOUR),
+      runId: 'run-2',
+      at: now,
+    })
+
+    assert.deepEqual([won, lost], [true, false])
+    assert.equal((await store.get('guest'))?.lastRunId, 'run-1')
+  })
+
+  test('a cadence survives a store with no serialize plugin', async () => {
+    const plainDb = bare()
+    await applyPikkuSchemas(plainDb, [virtualUserScheduleSchema])
+    const plain = new KyselyVirtualUserScheduleStore(plainDb)
+    await plain.set({
+      persona: 'guest',
+      enabled: true,
+      goals: ['book a room'],
+      budget: { steps: 40, duration: '30m' },
+      minIntervalMs: HOUR,
+      maxIntervalMs: 2 * HOUR,
+    })
+
+    const stored = await plain.get('guest')
+    assert.equal(stored?.enabled, true)
+    assert.deepEqual(stored?.goals, ['book a room'])
+    assert.deepEqual(stored?.budget, { steps: 40, duration: '30m' })
+    assert.ok(stored?.nextRunAt instanceof Date)
+  })
+
+  test('a persona can be taken off the clock entirely', async () => {
+    await store.set({ persona: 'guest', enabled: true })
+    await store.remove('guest')
+    assert.equal(await store.get('guest'), null)
+    assert.deepEqual(await store.list(), [])
+  })
+})