template-app/modules/agent/__tests__/claimed-job.ts

// 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
}