@shipfox/api-runners 5.0.0 → 6.0.0

package/src/db/provisioner-tokens.test.ts

@@ -1,5 +1,5 @@
 import {hashOpaqueToken} from '@shipfox/node-tokens';
-import {sql} from 'drizzle-orm';
+import {eq, sql} from 'drizzle-orm';
 import {db} from '#db/db.js';
 import {
   createProvisionerToken,
@@ -9,7 +9,12 @@ import {
   revokeProvisionerToken,
   touchProvisionerLastSeen,
 } from '#db/provisioner-tokens.js';
-import {provisionerTokenFactory} from '#test/index.js';
+import {provisionerTokens} from '#db/schema/provisioner-tokens.js';
+import {runnerActivationTokens} from '#db/schema/runner-activation-tokens.js';
+import {runnerBootstrapTokens, runnerControlSessions} from '#db/schema/runner-control-sessions.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+import {runnerSessions} from '#db/schema/runner-sessions.js';
+import {providerRunnerFactory, provisionerTokenFactory} from '#test/index.js';
 
 describe('provisioner token db', () => {
   it('creates and resolves a token', async () => {
@@ -17,6 +22,7 @@ describe('provisioner token db', () => {
     const rawToken = 'sf_pt_test-token';
     const createdByUserId = crypto.randomUUID();
     const token = await createProvisionerToken({
+      scope: 'workspace',
       workspaceId,
       hashedToken: hashOpaqueToken(rawToken),
       prefix: rawToken.slice(0, 12),
@@ -37,6 +43,28 @@ describe('provisioner token db', () => {
     expect(result).toBeUndefined();
   });
 
+  it('creates installation factory tokens without a workspace', async () => {
+    const token = await provisionerTokenFactory.create({scope: 'installation'});
+
+    expect(token).toMatchObject({scope: 'installation', workspaceId: null});
+  });
+
+  it('enforces valid scope and workspace combinations in PostgreSQL', async () => {
+    const insert = (scope: 'installation' | 'workspace', workspaceId: string | null) =>
+      db()
+        .insert(provisionerTokens)
+        .values({
+          scope,
+          workspaceId,
+          hashedToken: hashOpaqueToken(`${scope}-${workspaceId}`),
+          prefix: 'sf_pt_invalid',
+          createdByUserId: crypto.randomUUID(),
+        });
+
+    await expect(insert('installation', crypto.randomUUID())).rejects.toThrow();
+    await expect(insert('workspace', null)).rejects.toThrow();
+  });
+
   it('lists usable tokens for a workspace', async () => {
     const workspaceId = crypto.randomUUID();
     const usable = await provisionerTokenFactory.create({workspaceId, name: 'usable'});
@@ -74,6 +102,115 @@ describe('provisioner token db', () => {
     expect(revoked?.revokedByUserId).toBe(revokedByUserId);
   });
 
+  it("cascades revocation to the provisioner's unclaimed runner credentials and sessions", async () => {
+    const workspaceId = crypto.randomUUID();
+    const provisioner = await provisionerTokenFactory.create({workspaceId});
+    const runner = await providerRunnerFactory.create({
+      workspaceId,
+      provisionerId: provisioner.id,
+      runnerSessionId: null,
+    });
+    const future = new Date(Date.now() + 60_000);
+    const tokenValue = (name: string) => hashOpaqueToken(`${name}-${crypto.randomUUID()}`);
+    await db()
+      .insert(runnerBootstrapTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        provisionerId: provisioner.id,
+        hashedToken: tokenValue('bootstrap'),
+        prefix: 'bootstrap',
+        expiresAt: future,
+      });
+    await db()
+      .insert(runnerControlSessions)
+      .values({
+        runnerInstanceId: runner.id,
+        provisionerId: provisioner.id,
+        hashedToken: tokenValue('control'),
+        prefix: 'control',
+        expiresAt: future,
+      });
+    await db()
+      .insert(runnerActivationTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        hashedToken: tokenValue('activation'),
+        prefix: 'activation',
+        expiresAt: future,
+      });
+    const [unclaimedSession] = await db()
+      .insert(runnerSessions)
+      .values({
+        workspaceId,
+        scope: 'workspace',
+        registrationTokenId: crypto.randomUUID(),
+        registrationTokenKind: 'activation',
+        runnerInstanceId: runner.id,
+        provisionerId: provisioner.id,
+        providerRunnerId: runner.providerRunnerId,
+        labels: ['linux'],
+        maxClaims: 1,
+        claimsUsed: 0,
+      })
+      .returning();
+    const [claimedSession] = await db()
+      .insert(runnerSessions)
+      .values({
+        workspaceId,
+        scope: 'workspace',
+        registrationTokenId: crypto.randomUUID(),
+        registrationTokenKind: 'activation',
+        runnerInstanceId: crypto.randomUUID(),
+        provisionerId: provisioner.id,
+        providerRunnerId: crypto.randomUUID(),
+        labels: ['linux'],
+        maxClaims: 1,
+        claimsUsed: 1,
+      })
+      .returning();
+    if (!unclaimedSession || !claimedSession) throw new Error('Runner sessions were not created');
+
+    await revokeProvisionerToken({
+      tokenId: provisioner.id,
+      workspaceId,
+      revokedByUserId: crypto.randomUUID(),
+    });
+
+    const [bootstrap] = await db()
+      .select()
+      .from(runnerBootstrapTokens)
+      .where(eq(runnerBootstrapTokens.runnerInstanceId, runner.id));
+    const [control] = await db()
+      .select()
+      .from(runnerControlSessions)
+      .where(eq(runnerControlSessions.runnerInstanceId, runner.id));
+    const [activation] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.runnerInstanceId, runner.id));
+    const [terminatedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const [revokedSession] = await db()
+      .select()
+      .from(runnerSessions)
+      .where(eq(runnerSessions.id, unclaimedSession.id));
+    const [preservedSession] = await db()
+      .select()
+      .from(runnerSessions)
+      .where(eq(runnerSessions.id, claimedSession.id));
+
+    expect(bootstrap?.revokedAt).toBeInstanceOf(Date);
+    expect(control).toMatchObject({closeReason: 'provisioner-revoked'});
+    expect(control?.closedAt).toBeInstanceOf(Date);
+    expect(activation?.revokedAt).toBeInstanceOf(Date);
+    expect(terminatedRunner).toMatchObject({state: 'terminated'});
+    expect(terminatedRunner?.terminatedAt).toBeInstanceOf(Date);
+    expect(revokedSession?.revokedAt).toBeInstanceOf(Date);
+    expect(preservedSession?.revokedAt).toBeNull();
+  });
+
   it('preserves the original revocation audit fields on repeat revoke', async () => {
     const workspaceId = crypto.randomUUID();
     const token = await provisionerTokenFactory.create({workspaceId});

package/src/db/provisioner-tokens.ts

@@ -1,10 +1,19 @@
-import {and, desc, eq, gt, isNull, or, sql} from 'drizzle-orm';
-import type {ActiveProvisionerToken, ProvisionerToken} from '#core/entities/provisioner-token.js';
+import {and, desc, eq, gt, inArray, isNull, or, sql} from 'drizzle-orm';
+import type {
+  ActiveProvisionerToken,
+  ProvisionerScope,
+  ProvisionerToken,
+} from '#core/entities/provisioner-token.js';
 import {db} from './db.js';
 import {provisionerTokens, toProvisionerToken} from './schema/provisioner-tokens.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';
 
 export interface CreateProvisionerTokenParams {
-  workspaceId: string;
+  scope: ProvisionerScope;
+  workspaceId?: string | undefined;
   hashedToken: string;
   prefix: string;
   createdByUserId: string;
@@ -12,13 +21,68 @@ export interface CreateProvisionerTokenParams {
   expiresAt?: Date | undefined;
 }
 
+async function cascadeProvisionerRevocation(
+  tx: Parameters<Parameters<ReturnType<typeof db>['transaction']>[0]>[0],
+  provisionerId: string,
+) {
+  await tx
+    .update(runnerBootstrapTokens)
+    .set({revokedAt: sql`now()`})
+    .where(eq(runnerBootstrapTokens.provisionerId, provisionerId));
+  await tx
+    .update(runnerControlSessions)
+    .set({closedAt: sql`now()`, closeReason: 'provisioner-revoked'})
+    .where(
+      and(
+        eq(runnerControlSessions.provisionerId, provisionerId),
+        isNull(runnerControlSessions.closedAt),
+      ),
+    );
+  await tx
+    .update(runnerActivationTokens)
+    .set({revokedAt: sql`now()`})
+    .where(
+      and(
+        isNull(runnerActivationTokens.consumedAt),
+        isNull(runnerActivationTokens.revokedAt),
+        inArray(
+          runnerActivationTokens.runnerInstanceId,
+          tx
+            .select({id: providerRunners.id})
+            .from(providerRunners)
+            .where(eq(providerRunners.provisionerId, provisionerId)),
+        ),
+      ),
+    );
+  await tx
+    .update(providerRunners)
+    .set({state: 'terminated', terminatedAt: sql`now()`, updatedAt: sql`now()`})
+    .where(
+      and(
+        eq(providerRunners.provisionerId, provisionerId),
+        isNull(providerRunners.runnerSessionId),
+      ),
+    );
+  await tx
+    .update(runnerSessions)
+    .set({revokedAt: sql`now()`})
+    .where(
+      and(
+        eq(runnerSessions.provisionerId, provisionerId),
+        eq(runnerSessions.claimsUsed, 0),
+        isNull(runnerSessions.revokedAt),
+      ),
+    );
+}
+
 export async function createProvisionerToken(
   params: CreateProvisionerTokenParams,
 ): Promise<ProvisionerToken> {
   const rows = await db()
     .insert(provisionerTokens)
     .values({
-      workspaceId: params.workspaceId,
+      scope: params.scope,
+      workspaceId: params.workspaceId ?? null,
       hashedToken: params.hashedToken,
       prefix: params.prefix,
       createdByUserId: params.createdByUserId,
@@ -32,6 +96,42 @@ export async function createProvisionerToken(
   return toProvisionerToken(row);
 }
 
+export async function revokeInstallationProvisionerToken(params: {
+  tokenId: string;
+  revokedByUserId: string;
+}): Promise<ProvisionerToken | undefined> {
+  const rows = await db().transaction(async (tx) => {
+    const rows = await tx
+      .update(provisionerTokens)
+      .set({revokedAt: new Date(), revokedByUserId: params.revokedByUserId, updatedAt: new Date()})
+      .where(
+        and(
+          eq(provisionerTokens.id, params.tokenId),
+          eq(provisionerTokens.scope, 'installation'),
+          isNull(provisionerTokens.revokedAt),
+        ),
+      )
+      .returning();
+    if (rows[0]) {
+      await cascadeProvisionerRevocation(tx, params.tokenId);
+    }
+    return rows;
+  });
+
+  const row = rows[0];
+  if (row) return toProvisionerToken(row);
+
+  const existingRows = await db()
+    .select()
+    .from(provisionerTokens)
+    .where(
+      and(eq(provisionerTokens.id, params.tokenId), eq(provisionerTokens.scope, 'installation')),
+    )
+    .limit(1);
+  const existingRow = existingRows[0];
+  return existingRow ? toProvisionerToken(existingRow) : undefined;
+}
+
 export async function listUsableProvisionerTokensByWorkspaceId(
   workspaceId: string,
 ): Promise<ProvisionerToken[]> {
@@ -42,6 +142,7 @@ export async function listUsableProvisionerTokensByWorkspaceId(
     .where(
       and(
         eq(provisionerTokens.workspaceId, workspaceId),
+        eq(provisionerTokens.scope, 'workspace'),
         isNull(provisionerTokens.revokedAt),
         or(isNull(provisionerTokens.expiresAt), gt(provisionerTokens.expiresAt, now)),
       ),
@@ -70,17 +171,22 @@ export async function revokeProvisionerToken(params: {
   workspaceId: string;
   revokedByUserId: string;
 }): Promise<ProvisionerToken | undefined> {
-  const rows = await db()
-    .update(provisionerTokens)
-    .set({revokedAt: new Date(), revokedByUserId: params.revokedByUserId, updatedAt: new Date()})
-    .where(
-      and(
-        eq(provisionerTokens.id, params.tokenId),
-        eq(provisionerTokens.workspaceId, params.workspaceId),
-        isNull(provisionerTokens.revokedAt),
-      ),
-    )
-    .returning();
+  const rows = await db().transaction(async (tx) => {
+    const rows = await tx
+      .update(provisionerTokens)
+      .set({revokedAt: new Date(), revokedByUserId: params.revokedByUserId, updatedAt: new Date()})
+      .where(
+        and(
+          eq(provisionerTokens.id, params.tokenId),
+          eq(provisionerTokens.workspaceId, params.workspaceId),
+          eq(provisionerTokens.scope, 'workspace'),
+          isNull(provisionerTokens.revokedAt),
+        ),
+      )
+      .returning();
+    if (rows[0]) await cascadeProvisionerRevocation(tx, params.tokenId);
+    return rows;
+  });
 
   const row = rows[0];
   if (row) return toProvisionerToken(row);
@@ -92,6 +198,7 @@ export async function revokeProvisionerToken(params: {
       and(
         eq(provisionerTokens.id, params.tokenId),
         eq(provisionerTokens.workspaceId, params.workspaceId),
+        eq(provisionerTokens.scope, 'workspace'),
       ),
     )
     .limit(1);
@@ -129,6 +236,7 @@ export async function listActiveProvisionerTokens(params: {
     .where(
       and(
         eq(provisionerTokens.workspaceId, params.workspaceId),
+        eq(provisionerTokens.scope, 'workspace'),
         isNull(provisionerTokens.revokedAt),
         or(isNull(provisionerTokens.expiresAt), gt(provisionerTokens.expiresAt, sql`now()`)),
         sql`${provisionerTokens.lastSeenAt} > now() - (${params.windowSeconds} || ' seconds')::interval`,

package/src/db/reservations.test.ts

@@ -3,6 +3,7 @@ import {db} from '#db/db.js';
 import {
   deleteReservationsByIds,
   pollDemandAndReserve,
+  pollInstallationDemandAndReserve,
   releaseReservationUnits,
 } from '#db/reservations.js';
 import {reservations} from '#db/schema/reservations.js';
@@ -75,6 +76,54 @@ describe('pollDemandAndReserve', () => {
     expect(reserved).toBeLessThanOrEqual(5);
   });
 
+  it('allocates eligible workspace heads in oldest-demand order without overselling templates', async () => {
+    const olderWorkspaceId = crypto.randomUUID();
+    const newerWorkspaceId = crypto.randomUUID();
+    await pendingJobFactory.create({workspaceId: olderWorkspaceId, requiredLabels: ['linux']});
+    await pendingJobFactory.create({workspaceId: newerWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 5,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([olderWorkspaceId, newerWorkspaceId]),
+    });
+
+    expect(result.reservations).toEqual([
+      expect.objectContaining({workspaceId: olderWorkspaceId, labels: ['linux'], count: 1}),
+    ]);
+  });
+
+  it('reports committed installation grants before an aborted poll stops allocating', async () => {
+    const firstWorkspaceId = crypto.randomUUID();
+    const secondWorkspaceId = crypto.randomUUID();
+    const abortController = new AbortController();
+    const reportedReservations: string[] = [];
+    await pendingJobFactory.create({workspaceId: firstWorkspaceId, requiredLabels: ['linux']});
+    await pendingJobFactory.create({workspaceId: secondWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 2)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([firstWorkspaceId, secondWorkspaceId]),
+      signal: abortController.signal,
+      onReservations: (reservations) => {
+        reportedReservations.push(...reservations.map((reservation) => reservation.reservationId));
+        abortController.abort();
+      },
+    });
+
+    expect(result.reservations).toHaveLength(1);
+    expect(reportedReservations).toEqual(
+      result.reservations.map((reservation) => reservation.reservationId),
+    );
+  });
+
   it('does not count expired reservations against demand', async () => {
     await createPendingJobs(1, ['linux']);
     await reservationFactory.create({

package/src/db/reservations.ts

@@ -3,6 +3,8 @@ import {and, asc, eq, gt, inArray, lt, sql} from 'drizzle-orm';
 import type {Tx} from './db.js';
 import {db} from './db.js';
 import {pendingJobExecutions} from './schema/pending-job-executions.js';
+import {provisionerCapabilitySnapshots} from './schema/provisioner-capability-snapshots.js';
+import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {reservations} from './schema/reservations.js';
 
 export interface ReservationTemplate {
@@ -14,6 +16,7 @@ export interface ReservationTemplate {
 }
 
 export interface DemandStat {
+  workspaceId?: string;
   labels: string[];
   queued: number;
   reserved: number;
@@ -22,6 +25,7 @@ export interface DemandStat {
 
 export interface ReservationGrant {
   reservationId: string;
+  workspaceId?: string;
   labels: string[];
   count: number;
   expiresAt: Date;
@@ -33,6 +37,18 @@ export interface PollDemandAndReserveParams {
   maxReservations: number;
   ttlSeconds: number;
   templates: ReservationTemplate[];
+  capabilityWindowSeconds?: number;
+}
+
+export interface InstallationPollDemandAndReserveParams {
+  provisionerId: string;
+  maxReservations: number;
+  ttlSeconds: number;
+  templates: ReservationTemplate[];
+  capabilityWindowSeconds: number;
+  eligibleWorkspaceIds: ReadonlySet<string>;
+  signal?: AbortSignal;
+  onReservations?: (reservations: ReservationGrant[]) => void;
 }
 
 interface NormalizedTemplate {
@@ -60,8 +76,75 @@ export async function pollDemandAndReserveTx(
   params: PollDemandAndReserveParams,
 ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
   await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${params.workspaceId}))`);
+  return await pollDemandAndReserveLockedTx(tx, params);
+}
+
+export async function pollInstallationDemandAndReserve(
+  params: InstallationPollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  const candidateWorkspaceIds = await listInstallationDemandWorkspaceIds(
+    params.eligibleWorkspaceIds,
+  );
+  const results: Array<{stats: DemandStat[]; reservations: ReservationGrant[]}> = [];
+  let remainingMaxReservations = params.maxReservations;
+  const remainingTemplates = params.templates.map((template) => ({
+    ...template,
+    labels: [...canonicalizeLabels(template.labels)],
+    remainingSlots: template.availableSlots,
+  }));
+  for (const workspaceId of candidateWorkspaceIds) {
+    if (params.signal?.aborted || remainingMaxReservations === 0) break;
+    const result = await db().transaction(async (tx) => {
+      const lockResult = await tx.execute<{locked: boolean}>(
+        sql`select pg_try_advisory_xact_lock(hashtext(${workspaceId})) as locked`,
+      );
+      const locked = lockResult.rows[0];
+      if (!locked?.locked) return {stats: [], reservations: []};
+      return await pollDemandAndReserveLockedTx(tx, {
+        workspaceId,
+        provisionerId: params.provisionerId,
+        maxReservations: remainingMaxReservations,
+        ttlSeconds: params.ttlSeconds,
+        templates: remainingTemplates.map((template) => ({
+          ...template,
+          availableSlots: template.remainingSlots,
+        })),
+        capabilityWindowSeconds: params.capabilityWindowSeconds,
+      });
+    });
+    results.push(result);
+    consumeInstallationTemplateSlots(remainingTemplates, result.reservations);
+    params.onReservations?.(result.reservations);
+    remainingMaxReservations -= result.reservations.reduce(
+      (total, reservation) => total + reservation.count,
+      0,
+    );
+  }
+  return {
+    stats: results.flatMap((result) => result.stats),
+    reservations: results.flatMap((result) => result.reservations),
+  };
+}
 
-  const demandRows = (
+function consumeInstallationTemplateSlots(
+  templates: NormalizedTemplate[],
+  reservations: ReservationGrant[],
+): void {
+  for (const reservation of reservations) {
+    const satisfyingTemplates = templates
+      .filter((template) => isSubset(reservation.labels, template.labels))
+      .sort(
+        (a, b) => a.labels.length - b.labels.length || a.templateKey.localeCompare(b.templateKey),
+      );
+    drawSlots(satisfyingTemplates, reservation.count);
+  }
+}
+
+async function pollDemandAndReserveLockedTx(
+  tx: Tx,
+  params: PollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  let demandRows = (
     await tx
       .select({
         requiredLabels: pendingJobExecutions.requiredLabels,
@@ -75,6 +158,15 @@ export async function pollDemandAndReserveTx(
     ...row,
     oldestQueuedAt: new Date(row.oldestQueuedAt),
   }));
+  if (params.capabilityWindowSeconds !== undefined) {
+    const capabilityLabels = await listActiveWorkspaceCapabilityLabelsTx(tx, {
+      workspaceId: params.workspaceId,
+      windowSeconds: params.capabilityWindowSeconds,
+    });
+    demandRows = demandRows.filter(
+      (demand) => !capabilityLabels.some((labels) => isSubset(demand.requiredLabels, labels)),
+    );
+  }
 
   const activeReservationRows = await tx
     .select({
@@ -114,6 +206,8 @@ export async function pollDemandAndReserveTx(
   const stats: DemandStat[] = [];
   const grants: ReservationGrant[] = [];
   let remainingMaxReservations = params.maxReservations;
+  const workspaceIdField =
+    params.capabilityWindowSeconds === undefined ? {} : {workspaceId: params.workspaceId};
 
   for (const demand of sortDemandRows(demandRows)) {
     const satisfyingTemplates = templates
@@ -152,6 +246,7 @@ export async function pollDemandAndReserveTx(
       reservedAfterGrant += grant;
       grants.push({
         reservationId: inserted.id,
+        ...workspaceIdField,
         labels: demand.requiredLabels,
         count: grant,
         expiresAt: inserted.expiresAt,
@@ -159,6 +254,7 @@ export async function pollDemandAndReserveTx(
     }
 
     stats.push({
+      ...workspaceIdField,
       labels: demand.requiredLabels,
       queued: demand.queued,
       reserved: reservedAfterGrant,
@@ -169,6 +265,55 @@ export async function pollDemandAndReserveTx(
   return {stats, reservations: grants};
 }
 
+async function listInstallationDemandWorkspaceIds(eligibleWorkspaceIds: ReadonlySet<string>) {
+  if (eligibleWorkspaceIds.size === 0) return [];
+  const rows = await db()
+    .select({
+      workspaceId: pendingJobExecutions.workspaceId,
+      oldestQueuedAt: sql<Date>`min(${pendingJobExecutions.createdAt})`,
+    })
+    .from(pendingJobExecutions)
+    .where(inArray(pendingJobExecutions.workspaceId, [...eligibleWorkspaceIds]))
+    .groupBy(pendingJobExecutions.workspaceId)
+    .orderBy(
+      asc(sql`min(${pendingJobExecutions.createdAt})`),
+      asc(pendingJobExecutions.workspaceId),
+    );
+  return rows.map((row) => row.workspaceId);
+}
+
+export async function listQueuedDemandWorkspaceIds(): Promise<string[]> {
+  const rows = await db()
+    .select({workspaceId: pendingJobExecutions.workspaceId})
+    .from(pendingJobExecutions)
+    .groupBy(pendingJobExecutions.workspaceId);
+  return rows.map((row) => row.workspaceId);
+}
+
+async function listActiveWorkspaceCapabilityLabelsTx(
+  tx: Tx,
+  params: {workspaceId: string; windowSeconds: number},
+): Promise<string[][]> {
+  const rows = await tx
+    .select({labels: provisionerCapabilitySnapshots.labels})
+    .from(provisionerCapabilitySnapshots)
+    .innerJoin(
+      provisionerTokens,
+      eq(provisionerTokens.id, provisionerCapabilitySnapshots.provisionerId),
+    )
+    .where(
+      and(
+        eq(provisionerCapabilitySnapshots.workspaceId, params.workspaceId),
+        eq(provisionerTokens.workspaceId, params.workspaceId),
+        eq(provisionerTokens.scope, 'workspace'),
+        sql`${provisionerCapabilitySnapshots.advertisedAt} > now() - (${params.windowSeconds} || ' seconds')::interval`,
+        sql`${provisionerTokens.revokedAt} is null`,
+        sql`(${provisionerTokens.expiresAt} is null or ${provisionerTokens.expiresAt} > now())`,
+      ),
+    );
+  return rows.map((row) => row.labels);
+}
+
 export async function deleteExpiredReservations(params?: {limit?: number}): Promise<number> {
   const expiredIds = db()
     .select({id: reservations.id})