@shipfox/api-runners 12.1.0 → 12.3.0

package/src/db/runner-sessions.ts

@@ -1,7 +1,10 @@
 import type {RunnerToolCapabilitiesDto} from '@shipfox/api-runners-dto';
 import {and, asc, eq, gt, inArray, isNull, lt, notExists, or, sql} from 'drizzle-orm';
 import type {RunnerSession} from '#core/entities/runner-session.js';
+import {EmptyRunnerLabelsError} from '#core/errors.js';
+import {sanitizeRunnerLabelsOrThrow} from '#core/runner-labels.js';
 import {db} from './db.js';
+import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {runnerActivationTokens} from './schema/runner-activation-tokens.js';
 import {runnerControlSessions} from './schema/runner-control-sessions.js';
 import {providerRunners} from './schema/runner-instances.js';
@@ -56,40 +59,61 @@ export async function createRunnerSessionConsumingActivationToken(params: {
   toolCapabilities?: RunnerToolCapabilitiesDto | null;
 }) {
   return await db().transaction(async (tx) => {
-    const [token] = await tx
+    const [runner] = await tx
       .select({
-        id: runnerActivationTokens.id,
-        runnerInstanceId: runnerActivationTokens.runnerInstanceId,
+        activationTokenId: runnerActivationTokens.id,
+        runnerInstanceId: providerRunners.id,
         workspaceId: providerRunners.workspaceId,
         provisionerId: providerRunners.provisionerId,
         providerRunnerId: providerRunners.providerRunnerId,
+        runnerSessionId: providerRunners.runnerSessionId,
       })
       .from(runnerActivationTokens)
       .innerJoin(providerRunners, eq(providerRunners.id, runnerActivationTokens.runnerInstanceId))
+      .where(eq(runnerActivationTokens.id, params.activationTokenId))
+      .limit(1)
+      .for('update', {of: providerRunners});
+    if (!runner?.workspaceId || !runner.providerRunnerId || runner.runnerSessionId)
+      throw new Error('Runner activation token is invalid, expired, or has already been used');
+
+    const [activationToken] = await tx
+      .select({id: runnerActivationTokens.id})
+      .from(runnerActivationTokens)
       .where(
         and(
-          eq(runnerActivationTokens.id, params.activationTokenId),
+          eq(runnerActivationTokens.id, runner.activationTokenId),
+          eq(runnerActivationTokens.runnerInstanceId, runner.runnerInstanceId),
           isNull(runnerActivationTokens.consumedAt),
           isNull(runnerActivationTokens.revokedAt),
           gt(runnerActivationTokens.expiresAt, sql`now()`),
-          isNull(providerRunners.runnerSessionId),
         ),
       )
       .limit(1)
       .for('update');
-    if (!token?.workspaceId || !token.providerRunnerId)
+    if (!activationToken)
       throw new Error('Runner activation token is invalid, expired, or has already been used');
+
+    const [provisioner] = await tx
+      .select({scope: provisionerTokens.scope})
+      .from(provisionerTokens)
+      .where(eq(provisionerTokens.id, runner.provisionerId))
+      .limit(1);
+    const labels = sanitizeRunnerLabelsOrThrow(params.labels, {
+      scope: provisioner?.scope ?? 'workspace',
+      source: 'activation runner registration',
+    });
+    if (labels.length === 0) throw new EmptyRunnerLabelsError();
     const [session] = await tx
       .insert(runnerSessions)
       .values({
-        workspaceId: token.workspaceId,
+        workspaceId: runner.workspaceId,
         scope: 'workspace',
-        registrationTokenId: token.id,
+        registrationTokenId: activationToken.id,
         registrationTokenKind: 'activation',
-        runnerInstanceId: token.runnerInstanceId,
-        provisionerId: token.provisionerId,
-        providerRunnerId: token.providerRunnerId,
-        labels: params.labels,
+        runnerInstanceId: runner.runnerInstanceId,
+        provisionerId: runner.provisionerId,
+        providerRunnerId: runner.providerRunnerId,
+        labels,
         toolCapabilities: params.toolCapabilities ?? null,
         toolCapabilitiesReportedAt: params.toolCapabilities ? sql`now()` : null,
         maxClaims: 1,
@@ -100,13 +124,13 @@ export async function createRunnerSessionConsumingActivationToken(params: {
     await tx
       .update(runnerActivationTokens)
       .set({consumedAt: sql`now()`, consumedSessionId: session.id})
-      .where(eq(runnerActivationTokens.id, token.id));
+      .where(eq(runnerActivationTokens.id, activationToken.id));
     await tx
       .update(providerRunners)
       .set({runnerSessionId: session.id, updatedAt: sql`now()`})
       .where(
         and(
-          eq(providerRunners.id, token.runnerInstanceId),
+          eq(providerRunners.id, runner.runnerInstanceId),
           isNull(providerRunners.runnerSessionId),
         ),
       );
@@ -115,7 +139,7 @@ export async function createRunnerSessionConsumingActivationToken(params: {
       .set({closedAt: sql`now()`, closeReason: 'activated'})
       .where(
         and(
-          eq(runnerControlSessions.runnerInstanceId, token.runnerInstanceId),
+          eq(runnerControlSessions.runnerInstanceId, runner.runnerInstanceId),
           isNull(runnerControlSessions.closedAt),
         ),
       );

package/src/db/schema/runner-instances.ts

@@ -60,6 +60,9 @@ export const providerRunners = pgTable(
       .where(
         sql`${table.intendedReservationId} is not null and ${table.reservationReleasedAt} is null`,
       ),
+    index('runners_runner_instances_intended_reservation_idx')
+      .on(table.intendedReservationId)
+      .where(sql`${table.intendedReservationId} is not null`),
     index('runners_runner_instances_provisioner_reservation_idx')
       .on(table.provisionerId, table.reservationId)
       .where(sql`${table.reservationId} is not null`),

package/src/metrics/index.ts

@@ -1,3 +1,4 @@
+export type {RunnerReservationPromotionFailureReason} from './instance.js';
 export {
   jobExecutionClaimedCount,
   jobExecutionEnqueuedCount,
@@ -7,6 +8,7 @@ export {
   providerRunnerReconcileCallCount,
   providerRunnerTerminateIntentHonoredCount,
   providerRunnerTerminateIntentIssuedCount,
+  recordRunnerReservationPromotionFailure,
   recordRunnersRateLimitCheck,
   recordRunnersRateLimitPruneFailure,
 } from './instance.js';

package/src/metrics/instance.ts

@@ -83,6 +83,18 @@ export const reservationReleasedCount = meter.createCounter<Record<string, never
   {description: 'Reservation units released from terminal provisioned runner reports'},
 );
 
+export type RunnerReservationPromotionFailureReason =
+  | 'reservation-expired'
+  | 'reservation-not-found'
+  | 'already-assigned'
+  | 'not-assignable';
+
+export const runnerReservationPromotionFailureCount = meter.createCounter<{
+  reason: RunnerReservationPromotionFailureReason;
+}>('runners_reservation_promotion_failures', {
+  description: 'Runner reservation promotion failures during enrollment by reason',
+});
+
 export type RunnersRateLimitAction = 'provisioner-mint' | 'ephemeral-register';
 export type RunnersRateLimitScope = 'provisioner' | 'ephemeral-token';
 export type RunnersRateLimitOutcome = 'allowed' | 'blocked' | 'unavailable';
@@ -107,6 +119,12 @@ function recordMetric(record: () => void): void {
   }
 }
 
+export function recordRunnerReservationPromotionFailure(
+  reason: RunnerReservationPromotionFailureReason,
+): void {
+  recordMetric(() => runnerReservationPromotionFailureCount.add(1, {reason}));
+}
+
 export function recordRunnersRateLimitCheck(params: {
   action: RunnersRateLimitAction;
   scope: RunnersRateLimitScope;

package/src/metrics/service.test.ts

@@ -0,0 +1,99 @@
+const mocks = vi.hoisted(() => {
+  const gauges = {
+    enrolledRunnersWithoutRecentReport: {},
+    pendingJobExecutions: {},
+    runningJobExecutions: {},
+  };
+  const gaugeByName = {
+    runners_enrolled_without_recent_report: gauges.enrolledRunnersWithoutRecentReport,
+    runners_pending_job_executions: gauges.pendingJobExecutions,
+    runners_running_job_executions: gauges.runningJobExecutions,
+  };
+  return {
+    addBatchObservableCallback: vi.fn(),
+    countStaleEnrolledRunnerInstances: vi.fn(),
+    createObservableGauge: vi.fn((name: string) => gaugeByName[name as keyof typeof gaugeByName]),
+    gauges,
+    getMeter: vi.fn(),
+    getJobExecutionQueueDepth: vi.fn(),
+    getServiceMetricsProvider: vi.fn(),
+  };
+});
+
+vi.mock('@shipfox/node-opentelemetry', () => ({
+  getServiceMetricsProvider: mocks.getServiceMetricsProvider,
+}));
+vi.mock('#config.js', () => ({
+  config: {RUNNER_STALE_PROVISIONED_RUNNER_THRESHOLD_SECONDS: 300},
+}));
+vi.mock('#db/job-executions.js', () => ({
+  getJobExecutionQueueDepth: mocks.getJobExecutionQueueDepth,
+}));
+vi.mock('#db/runner-instances.js', () => ({
+  countStaleEnrolledRunnerInstances: mocks.countStaleEnrolledRunnerInstances,
+}));
+
+import {registerRunnersServiceMetrics} from './service.js';
+
+describe('registerRunnersServiceMetrics', () => {
+  beforeEach(() => {
+    mocks.addBatchObservableCallback.mockReset();
+    mocks.countStaleEnrolledRunnerInstances.mockReset();
+    mocks.createObservableGauge.mockClear();
+    mocks.getJobExecutionQueueDepth.mockReset();
+    mocks.getMeter.mockReset();
+    mocks.getServiceMetricsProvider.mockReset();
+    mocks.getJobExecutionQueueDepth.mockResolvedValue({
+      pendingJobExecutions: 0,
+      runningJobExecutions: 0,
+    });
+    mocks.getMeter.mockReturnValue({
+      createObservableGauge: mocks.createObservableGauge,
+      addBatchObservableCallback: mocks.addBatchObservableCallback,
+    });
+    mocks.getServiceMetricsProvider.mockReturnValue({getMeter: mocks.getMeter});
+  });
+
+  it('observes enrolled runners without recent reports after the grace window', async () => {
+    mocks.countStaleEnrolledRunnerInstances.mockResolvedValue(2);
+
+    registerRunnersServiceMetrics();
+    const callback = mocks.addBatchObservableCallback.mock.calls[0]?.[0];
+    if (typeof callback !== 'function') throw new Error('Expected metrics callback');
+    const observer = {observe: vi.fn()};
+
+    await callback(observer);
+
+    expect(mocks.createObservableGauge).toHaveBeenCalledWith(
+      'runners_enrolled_without_recent_report',
+      {
+        description:
+          'Running enrolled runners with a live control session, no workspace or runner session, and no recent provisioner report after the stale-runner grace window',
+      },
+    );
+    expect(mocks.countStaleEnrolledRunnerInstances).toHaveBeenCalledWith({graceSeconds: 300});
+    expect(observer.observe).toHaveBeenCalledWith(
+      mocks.gauges.enrolledRunnersWithoutRecentReport,
+      2,
+    );
+  });
+
+  it('keeps queue gauges observable when the enrolled-runner query fails', async () => {
+    mocks.getJobExecutionQueueDepth.mockResolvedValue({
+      pendingJobExecutions: 3,
+      runningJobExecutions: 4,
+    });
+    mocks.countStaleEnrolledRunnerInstances.mockRejectedValue(new Error('database unavailable'));
+
+    registerRunnersServiceMetrics();
+    const callback = mocks.addBatchObservableCallback.mock.calls[0]?.[0];
+    if (typeof callback !== 'function') throw new Error('Expected metrics callback');
+    const observer = {observe: vi.fn()};
+
+    await callback(observer);
+
+    expect(observer.observe).toHaveBeenCalledWith(mocks.gauges.pendingJobExecutions, 3);
+    expect(observer.observe).toHaveBeenCalledWith(mocks.gauges.runningJobExecutions, 4);
+    expect(observer.observe).toHaveBeenCalledTimes(2);
+  });
+});

package/src/metrics/service.ts

@@ -1,5 +1,7 @@
 import {getServiceMetricsProvider} from '@shipfox/node-opentelemetry';
+import {config} from '#config.js';
 import {getJobExecutionQueueDepth} from '#db/job-executions.js';
+import {countStaleEnrolledRunnerInstances} from '#db/runner-instances.js';
 
 export function registerRunnersServiceMetrics(): void {
   const meter = getServiceMetricsProvider().getMeter('runners');
@@ -10,13 +12,31 @@ export function registerRunnersServiceMetrics(): void {
   const runningJobExecutions = meter.createObservableGauge('runners_running_job_executions', {
     description: 'Job executions currently claimed by a runner and in progress',
   });
+  const enrolledRunnersWithoutRecentReport = meter.createObservableGauge(
+    'runners_enrolled_without_recent_report',
+    {
+      description:
+        'Running enrolled runners with a live control session, no workspace or runner session, and no recent provisioner report after the stale-runner grace window',
+    },
+  );
 
   meter.addBatchObservableCallback(
     async (observer) => {
-      const depth = await getJobExecutionQueueDepth();
-      observer.observe(pendingJobExecutions, depth.pendingJobExecutions);
-      observer.observe(runningJobExecutions, depth.runningJobExecutions);
+      const [depthResult, staleEnrolledRunnerCountResult] = await Promise.allSettled([
+        getJobExecutionQueueDepth(),
+        countStaleEnrolledRunnerInstances({
+          graceSeconds: config.RUNNER_STALE_PROVISIONED_RUNNER_THRESHOLD_SECONDS,
+        }),
+      ]);
+
+      if (depthResult.status === 'fulfilled') {
+        observer.observe(pendingJobExecutions, depthResult.value.pendingJobExecutions);
+        observer.observe(runningJobExecutions, depthResult.value.runningJobExecutions);
+      }
+      if (staleEnrolledRunnerCountResult.status === 'fulfilled') {
+        observer.observe(enrolledRunnersWithoutRecentReport, staleEnrolledRunnerCountResult.value);
+      }
     },
-    [pendingJobExecutions, runningJobExecutions],
+    [pendingJobExecutions, runningJobExecutions, enrolledRunnersWithoutRecentReport],
   );
 }

package/src/presentation/routes/late-runner-enrollment.test.ts

@@ -0,0 +1,230 @@
+import {AUTH_PROVISIONER_TOKEN, AUTH_USER, setProvisionerContext} from '@shipfox/api-auth-context';
+import {
+  type AuthMethod,
+  ClientError,
+  closeApp,
+  createApp,
+  extractBearerToken,
+} from '@shipfox/node-fastify';
+import {vi} from '@shipfox/vitest/vi';
+import {eq} from 'drizzle-orm';
+import type {FastifyInstance, FastifyRequest} from 'fastify';
+import {db} from '#db/db.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+import {runnerSessions} from '#db/schema/runner-sessions.js';
+import {runnerReservationPromotionFailureCount} from '#metrics/instance.js';
+import {
+  createRunnerControlSessionAuthMethod,
+  createRunnerRegistrationTokenAuthMethod,
+} from '#presentation/auth/index.js';
+import {
+  arrangeDeletedRunnerEnrollment,
+  arrangeExpiredRunnerEnrollment,
+  fakeLeaseTokenAuthMethod,
+  fakeRunnerSessionAuthMethod,
+  pendingJobFactory,
+  provisionerTokenFactory,
+  runnersTestAuthClient,
+} from '#test/index.js';
+import {createRunnerRoutes} from './index.js';
+
+let workspaceToken: string;
+const fakeUserAuth: AuthMethod = {name: AUTH_USER, authenticate: () => Promise.resolve()};
+
+describe('late runner enrollment recovery', () => {
+  let app: FastifyInstance;
+  let provisionerId: string;
+  let workspaceId: string;
+
+  const provisionerAuth: AuthMethod = {
+    name: AUTH_PROVISIONER_TOKEN,
+    authenticate: (request: FastifyRequest) => {
+      if (extractBearerToken(request.headers.authorization) !== workspaceToken)
+        throw new ClientError('Invalid provisioner token', 'unauthorized', {status: 401});
+      setProvisionerContext(request, {
+        scope: 'workspace',
+        workspaceId,
+        provisionerTokenId: provisionerId,
+      });
+      return Promise.resolve();
+    },
+  };
+
+  beforeAll(async () => {
+    app = await createApp({
+      auth: [
+        fakeUserAuth,
+        provisionerAuth,
+        createRunnerRegistrationTokenAuthMethod(),
+        createRunnerControlSessionAuthMethod(),
+        fakeRunnerSessionAuthMethod,
+        fakeLeaseTokenAuthMethod,
+      ],
+      routes: createRunnerRoutes(runnersTestAuthClient),
+      swagger: false,
+    });
+    await app.ready();
+  });
+
+  afterAll(async () => {
+    await closeApp();
+  });
+
+  beforeEach(async () => {
+    workspaceId = crypto.randomUUID();
+    workspaceToken = `late-runner-workspace-provisioner-${crypto.randomUUID()}`;
+    const provisioner = await provisionerTokenFactory.create(
+      {scope: 'workspace', workspaceId},
+      {transient: {rawToken: workspaceToken}},
+    );
+    provisionerId = provisioner.id;
+  });
+
+  it('keeps an enrolled runner recoverable when its reservation expires first', async () => {
+    const arrangement = await arrangeExpiredRunnerEnrollment({provisionerId, workspaceId});
+    const failureSpy = vi.spyOn(runnerReservationPromotionFailureCount, 'add');
+
+    try {
+      await enrollRunner(arrangement.controlSessionToken);
+
+      expect(failureSpy).toHaveBeenCalledWith(1, {reason: 'reservation-expired'});
+    } finally {
+      failureSpy.mockRestore();
+    }
+
+    const [runner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(runner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: arrangement.reservationId,
+      assignedAt: null,
+      runnerSessionId: null,
+      state: 'running',
+      providerRunnerId: expect.any(String),
+    });
+  });
+
+  it('keeps an enrolled runner recoverable when its reservation was swept away', async () => {
+    const arrangement = await arrangeDeletedRunnerEnrollment({provisionerId, workspaceId});
+    const failureSpy = vi.spyOn(runnerReservationPromotionFailureCount, 'add');
+
+    try {
+      await enrollRunner(arrangement.controlSessionToken);
+
+      expect(failureSpy).toHaveBeenCalledWith(1, {reason: 'reservation-not-found'});
+    } finally {
+      failureSpy.mockRestore();
+    }
+
+    const [runner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(runner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: arrangement.reservationId,
+      assignedAt: null,
+      runnerSessionId: null,
+      state: 'running',
+      providerRunnerId: expect.any(String),
+    });
+  });
+
+  it('rebinds a stranded runner through activation and claims its pending job', async () => {
+    const arrangement = await arrangeExpiredRunnerEnrollment({provisionerId, workspaceId});
+    await enrollRunner(arrangement.controlSessionToken);
+    const pendingJob = await pendingJobFactory.create({workspaceId, requiredLabels: ['linux']});
+
+    const demand = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        wait_seconds: 0,
+        max_reservations: 1,
+        reservation_ttl_seconds: 60,
+        templates: [
+          {
+            template_key: 'linux',
+            labels: ['linux'],
+            available_slots: 1,
+            starting: 0,
+            running: 1,
+          },
+        ],
+      },
+    });
+
+    expect(demand.statusCode).toBe(200);
+    const reboundReservationId = demand.json().reservations[0]?.reservation_id;
+    expect(reboundReservationId).toEqual(expect.any(String));
+
+    const [reboundRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(reboundRunner).toMatchObject({
+      workspaceId,
+      reservationId: reboundReservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+      runnerSessionId: null,
+      state: 'running',
+    });
+
+    const assignment = await app.inject({
+      method: 'GET',
+      url: '/runner-control/assignment?wait_seconds=1',
+      headers: {authorization: `Bearer ${arrangement.controlSessionToken}`},
+    });
+    expect(assignment.statusCode).toBe(200);
+    const activationToken = assignment.json().activation_token;
+    expect(activationToken).toEqual(expect.any(String));
+
+    const registered = await app.inject({
+      method: 'POST',
+      url: '/runners/register',
+      headers: {authorization: `Bearer ${activationToken}`},
+      payload: {labels: ['linux']},
+    });
+    expect(registered.statusCode).toBe(200);
+    expect(registered.json()).toMatchObject({mode: 'activation', max_claims: 1});
+
+    const claimed = await app.inject({
+      method: 'POST',
+      url: '/runners/jobs/request',
+      headers: {authorization: `Bearer ${registered.json().session_token}`},
+    });
+    expect(claimed.statusCode).toBe(200);
+    expect(claimed.json()).toMatchObject({job_id: pendingJob.jobId});
+
+    const [session] = await db()
+      .select()
+      .from(runnerSessions)
+      .where(eq(runnerSessions.id, registered.json().session_id));
+    expect(session).toMatchObject({
+      registrationTokenKind: 'activation',
+      runnerInstanceId: arrangement.runnerInstanceId,
+      provisionerId,
+      providerRunnerId: reboundRunner?.providerRunnerId,
+      maxClaims: 1,
+      claimsUsed: 1,
+    });
+  });
+
+  async function enrollRunner(controlSessionToken: string): Promise<void> {
+    const enrolled = await app.inject({
+      method: 'POST',
+      url: '/runner-control/enrollment',
+      headers: {authorization: `Bearer ${controlSessionToken}`},
+      payload: {labels: ['linux'], provider_kind: 'docker', protocol_version: '1'},
+    });
+
+    expect(enrolled.statusCode).toBe(200);
+    expect(enrolled.json()).toEqual({activation_token: null});
+  }
+});