@shipfox/api-runners 5.0.0 → 6.0.0

package/src/db/capacity-assignments.test.ts

@@ -0,0 +1,159 @@
+import {eq, inArray} from 'drizzle-orm';
+import {assignCapacityBatch} from '#db/capacity-assignments.js';
+import {db} from '#db/db.js';
+import {capacityAssignments} from '#db/schema/capacity-assignments.js';
+import {reservations} from '#db/schema/reservations.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+
+describe('assignCapacityBatch', () => {
+  let workspaceId: string;
+  let provisionerId: string;
+
+  beforeEach(() => {
+    workspaceId = crypto.randomUUID();
+    provisionerId = crypto.randomUUID();
+  });
+
+  it('assigns compatible live capacity once and returns the same assignment on retry', async () => {
+    const reservation = await createReservation();
+    const capacity = await createCapacity();
+
+    const first = await assignCapacityBatch({
+      provisionerId,
+      reservationId: reservation.id,
+      capacityIds: [capacity.id],
+    });
+    const retry = await assignCapacityBatch({
+      provisionerId,
+      reservationId: reservation.id,
+      capacityIds: [capacity.id],
+    });
+
+    expect(first).toHaveLength(1);
+    expect(retry).toEqual(first);
+    expect(first[0]).toMatchObject({
+      capacityId: capacity.id,
+      reservationId: reservation.id,
+      workspaceId,
+      provisionerId,
+    });
+  });
+
+  it('does not rebind capacity or a reservation under concurrent assignment requests', async () => {
+    const reservation = await createReservation();
+    const firstCapacity = await createCapacity();
+    const secondCapacity = await createCapacity();
+
+    const result = await Promise.allSettled([
+      assignCapacityBatch({
+        provisionerId,
+        reservationId: reservation.id,
+        capacityIds: [firstCapacity.id],
+      }),
+      assignCapacityBatch({
+        provisionerId,
+        reservationId: reservation.id,
+        capacityIds: [secondCapacity.id],
+      }),
+    ]);
+
+    expect(result.filter((entry) => entry.status === 'fulfilled')).toHaveLength(1);
+    const rows = await db()
+      .select()
+      .from(capacityAssignments)
+      .where(eq(capacityAssignments.reservationId, reservation.id));
+    expect(rows).toHaveLength(1);
+  });
+
+  it('consumes each slot in a multi-capacity reservation without rebinding capacity', async () => {
+    const reservation = await createReservation({count: 2});
+    const firstCapacity = await createCapacity();
+    const secondCapacity = await createCapacity();
+
+    const assignments = await assignCapacityBatch({
+      provisionerId,
+      reservationId: reservation.id,
+      capacityIds: [firstCapacity.id, secondCapacity.id],
+    });
+    const retry = await assignCapacityBatch({
+      provisionerId,
+      reservationId: reservation.id,
+      capacityIds: [firstCapacity.id, secondCapacity.id],
+    });
+
+    expect(assignments).toHaveLength(2);
+    expect(retry).toEqual(assignments);
+  });
+
+  it('rejects expired, incompatible, and foreign capacity without consuming it', async () => {
+    const expiredReservation = await createReservation({expiresAt: new Date(Date.now() - 1_000)});
+    const compatibleCapacity = await createCapacity();
+    const incompatibleReservation = await createReservation({requiredLabels: ['gpu']});
+    const foreignCapacity = await createCapacity({provisionerId: crypto.randomUUID()});
+
+    await expect(
+      assignCapacityBatch({
+        provisionerId,
+        reservationId: expiredReservation.id,
+        capacityIds: [compatibleCapacity.id],
+      }),
+    ).rejects.toMatchObject({name: 'ReservationExpiredError'});
+    await expect(
+      assignCapacityBatch({
+        provisionerId,
+        reservationId: incompatibleReservation.id,
+        capacityIds: [compatibleCapacity.id],
+      }),
+    ).rejects.toMatchObject({name: 'CapacityNotAssignableError'});
+    await expect(
+      assignCapacityBatch({
+        provisionerId,
+        reservationId: incompatibleReservation.id,
+        capacityIds: [foreignCapacity.id],
+      }),
+    ).rejects.toMatchObject({name: 'CapacityNotAssignableError'});
+
+    const rows = await db()
+      .select()
+      .from(capacityAssignments)
+      .where(
+        inArray(capacityAssignments.reservationId, [
+          expiredReservation.id,
+          incompatibleReservation.id,
+        ]),
+      );
+    expect(rows).toEqual([]);
+  });
+
+  async function createReservation(overrides: Partial<typeof reservations.$inferInsert> = {}) {
+    const [row] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+        ...overrides,
+      })
+      .returning();
+    if (!row) throw new Error('Reservation insert returned no row');
+    return row;
+  }
+
+  async function createCapacity(overrides: Partial<typeof providerRunners.$inferInsert> = {}) {
+    const [row] = await db()
+      .insert(providerRunners)
+      .values({
+        provisionerId,
+        providerRunnerId: crypto.randomUUID(),
+        labels: ['linux'],
+        state: 'running',
+        reportedAt: new Date(),
+        ...overrides,
+      })
+      .returning();
+    if (!row) throw new Error('Capacity insert returned no row');
+    return row;
+  }
+});

package/src/db/capacity-assignments.ts

@@ -0,0 +1,116 @@
+import {and, eq, inArray, sql} from 'drizzle-orm';
+import type {CapacityAssignment} from '#core/entities/capacity-assignment.js';
+import {
+  CapacityAlreadyAssignedError,
+  CapacityNotAssignableError,
+  ReservationExpiredError,
+  ReservationNotFoundError,
+} from '#core/errors.js';
+import {db} from './db.js';
+import {capacityAssignments, toCapacityAssignment} from './schema/capacity-assignments.js';
+import {reservations} from './schema/reservations.js';
+import {providerRunners} from './schema/runner-instances.js';
+
+export async function assignCapacityBatch(params: {
+  provisionerId: string;
+  reservationId: string;
+  capacityIds: string[];
+}): Promise<CapacityAssignment[]> {
+  return await db().transaction(async (tx) => {
+    const lockKey = [
+      'runners_capacity_assignments',
+      params.provisionerId,
+      params.reservationId,
+    ].join(':');
+    await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${lockKey}))`);
+    for (const capacityId of [...params.capacityIds].sort()) {
+      await tx.execute(
+        sql`select pg_advisory_xact_lock(hashtext(${['runners_capacity_assignments', capacityId].join(':')}))`,
+      );
+    }
+
+    const [reservation] = await tx
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.id, params.reservationId),
+          eq(reservations.provisionerId, params.provisionerId),
+        ),
+      )
+      .limit(1);
+    if (!reservation) throw new ReservationNotFoundError(params.reservationId);
+    if (reservation.expiresAt <= new Date())
+      throw new ReservationExpiredError(params.reservationId);
+
+    const existing = await tx
+      .select()
+      .from(capacityAssignments)
+      .where(eq(capacityAssignments.reservationId, reservation.id));
+    const existingByCapacityId = new Map(
+      existing.map((assignment) => [assignment.capacityId, assignment]),
+    );
+    const newCapacityIds = params.capacityIds.filter(
+      (capacityId) => !existingByCapacityId.has(capacityId),
+    );
+    if (existing.length + newCapacityIds.length > reservation.count)
+      throw new CapacityNotAssignableError(params.capacityIds[0] ?? '');
+
+    if (newCapacityIds.length === 0) {
+      return params.capacityIds.flatMap((capacityId) => {
+        const assignment = existingByCapacityId.get(capacityId);
+        return assignment ? [toCapacityAssignment(assignment)] : [];
+      });
+    }
+
+    const capacities = await tx
+      .select()
+      .from(providerRunners)
+      .where(
+        and(
+          eq(providerRunners.provisionerId, params.provisionerId),
+          inArray(providerRunners.id, newCapacityIds),
+        ),
+      );
+    if (capacities.length !== newCapacityIds.length)
+      throw new CapacityNotAssignableError(params.capacityIds[0] ?? '');
+    for (const capacity of capacities) {
+      if (capacity.state !== 'starting' && capacity.state !== 'running')
+        throw new CapacityNotAssignableError(capacity.id);
+      if (
+        !capacity.providerRunnerId ||
+        capacity.labels.length === 0 ||
+        !reservation.requiredLabels.every((label) => capacity.labels.includes(label))
+      ) {
+        throw new CapacityNotAssignableError(capacity.id);
+      }
+    }
+
+    const assignedCapacity = await tx
+      .select({capacityId: capacityAssignments.capacityId})
+      .from(capacityAssignments)
+      .where(inArray(capacityAssignments.capacityId, newCapacityIds));
+    if (assignedCapacity.length > 0)
+      throw new CapacityAlreadyAssignedError(assignedCapacity[0]?.capacityId ?? '');
+
+    const inserted = await tx
+      .insert(capacityAssignments)
+      .values(
+        newCapacityIds.map((capacityId) => ({
+          capacityId,
+          reservationId: reservation.id,
+          workspaceId: reservation.workspaceId,
+          provisionerId: params.provisionerId,
+        })),
+      )
+      .returning();
+    const insertedByCapacityId = new Map(
+      inserted.map((assignment) => [assignment.capacityId, assignment]),
+    );
+    return params.capacityIds.flatMap((capacityId) => {
+      const assignment =
+        existingByCapacityId.get(capacityId) ?? insertedByCapacityId.get(capacityId);
+      return assignment ? [toCapacityAssignment(assignment)] : [];
+    });
+  });
+}

package/src/db/db.ts

@@ -1,24 +1,33 @@
 import {drizzle, type NodePgDatabase} from '@shipfox/node-drizzle';
 import {pgClient} from '@shipfox/node-postgres';
+import {capacityAssignments} from './schema/capacity-assignments.js';
 import {ephemeralRegistrationTokens} from './schema/ephemeral-registration-tokens.js';
 import {manualRegistrationTokens} from './schema/manual-registration-tokens.js';
 import {runnersOutbox} from './schema/outbox.js';
 import {pendingJobExecutions} from './schema/pending-job-executions.js';
-import {provisionedRunners} from './schema/provisioned-runners.js';
+import {provisionerCapabilitySnapshots} from './schema/provisioner-capability-snapshots.js';
 import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {runnersRateLimits} from './schema/rate-limits.js';
 import {reservations} from './schema/reservations.js';
+import {runnerActivationTokens} from './schema/runner-activation-tokens.js';
+import {runnerBootstrapTokens, runnerControlSessions} from './schema/runner-control-sessions.js';
+import {providerRunners} from './schema/runner-instances.js';
 import {runnerSessions} from './schema/runner-sessions.js';
 import {runningJobExecutions} from './schema/running-job-executions.js';
 
 export const schema = {
+  capacityAssignments,
   ephemeralRegistrationTokens,
   pendingJobExecutions,
-  provisionedRunners,
+  providerRunners,
+  provisionerCapabilitySnapshots,
   provisionerTokens,
   runnersRateLimits,
   reservations,
   runnerSessions,
+  runnerBootstrapTokens,
+  runnerControlSessions,
+  runnerActivationTokens,
   manualRegistrationTokens,
   runningJobExecutions,
   runnersOutbox,

package/src/db/ephemeral-registration-tokens.test.ts

@@ -44,7 +44,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     });
 
     expect(result).toHaveLength(2);
-    expect(result.map((token) => token.provisionedRunnerId).sort()).toEqual([
+    expect(result.map((token) => token.providerRunnerId).sort()).toEqual([
       'provisioned-runner-a',
       'provisioned-runner-b',
     ]);
@@ -53,14 +53,14 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     expect(result.every((token) => token.provisionerId === provisionerId)).toBe(true);
     expect(
       await resolveEphemeralRegistrationTokenByHash(hashOpaqueToken(rawTokens[0] ?? '')),
-    ).toEqual(expect.objectContaining({provisionedRunnerId: 'provisioned-runner-a'}));
+    ).toEqual(expect.objectContaining({providerRunnerId: 'provisioned-runner-a'}));
   });
 
   it('rejects the batch and inserts no rows when any provisioned runner already has an active token', async () => {
     await ephemeralRegistrationTokenFactory.create({
       workspaceId,
       provisionerId,
-      provisionedRunnerId: 'provisioned-runner-a',
+      providerRunnerId: 'provisioned-runner-a',
       expiresAt: new Date(Date.now() + 60_000),
     });
     const rawToken = generateOpaqueToken('ephemeralRegistrationToken');
@@ -85,7 +85,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
         ),
       );
     expect(rows).toHaveLength(1);
-    expect(rows[0]?.provisionedRunnerId).toBe('provisioned-runner-a');
+    expect(rows[0]?.providerRunnerId).toBe('provisioned-runner-a');
   });
 
   it('allows sequential batches that exactly reach the reservation count', async () => {
@@ -152,7 +152,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
       workspaceId,
       provisionerId,
       reservationId,
-      provisionedRunnerId: 'provisioned-runner-consumed',
+      providerRunnerId: 'provisioned-runner-consumed',
       expiresAt: new Date(Date.now() + 60_000),
     });
     await db()
@@ -163,7 +163,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
       workspaceId,
       provisionerId,
       reservationId,
-      provisionedRunnerId: 'provisioned-runner-expired',
+      providerRunnerId: 'provisioned-runner-expired',
       expiresAt: new Date(Date.now() - 60_000),
     });
 
@@ -209,9 +209,9 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     return rows.length;
   }
 
-  function row(provisionedRunnerId: string, rawToken: string) {
+  function row(providerRunnerId: string, rawToken: string) {
     return {
-      provisionedRunnerId,
+      providerRunnerId,
       hashedToken: hashOpaqueToken(rawToken),
       prefix: extractDisplayPrefix(rawToken),
     };
@@ -303,7 +303,7 @@ describe('deleteExpiredEphemeralRegistrationTokens', () => {
         id,
         workspaceId,
         provisionerId,
-        provisionedRunnerId: `provisioned-${id}`,
+        providerRunnerId: `provisioned-${id}`,
         hashedToken: crypto.randomUUID(),
         prefix: 'sfxr_test',
         expiresAt: params.expiresAt,

package/src/db/ephemeral-registration-tokens.ts

@@ -16,22 +16,22 @@ import {
   ephemeralRegistrationTokens,
   toEphemeralRegistrationToken,
 } from './schema/ephemeral-registration-tokens.js';
-import {provisionedRunners} from './schema/provisioned-runners.js';
 import {reservations} from './schema/reservations.js';
+import {providerRunners} from './schema/runner-instances.js';
 import {runnerSessions, toRunnerSession} from './schema/runner-sessions.js';
 
 export interface CreateEphemeralRegistrationTokenParams {
   workspaceId: string;
   provisionerId: string;
   reservationId?: string | null | undefined;
-  provisionedRunnerId: string;
+  providerRunnerId: string;
   hashedToken: string;
   prefix: string;
   expiresAt: Date;
 }
 
 export interface CreateEphemeralRegistrationTokensBatchRow {
-  provisionedRunnerId: string;
+  providerRunnerId: string;
   hashedToken: string;
   prefix: string;
 }
@@ -48,13 +48,13 @@ export async function createEphemeralRegistrationToken(
   params: CreateEphemeralRegistrationTokenParams,
 ): Promise<EphemeralRegistrationToken> {
   const rows = await db().transaction(async (tx) => {
-    const provisionedRunnerLockKey = [
+    const providerRunnerLockKey = [
       'runners_ephemeral_registration_tokens',
       params.workspaceId,
       params.provisionerId,
-      params.provisionedRunnerId,
+      params.providerRunnerId,
     ].join(':');
-    await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${provisionedRunnerLockKey}))`);
+    await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${providerRunnerLockKey}))`);
 
     const [active] = await tx
       .select({id: ephemeralRegistrationTokens.id})
@@ -63,7 +63,7 @@ export async function createEphemeralRegistrationToken(
         and(
           eq(ephemeralRegistrationTokens.workspaceId, params.workspaceId),
           eq(ephemeralRegistrationTokens.provisionerId, params.provisionerId),
-          eq(ephemeralRegistrationTokens.provisionedRunnerId, params.provisionedRunnerId),
+          eq(ephemeralRegistrationTokens.providerRunnerId, params.providerRunnerId),
           isNull(ephemeralRegistrationTokens.consumedAt),
           gt(ephemeralRegistrationTokens.expiresAt, sql`now()`),
         ),
@@ -74,7 +74,7 @@ export async function createEphemeralRegistrationToken(
       throw new ActiveEphemeralRegistrationTokenExistsError(
         params.workspaceId,
         params.provisionerId,
-        params.provisionedRunnerId,
+        params.providerRunnerId,
       );
     }
 
@@ -84,7 +84,7 @@ export async function createEphemeralRegistrationToken(
         workspaceId: params.workspaceId,
         provisionerId: params.provisionerId,
         reservationId: params.reservationId ?? null,
-        provisionedRunnerId: params.provisionedRunnerId,
+        providerRunnerId: params.providerRunnerId,
         hashedToken: params.hashedToken,
         prefix: params.prefix,
         expiresAt: params.expiresAt,
@@ -107,16 +107,16 @@ export async function createEphemeralRegistrationTokensBatch(
       params.provisionerId,
     ].join(':');
     await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${provisionerLockKey}))`);
-    for (const provisionedRunnerId of [
-      ...new Set(params.rows.map((row) => row.provisionedRunnerId)),
+    for (const providerRunnerId of [
+      ...new Set(params.rows.map((row) => row.providerRunnerId)),
     ].sort()) {
-      const provisionedRunnerLockKey = [
+      const providerRunnerLockKey = [
         'runners_ephemeral_registration_tokens',
         params.workspaceId,
         params.provisionerId,
-        provisionedRunnerId,
+        providerRunnerId,
       ].join(':');
-      await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${provisionedRunnerLockKey}))`);
+      await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${providerRunnerLockKey}))`);
     }
 
     const [reservation] = await tx
@@ -162,15 +162,15 @@ export async function createEphemeralRegistrationTokensBatch(
       );
     }
 
-    const provisionedRunnerIds = params.rows.map((row) => row.provisionedRunnerId);
+    const providerRunnerIds = params.rows.map((row) => row.providerRunnerId);
     const activeRows = await tx
-      .select({provisionedRunnerId: ephemeralRegistrationTokens.provisionedRunnerId})
+      .select({providerRunnerId: ephemeralRegistrationTokens.providerRunnerId})
       .from(ephemeralRegistrationTokens)
       .where(
         and(
           eq(ephemeralRegistrationTokens.workspaceId, params.workspaceId),
           eq(ephemeralRegistrationTokens.provisionerId, params.provisionerId),
-          inArray(ephemeralRegistrationTokens.provisionedRunnerId, provisionedRunnerIds),
+          inArray(ephemeralRegistrationTokens.providerRunnerId, providerRunnerIds),
           isNull(ephemeralRegistrationTokens.consumedAt),
           gt(ephemeralRegistrationTokens.expiresAt, sql`now()`),
         ),
@@ -178,7 +178,7 @@ export async function createEphemeralRegistrationTokensBatch(
 
     if (activeRows.length > 0) {
       throw new ActiveEphemeralRegistrationTokensExistError(
-        activeRows.map((row) => row.provisionedRunnerId),
+        activeRows.map((row) => row.providerRunnerId),
       );
     }
 
@@ -189,7 +189,7 @@ export async function createEphemeralRegistrationTokensBatch(
           workspaceId: params.workspaceId,
           provisionerId: params.provisionerId,
           reservationId: params.reservationId,
-          provisionedRunnerId: row.provisionedRunnerId,
+          providerRunnerId: row.providerRunnerId,
           hashedToken: row.hashedToken,
           prefix: row.prefix,
           expiresAt: params.expiresAt,
@@ -277,7 +277,7 @@ export async function createRunnerSessionConsumingEphemeralToken(params: {
       .returning({
         id: ephemeralRegistrationTokens.id,
         provisionerId: ephemeralRegistrationTokens.provisionerId,
-        provisionedRunnerId: ephemeralRegistrationTokens.provisionedRunnerId,
+        providerRunnerId: ephemeralRegistrationTokens.providerRunnerId,
       });
 
     if (!consumed[0]) {
@@ -315,7 +315,7 @@ export async function createRunnerSessionConsumingEphemeralToken(params: {
         registrationTokenId: params.ephemeralTokenId,
         registrationTokenKind: 'ephemeral',
         provisionerId: consumed[0].provisionerId,
-        provisionedRunnerId: consumed[0].provisionedRunnerId,
+        providerRunnerId: consumed[0].providerRunnerId,
         labels: params.labels,
         toolCapabilities: params.toolCapabilities ?? null,
         toolCapabilitiesReportedAt: params.toolCapabilities ? sql`now()` : null,
@@ -332,14 +332,14 @@ export async function createRunnerSessionConsumingEphemeralToken(params: {
       .where(eq(ephemeralRegistrationTokens.id, params.ephemeralTokenId));
 
     await tx
-      .update(provisionedRunners)
+      .update(providerRunners)
       .set({runnerSessionId: session.id, updatedAt: sql`now()`})
       .where(
         and(
-          eq(provisionedRunners.workspaceId, params.workspaceId),
-          eq(provisionedRunners.provisionerId, consumed[0].provisionerId),
-          eq(provisionedRunners.provisionedRunnerId, consumed[0].provisionedRunnerId),
-          isNull(provisionedRunners.reservationReleasedAt),
+          eq(providerRunners.workspaceId, params.workspaceId),
+          eq(providerRunners.provisionerId, consumed[0].provisionerId),
+          eq(providerRunners.providerRunnerId, consumed[0].providerRunnerId),
+          isNull(providerRunners.reservationReleasedAt),
         ),
       );
 

package/src/db/index.ts

@@ -1,6 +1,7 @@
 import {dirname, resolve} from 'node:path';
 import {fileURLToPath} from 'node:url';
 
+export {assignCapacityBatch} from './capacity-assignments.js';
 export {closeDb, db, schema} from './db.js';
 export type {
   CreateEphemeralRegistrationTokenParams,
@@ -15,7 +16,7 @@ export {
 export type {
   ClaimedJobExecution,
   EnqueueJobExecutionParams,
-  ProvisionedRunnerBoundJobExecution,
+  RunnerInstanceBoundJobExecution,
 } from './job-executions.js';
 export {
   cancelRunnerJobs,
@@ -24,7 +25,7 @@ export {
   expireStuckJobExecutions,
   isJobLeaseActive,
   listActiveRunningJobExecutions,
-  listRunningJobExecutionsByProvisionedRunnerTx,
+  listRunningJobExecutionsByRunnerInstanceTx,
   recordHeartbeat,
   releaseJobExecution,
 } from './job-executions.js';
@@ -35,26 +36,12 @@ export {
   resolveManualRegistrationTokenByHash,
   revokeManualRegistrationToken,
 } from './manual-registration-tokens.js';
-export type {
-  ActiveProvisionedRunnerTemplateCount,
-  ProvisionedRunnerReportEvent,
-  ProvisionedRunnerTerminateIntent,
-  ProvisionedRunnerTerminateIntentReason,
-  ReapStaleProvisionedRunnersResult,
-  ReconcileProvisionedRunnersDbResult,
-  ReconcileProvisionedRunnersParams,
-  ReportProvisionedRunnersParams,
-} from './provisioned-runners.js';
 export {
-  isTerminalState,
-  listActiveProvisionedRunnerCountsByTemplateTx,
-  listActiveProvisionedRunners,
-  listProvisionerTerminateIntentRowsTx,
-  listProvisionerTerminateIntents,
-  reapStaleProvisionedRunners,
-  reconcileProvisionedRunners,
-  reportProvisionedRunners,
-} from './provisioned-runners.js';
+  hasActiveWorkspaceProvisionerCapability,
+  listActiveWorkspaceProvisionerCapabilitySnapshots,
+  listStaleWorkspaceProvisionerCapabilitySnapshots,
+  publishWorkspaceProvisionerCapabilitySnapshot,
+} from './provisioner-capability-snapshots.js';
 export type {CreateProvisionerTokenParams} from './provisioner-tokens.js';
 export {
   createProvisionerToken,
@@ -80,11 +67,34 @@ export {
   pollDemandAndReserve,
   releaseReservationUnits,
 } from './reservations.js';
+export type {
+  ActiveRunnerInstanceTemplateCount,
+  ReapStaleRunnerInstancesResult,
+  ReconcileRunnerInstancesDbResult,
+  ReconcileRunnerInstancesParams,
+  ReportRunnerInstancesParams,
+  RunnerInstanceReportEvent,
+  RunnerInstanceTerminateIntent,
+  RunnerInstanceTerminateIntentReason,
+} from './runner-instances.js';
+export {
+  attachProviderRunnerId,
+  createPlannedProvisionedCapacity,
+  isTerminalState,
+  listActiveRunnerInstanceCountsByTemplateTx,
+  listActiveRunnerInstances,
+  listProvisionerTerminateIntentRowsTx,
+  listProvisionerTerminateIntents,
+  reapStaleRunnerInstances,
+  reconcileRunnerInstances,
+  reportRunnerInstances,
+} from './runner-instances.js';
 export type {
   CreateRunnerSessionParams,
   DeleteExpiredRunnerSessionsParams,
 } from './runner-sessions.js';
 export {createRunnerSession, deleteExpiredRunnerSessions} from './runner-sessions.js';
 export {runnersOutbox} from './schema/outbox.js';
+export {runnerBootstrapTokens, runnerControlSessions} from './schema/runner-control-sessions.js';
 
 export const migrationsPath = resolve(dirname(fileURLToPath(import.meta.url)), '../../drizzle');