ConusAI vendor v0.1.0-beta.114

modules/agent/__tests__/claimed-job.ts

@@ -0,0 +1,173 @@
+// SPDX-License-Identifier: LicenseRef-ConusAI-Template-1.0
+// Real pg-boss claims for Agent tests, including local engine/provider tests.
+import { buildRegistry, Container, defineModule, type JobContext, newId, type ProjectScope } from "@conusai-kit/core"
+import { systemActor } from "@conusai-kit/core/worker"
+import { type DataSource, entitiesOf } from "@conusai-kit/data"
+import { connectKyselyDataSourceFromEnv, type KyselyDataSource } from "@conusai-kit/data-kysely"
+import {
+  createPgBoss,
+  createPgBossClaimGuard,
+  createPgBossJobQueue,
+  createWorker,
+  fromKyselyTransaction,
+  withJobTransaction,
+} from "@conusai-kit/jobs-pgboss"
+import { Pool, type PoolClient } from "pg"
+import { afterAll, beforeAll } from "vitest"
+import { z } from "zod"
+import { runJob } from "../adapters/procedures.js"
+import { RUN_JOB } from "../contract/index.js"
+import { AGENT_SCOPE } from "../core/ports.js"
+import { grantPgBoss, migrate, type Postgres, startPostgres } from "./claimed-postgres.js"
+import { workerRegistry } from "./worker-fixture.js"
+
+let postgres: Postgres
+let boss: ReturnType<typeof createPgBoss>
+let guard: ReturnType<typeof createPgBossClaimGuard>
+let sql: Pool
+let ownedClient: PoolClient
+/** Terminate only the client acquired from this test's dedicated guard pool. */
+export async function terminateClaimSession() {
+  const {
+    rows: [{ pid }],
+  } = await ownedClient.query("select pg_backend_pid() as pid")
+  await sql.query("select pg_terminate_backend($1)", [pid])
+}
+export type ClaimChange = "cancel" | "reclaim" | "delete" | "scope" | "actor" | "owner"
+const workers: ReturnType<typeof createWorker>[] = []
+const pending = new Map<
+  string,
+  {
+    container: Container
+    extra: Partial<JobContext>
+    change?: ClaimChange
+    resolve: () => void
+    reject: (error: unknown) => void
+  }
+>()
+let domainData: KyselyDataSource | undefined
+/** Queue and domain tables share the actual database, as in the supported host. */
+export async function claimedDataSource(): Promise<KyselyDataSource> {
+  if (!domainData) {
+    const entities = entitiesOf(workerRegistry, "postgres")
+    await migrate(postgres, entities)
+    // Migrations grant access to public tables; queue objects remain owned by pg-boss.
+    domainData = await connectKyselyDataSourceFromEnv(
+      { POSTGRES_URL: postgres.runtimeUrl, POSTGRES_READER_URL: postgres.readerUrl },
+      { entities, transactionContext: (trx, run) => withJobTransaction(fromKyselyTransaction(trx), run) },
+    )
+  }
+  return domainData
+}
+let currentData: DataSource
+const registry = buildRegistry([
+  defineModule({
+    id: "agent",
+    version: "1.0.0",
+    kit: "^0.1.0",
+    jobs: {
+      [RUN_JOB]: {
+        payload: z.object({ runId: z.string(), delivery: z.string() }),
+        async handler(payload, ctx) {
+          const target = pending.get(payload.delivery)
+          if (!target) throw new Error("unregistered test delivery")
+          try {
+            if (!ctx.execution) throw new Error("supported adapter did not supply execution")
+            const id = ctx.execution.id
+            if (target.change === "delete") await sql.query("delete from pgboss.job where id = $1", [id])
+            if (target.change === "cancel")
+              await sql.query("update pgboss.job set state = 'cancelled' where id = $1", [id])
+            if (target.change === "reclaim")
+              await sql.query("update pgboss.job set retry_count = retry_count + 1 where id = $1", [id])
+            if (target.change === "scope")
+              await sql.query(
+                "update pgboss.job set data = jsonb_set(data, '{scope,projectId}', to_jsonb($2::text)) where id = $1",
+                [id, newId()],
+              )
+            if (target.change === "actor")
+              await sql.query(
+                "update pgboss.job set data = jsonb_set(data, '{onBehalfOf}', to_jsonb($2::text)) where id = $1",
+                [id, newId()],
+              )
+            await runJob.handler(
+              { runId: payload.runId },
+              { ...ctx, ...target.extra, execution: ctx.execution, container: target.container },
+            )
+            target.resolve()
+          } catch (error) {
+            target.reject(error)
+          } finally {
+            pending.delete(payload.delivery)
+          }
+        },
+      },
+    },
+  }),
+])
+beforeAll(async () => {
+  postgres = await startPostgres()
+  boss = createPgBoss({ connectionString: postgres.ownerUrl })
+  await boss.start()
+  await boss.createQueue(RUN_JOB)
+  await grantPgBoss(postgres)
+  sql = new Pool({ connectionString: postgres.runtimeUrl, max: 2 })
+  const guardPool = new Pool({ connectionString: postgres.runtimeUrl, max: 2 })
+  guardPool.on("acquire", (client) => {
+    ownedClient = client
+  })
+  guard = createPgBossClaimGuard(guardPool)
+}, 300_000)
+afterAll(async () => {
+  for (const worker of workers) await worker.stop()
+  await guard?.close()
+  await sql?.end()
+  await boss?.stop({ graceful: true })
+  await domainData?.close()
+  await postgres?.stop()
+})
+let started = false
+export async function claimedDelivery(
+  data: DataSource,
+  container: Container,
+  scope: ProjectScope,
+  runId: string,
+  extra: Partial<JobContext> = {},
+  change?: ClaimChange,
+) {
+  currentData = data
+  if (!started) {
+    // Source forwarding is needed because a file exercises several independent tenancy fixtures.
+    const source = new Proxy({} as DataSource, {
+      get(_target, property) {
+        const value = Reflect.get(currentData, property)
+        return typeof value === "function" ? value.bind(currentData) : value
+      },
+    })
+    workers.length = 0
+    const worker = createWorker({
+      registry,
+      boss,
+      claimGuard: guard,
+      dataSource: source,
+      container: new Container(),
+      relay: { enabled: false },
+    })
+    workers.push(worker)
+    await worker.start()
+    started = true
+  }
+  // Mirror the actual adapter scope into the app's independently bound scope port.
+  const base = extra.container ?? container
+  const child = base.child()
+  child.value(AGENT_SCOPE, { ...base.resolve(AGENT_SCOPE), actor: systemActor(RUN_JOB) })
+  const delivery = newId()
+  const settled = new Promise<void>((resolve, reject) =>
+    pending.set(delivery, { container: child, extra, ...(change ? { change } : {}), resolve, reject }),
+  )
+  await createPgBossJobQueue({ boss }).enqueue(
+    change === "owner" ? { ...scope, actor: { userId: newId() } } : scope,
+    RUN_JOB,
+    { runId, delivery },
+  )
+  return settled
+}