@shipfox/api-runners 12.1.0 → 12.2.0

package/src/core/runner-labels.ts

@@ -0,0 +1,51 @@
+import {logger} from '@shipfox/node-opentelemetry';
+import {canonicalizeLabels} from '@shipfox/runner-labels';
+import {runnerReservedLabels} from '#config.js';
+import {RunnerLabelsReservedError} from './errors.js';
+
+export type RunnerLabelScope = 'installation' | 'workspace' | 'manual';
+
+type SanitizeRunnerLabelsParams = {
+  scope: RunnerLabelScope;
+  source: string;
+  logLevel?: 'debug' | 'warn';
+  reservedLabels?: readonly string[];
+};
+
+function sanitizeCanonicalRunnerLabels(
+  canonicalLabels: readonly string[],
+  params: SanitizeRunnerLabelsParams,
+): string[] {
+  if (params.scope === 'installation') return [...canonicalLabels];
+
+  const reservedLabels = new Set(canonicalizeLabels(params.reservedLabels ?? runnerReservedLabels));
+  if (reservedLabels.size === 0) return [...canonicalLabels];
+
+  const removedLabels = canonicalLabels.filter((label) => reservedLabels.has(label));
+  if (removedLabels.length === 0) return [...canonicalLabels];
+
+  const details = {source: params.source, removedLabels};
+  if (params.logLevel === 'debug')
+    logger().debug(details, 'Removed reserved runner labels from a non-installation source');
+  else logger().warn(details, 'Removed reserved runner labels from a non-installation source');
+  return canonicalLabels.filter((label) => !reservedLabels.has(label));
+}
+
+export function sanitizeRunnerLabels(
+  labels: readonly string[],
+  params: SanitizeRunnerLabelsParams,
+): string[] {
+  return sanitizeCanonicalRunnerLabels([...canonicalizeLabels(labels)], params);
+}
+
+export function sanitizeRunnerLabelsOrThrow(
+  labels: readonly string[],
+  params: SanitizeRunnerLabelsParams,
+): string[] {
+  const canonicalLabels = [...canonicalizeLabels(labels)];
+  const sanitizedLabels = sanitizeCanonicalRunnerLabels(canonicalLabels, params);
+  if (sanitizedLabels.length === 0 && canonicalLabels.length > 0) {
+    throw new RunnerLabelsReservedError(canonicalLabels);
+  }
+  return sanitizedLabels;
+}

package/src/core/runner-sessions.test.ts

@@ -9,12 +9,14 @@ import {
   getRunnerSessionTokenClaims,
   manualRegistrationTokenFactory,
   providerRunnerFactory,
+  provisionerTokenFactory,
   runnersTestAuthClient,
 } from '#test/index.js';
 import {
   EmptyRunnerLabelsError,
   RegistrationTokenConsumedError,
   RegistrationTokenWorkspaceMismatchError,
+  RunnerLabelsReservedError,
 } from './errors.js';
 import {issueRunnerActivationToken} from './runner-activation.js';
 import {registerRunnerSession} from './runner-sessions.js';
@@ -66,6 +68,15 @@ describe('registerRunnerSession', () => {
     ).rejects.toBeInstanceOf(EmptyRunnerLabelsError);
   });
 
+  it('names labels removed by the reserved-label policy', () => {
+    const error = new RunnerLabelsReservedError(['shipfox-managed']);
+
+    expect(error.message).toBe(
+      'All supplied runner labels are reserved for installation-scope provisioners: shipfox-managed',
+    );
+    expect(error.labels).toEqual(['shipfox-managed']);
+  });
+
   it('consumes an ephemeral token and creates a one-claim session', async () => {
     const token = await ephemeralRegistrationTokenFactory.create({workspaceId});
 
@@ -223,7 +234,8 @@ describe('activation runner sessions', () => {
 
   beforeEach(async () => {
     workspaceId = crypto.randomUUID();
-    provisionerId = crypto.randomUUID();
+    const provisioner = await provisionerTokenFactory.create({workspaceId});
+    provisionerId = provisioner.id;
     const runner = await providerRunnerFactory.create({workspaceId, provisionerId});
     runnerInstanceId = runner.id;
   });
@@ -296,4 +308,25 @@ describe('activation runner sessions', () => {
     expect(storedToken?.consumedAt).toBeInstanceOf(Date);
     expect(sessions).toHaveLength(1);
   });
+
+  it('strips reserved labels from workspace activation registration', async () => {
+    const rawToken = await issueRunnerActivationToken({
+      runnerInstanceId,
+      provisionerId,
+      ttlSeconds: 60,
+    });
+    const [activationToken] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.runnerInstanceId, runnerInstanceId));
+    if (!rawToken || !activationToken) throw new Error('Activation token was not created');
+
+    const result = await registerRunnerSession({
+      auth: runnersTestAuthClient,
+      credential: {kind: 'activation', activationTokenId: activationToken.id, workspaceId},
+      labels: ['linux', 'shipfox-managed'],
+    });
+
+    expect(result.session.labels).toEqual(['linux']);
+  });
 });

package/src/core/runner-sessions.ts

@@ -8,6 +8,7 @@ import {
 } from '#db/runner-sessions.js';
 import type {RunnerSession} from './entities/runner-session.js';
 import {EmptyRunnerLabelsError} from './errors.js';
+import {sanitizeRunnerLabelsOrThrow} from './runner-labels.js';
 
 export interface RegisterRunnerSessionResult {
   session: RunnerSession;
@@ -38,7 +39,13 @@ export async function registerRunnerSession(params: {
   labels: string[];
   toolCapabilities?: RunnerToolCapabilitiesDto | null;
 }): Promise<RegisterRunnerSessionResult> {
-  const labels = [...canonicalizeLabels(params.labels)];
+  const labels =
+    params.credential.kind === 'manual'
+      ? sanitizeRunnerLabelsOrThrow(params.labels, {
+          scope: 'manual',
+          source: 'manual runner registration',
+        })
+      : [...canonicalizeLabels(params.labels)];
   if (labels.length === 0) throw new EmptyRunnerLabelsError();
 
   const mode = params.credential.kind;

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

@@ -4,6 +4,7 @@ import type {EphemeralRegistrationToken} from '#core/entities/ephemeral-registra
 import {
   ActiveEphemeralRegistrationTokenExistsError,
   ActiveEphemeralRegistrationTokensExistError,
+  EmptyRunnerLabelsError,
   RegistrationTokenBatchExceedsReservationError,
   RegistrationTokenConsumedError,
   RegistrationTokenExpiredError,
@@ -11,11 +12,13 @@ import {
   ReservationExpiredError,
   ReservationNotFoundError,
 } from '#core/errors.js';
+import {sanitizeRunnerLabelsOrThrow} from '#core/runner-labels.js';
 import {db} from './db.js';
 import {
   ephemeralRegistrationTokens,
   toEphemeralRegistrationToken,
 } from './schema/ephemeral-registration-tokens.js';
+import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {reservations} from './schema/reservations.js';
 import {providerRunners} from './schema/runner-instances.js';
 import {runnerSessions, toRunnerSession} from './schema/runner-sessions.js';
@@ -307,6 +310,17 @@ export async function createRunnerSessionConsumingEphemeralToken(params: {
       );
     }
 
+    const [provisioner] = await tx
+      .select({scope: provisionerTokens.scope})
+      .from(provisionerTokens)
+      .where(eq(provisionerTokens.id, consumed[0].provisionerId))
+      .limit(1);
+    const labels = sanitizeRunnerLabelsOrThrow(params.labels, {
+      scope: provisioner?.scope ?? 'workspace',
+      source: 'ephemeral runner registration',
+    });
+    if (labels.length === 0) throw new EmptyRunnerLabelsError();
+
     const [session] = await tx
       .insert(runnerSessions)
       .values({
@@ -316,7 +330,7 @@ export async function createRunnerSessionConsumingEphemeralToken(params: {
         registrationTokenKind: 'ephemeral',
         provisionerId: consumed[0].provisionerId,
         providerRunnerId: consumed[0].providerRunnerId,
-        labels: params.labels,
+        labels,
         toolCapabilities: params.toolCapabilities ?? null,
         toolCapabilitiesReportedAt: params.toolCapabilities ? sql`now()` : null,
         maxClaims: params.maxClaims,

package/src/db/reservations.test.ts

@@ -461,6 +461,64 @@ describe('pollDemandAndReserve', () => {
     expect(storedToken?.revokedAt).toBeInstanceOf(Date);
   });
 
+  it('preserves activation tokens for active runners with an intended reservation', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Expected reservation');
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        workspaceId,
+        provisionerId,
+        providerRunnerId: crypto.randomUUID(),
+        reservationId: reservation.id,
+        intendedReservationId: reservation.id,
+        runnerSessionId: crypto.randomUUID(),
+        labels: ['linux'],
+        state: 'running',
+        reportedAt: new Date(),
+      })
+      .returning();
+    if (!runner) throw new Error('Expected runner instance');
+    const [activationToken] = await db()
+      .insert(runnerActivationTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!activationToken) throw new Error('Expected activation token');
+
+    const deleted = await deleteReservationsByIds([reservation.id]);
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const [storedToken] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.id, activationToken.id));
+    expect(deleted).toBe(1);
+    expect(storedRunner).toMatchObject({
+      reservationId: reservation.id,
+      intendedReservationId: null,
+      runnerSessionId: expect.any(String),
+      reservationReleasedAt: null,
+    });
+    expect(storedToken?.revokedAt).toBeNull();
+  });
+
   it('unbinds prewarmed runners and revokes activation tokens when deleting expired reservations', async () => {
     const runner = await createIdleRunner({labels: ['linux']});
     await createPendingJobs(1, ['linux']);
@@ -509,6 +567,115 @@ describe('pollDemandAndReserve', () => {
     expect(storedToken?.revokedAt).toBeInstanceOf(Date);
   });
 
+  it('clears intended reservations and rebinds runners when deleting a reservation by id', async () => {
+    const [reservation, survivingReservation] = await db()
+      .insert(reservations)
+      .values([
+        {
+          workspaceId,
+          provisionerId,
+          requiredLabels: ['linux'],
+          count: 1,
+          expiresAt: new Date(Date.now() + 60_000),
+        },
+        {
+          workspaceId: crypto.randomUUID(),
+          provisionerId,
+          requiredLabels: ['linux'],
+          count: 1,
+          expiresAt: new Date(Date.now() + 60_000),
+        },
+      ])
+      .returning();
+    if (!reservation || !survivingReservation) throw new Error('Expected reservations');
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: reservation.id,
+    });
+    const survivingRunner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: survivingReservation.id,
+    });
+
+    const deleted = await deleteReservationsByIds([reservation.id]);
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(deleted).toBe(1);
+    expect(storedRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId: null,
+      reservationId: null,
+      assignedAt: null,
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+    const [storedSurvivingRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, survivingRunner.id));
+    expect(storedSurvivingRunner?.intendedReservationId).toBe(survivingReservation.id);
+
+    await createPendingJobs(1, ['linux']);
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+    const reboundReservationId = result.reservations[0]?.reservationId;
+    const [reboundRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(reboundReservationId).toEqual(expect.any(String));
+    expect(reboundRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId,
+      reservationId: reboundReservationId,
+      assignedAt: expect.any(Date),
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+  });
+
+  it('clears intended reservations when deleting expired reservations', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() - 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Expected reservation');
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: reservation.id,
+    });
+
+    const deleted = await deleteExpiredReservations();
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(deleted).toBeGreaterThanOrEqual(1);
+    expect(storedRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId: null,
+      reservationId: null,
+      assignedAt: null,
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+  });
+
   it('decrements reservation units inside a caller transaction', async () => {
     await reservationFactory.create({
       workspaceId,
@@ -645,6 +812,7 @@ describe('pollDemandAndReserve', () => {
     labels: string[];
     createdAt?: Date;
     controlSessionExpiresAt?: Date;
+    intendedReservationId?: string;
   }) {
     const createdAt = params.createdAt ?? new Date();
     const [runner] = await db()
@@ -652,6 +820,7 @@ describe('pollDemandAndReserve', () => {
       .values({
         provisionerId,
         providerRunnerId: crypto.randomUUID(),
+        intendedReservationId: params.intendedReservationId,
         labels: params.labels,
         state: 'running',
         reportedAt: createdAt,

package/src/db/reservations.ts

@@ -10,6 +10,7 @@ import {
   isNotNull,
   isNull,
   lt,
+  or,
   sql,
 } from 'drizzle-orm';
 import type {Tx} from './db.js';
@@ -432,16 +433,26 @@ async function deleteReservationsWithCleanupTx(
 ): Promise<number> {
   if (ids.length === 0) return 0;
 
-  const assignedRunners = await tx
-    .select({id: providerRunners.id, reservationId: providerRunners.reservationId})
+  const affectedRunners = await tx
+    .select({
+      id: providerRunners.id,
+      reservationId: providerRunners.reservationId,
+      intendedReservationId: providerRunners.intendedReservationId,
+      runnerSessionId: providerRunners.runnerSessionId,
+      reservationReleasedAt: providerRunners.reservationReleasedAt,
+    })
     .from(providerRunners)
     .where(
-      and(
-        inArray(providerRunners.reservationId, ids),
-        isNull(providerRunners.runnerSessionId),
-        isNull(providerRunners.reservationReleasedAt),
+      or(
+        and(
+          inArray(providerRunners.reservationId, ids),
+          isNull(providerRunners.runnerSessionId),
+          isNull(providerRunners.reservationReleasedAt),
+        ),
+        inArray(providerRunners.intendedReservationId, ids),
       ),
     )
+    .orderBy(asc(providerRunners.id))
     .for('update');
   const reservationRows = await tx
     .select({id: reservations.id})
@@ -457,8 +468,20 @@ async function deleteReservationsWithCleanupTx(
 
   if (reservationIds.length === 0) return 0;
 
-  const assignedRunnerIds = assignedRunners
-    .filter((runner) => runner.reservationId && reservationIds.includes(runner.reservationId))
+  const assignedRunnerIds = affectedRunners
+    .filter(
+      (runner) =>
+        runner.reservationId &&
+        reservationIds.includes(runner.reservationId) &&
+        !runner.runnerSessionId &&
+        !runner.reservationReleasedAt,
+    )
+    .map((runner) => runner.id);
+  const intendedRunnerIds = affectedRunners
+    .filter(
+      (runner) =>
+        runner.intendedReservationId && reservationIds.includes(runner.intendedReservationId),
+    )
     .map((runner) => runner.id);
   if (assignedRunnerIds.length > 0) {
     await tx
@@ -487,6 +510,17 @@ async function deleteReservationsWithCleanupTx(
         ),
       );
   }
+  if (intendedRunnerIds.length > 0) {
+    await tx
+      .update(providerRunners)
+      .set({intendedReservationId: null, updatedAt: sql`now()`})
+      .where(
+        and(
+          inArray(providerRunners.id, intendedRunnerIds),
+          inArray(providerRunners.intendedReservationId, reservationIds),
+        ),
+      );
+  }
 
   const deleted = await tx
     .delete(reservations)