@shipfox/api-runners 12.2.0 → 12.3.0

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

@@ -4,6 +4,7 @@ import {db} from '#db/db.js';
 import {createRunnerSessionConsumingEphemeralToken} from '#db/ephemeral-registration-tokens.js';
 import {
   attachRunnerInstanceProviderId,
+  countStaleEnrolledRunnerInstances,
   listActiveRunnerInstanceCountsByTemplateTx,
   listActiveRunnerInstances,
   listProvisionerTerminateIntentRowsTx,
@@ -14,6 +15,7 @@ import {
 } from '#db/runner-instances.js';
 import {provisionerTokens} from '#db/schema/provisioner-tokens.js';
 import {reservations} from '#db/schema/reservations.js';
+import {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 {runningJobExecutions} from '#db/schema/running-job-executions.js';
@@ -231,6 +233,139 @@ describe('reportRunnerInstances', () => {
     expect(rows[0]?.reason).toBe('fresh');
   });
 
+  it('preserves a rebound reservation through stale reports and releases it on termination', async () => {
+    const staleReservationId = await createReservation(1);
+    await db()
+      .update(reservations)
+      .set({expiresAt: new Date(Date.now() - 60_000)})
+      .where(eq(reservations.id, staleReservationId));
+    const reboundReservationId = await createReservation(2);
+    const initialReportedAt = new Date(Date.now() - 120_000);
+    const runningReportedAt = new Date(Date.now() - 60_000);
+    await providerRunnerFactory.create({
+      workspaceId,
+      provisionerId,
+      providerRunnerId: 'rebound-runner',
+      reservationId: reboundReservationId,
+      state: 'running',
+      reportedAt: initialReportedAt,
+    });
+
+    const runningReport = await reportRunnerInstances({
+      scope: 'workspace',
+      workspaceId,
+      provisionerId,
+      events: [
+        event({
+          providerRunnerId: 'rebound-runner',
+          reservationId: staleReservationId,
+          reportedAt: runningReportedAt,
+        }),
+      ],
+    });
+    const terminalReport = await reportRunnerInstances({
+      scope: 'workspace',
+      workspaceId,
+      provisionerId,
+      events: [
+        event({
+          providerRunnerId: 'rebound-runner',
+          reservationId: staleReservationId,
+          state: 'failed',
+          reportedAt: new Date(runningReportedAt.getTime() + 1_000),
+        }),
+      ],
+    });
+
+    const [runner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.providerRunnerId, 'rebound-runner'));
+    const reservationRows = await reservationRowsFor({workspaceId, provisionerId});
+    expect(runningReport).toEqual({
+      accepted: 1,
+      reservationsReleased: 0,
+      terminateIntentsHonored: [],
+    });
+    expect(terminalReport).toEqual({
+      accepted: 1,
+      reservationsReleased: 1,
+      terminateIntentsHonored: [],
+    });
+    expect(runner).toMatchObject({
+      reservationId: reboundReservationId,
+      reservationReleasedAt: expect.any(Date),
+    });
+    expect(
+      reservationRows.find((reservation) => reservation.id === reboundReservationId)?.count,
+    ).toBe(1);
+    expect(
+      reservationRows.find((reservation) => reservation.id === staleReservationId)?.count,
+    ).toBe(1);
+  });
+
+  it('keeps the stored reservation when a report carries a stale reservation id', async () => {
+    const staleReservationId = await createReservation(1);
+    const boundReservationId = await createReservation(1);
+    await db()
+      .insert(providerRunners)
+      .values({
+        workspaceId,
+        provisionerId,
+        reservationId: boundReservationId,
+        providerRunnerId: 'provisioned-runner-1',
+        state: 'running',
+        reportedAt: new Date('2025-01-01T00:00:00.000Z'),
+      });
+
+    await reportRunnerInstances({
+      scope: 'workspace',
+      workspaceId,
+      provisionerId,
+      events: [
+        event({
+          providerRunnerId: 'provisioned-runner-1',
+          reservationId: staleReservationId,
+          state: 'running',
+          reportedAt: new Date('2025-01-01T00:01:00.000Z'),
+        }),
+      ],
+    });
+
+    const rows = await providerRunnerRowsFor({workspaceId, provisionerId});
+    expect(rows[0]?.reservationId).toBe(boundReservationId);
+  });
+
+  it('adopts a reported reservation id when the runner carries none', async () => {
+    const reservationId = await createReservation(1);
+    await db()
+      .insert(providerRunners)
+      .values({
+        workspaceId,
+        provisionerId,
+        providerRunnerId: 'provisioned-runner-1',
+        state: 'starting',
+        reportedAt: new Date('2025-01-01T00:00:00.000Z'),
+      });
+
+    await reportRunnerInstances({
+      scope: 'workspace',
+      workspaceId,
+      provisionerId,
+      events: [
+        event({
+          providerRunnerId: 'provisioned-runner-1',
+          reservationId,
+          state: 'running',
+          reportedAt: new Date('2025-01-01T00:01:00.000Z'),
+        }),
+      ],
+    });
+
+    const rows = await providerRunnerRowsFor({workspaceId, provisionerId});
+    expect(rows[0]?.reservationId).toBe(reservationId);
+  });
+
   it('does not let equal-timestamp lower-priority reports flip terminal state', async () => {
     const reservationId = await createReservation(1);
     const reportedAt = new Date('2025-01-01T00:00:00.000Z');
@@ -2046,6 +2181,102 @@ describe('reconcileRunnerInstances', () => {
   }
 });
 
+describe('countStaleEnrolledRunnerInstances', () => {
+  let provisionerId: string;
+
+  beforeEach(() => {
+    provisionerId = crypto.randomUUID();
+  });
+
+  function staleAt(): Date {
+    return new Date(Date.now() - 120_000);
+  }
+
+  it('counts only stale reported runners with a live control session and no assignment', async () => {
+    const baseline = await countStaleEnrolledRunnerInstances({graceSeconds: 60});
+    const stale = await createRunner({updatedAt: staleAt()});
+    await createControlSession(stale);
+
+    const fresh = await createRunner({updatedAt: new Date()});
+    await createControlSession(fresh);
+
+    const recentlyMutated = await createRunner({reportedAt: staleAt(), updatedAt: new Date()});
+    await createControlSession(recentlyMutated);
+
+    const freshReport = await createRunner({reportedAt: new Date(), updatedAt: staleAt()});
+    await createControlSession(freshReport);
+
+    const assigned = await createRunner({workspaceId: crypto.randomUUID(), updatedAt: staleAt()});
+    await createControlSession(assigned);
+
+    const activated = await createRunner({
+      runnerSessionId: crypto.randomUUID(),
+      updatedAt: staleAt(),
+    });
+    await createControlSession(activated);
+
+    await createRunner({updatedAt: staleAt()});
+    const expiredControlSession = await createRunner({updatedAt: staleAt()});
+    await createControlSession(expiredControlSession, {
+      expiresAt: new Date(Date.now() - 1_000),
+    });
+    const closedControlSession = await createRunner({updatedAt: staleAt()});
+    await createControlSession(closedControlSession, {
+      closedAt: new Date(),
+    });
+    const stopped = await createRunner({state: 'stopped', updatedAt: staleAt()});
+    await createControlSession(stopped);
+
+    const count = await countStaleEnrolledRunnerInstances({graceSeconds: 60});
+
+    expect(count - baseline).toBe(2);
+  });
+
+  async function createRunner(params: {
+    state?: 'running' | 'stopped';
+    workspaceId?: string | null;
+    runnerSessionId?: string | null;
+    reportedAt?: Date;
+    updatedAt?: Date;
+  }) {
+    const reportedAt = params.reportedAt ?? params.updatedAt ?? new Date();
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        workspaceId: params.workspaceId ?? null,
+        provisionerId,
+        providerRunnerId: crypto.randomUUID(),
+        templateKey: 'linux',
+        labels: ['linux'],
+        state: params.state ?? 'running',
+        runnerSessionId: params.runnerSessionId ?? null,
+        providerKind: 'docker',
+        reportedAt,
+        updatedAt: params.updatedAt ?? reportedAt,
+      })
+      .returning({id: providerRunners.id});
+    if (!runner) throw new Error('Expected runner instance');
+    return runner.id;
+  }
+
+  async function createControlSession(
+    runnerInstanceId: string,
+    params: {expiresAt?: Date; closedAt?: Date} = {},
+  ) {
+    await db()
+      .insert(runnerControlSessions)
+      .values({
+        runnerInstanceId,
+        provisionerId,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: params.expiresAt ?? new Date(Date.now() + 60_000),
+        closedAt: params.closedAt,
+        closeReason: params.closedAt ? 'test' : null,
+      });
+  }
+});
+
 describe('runner instance provider attachment', () => {
   it('updates the pre-created runner instance before its first lifecycle report', async () => {
     const provisionerId = crypto.randomUUID();

package/src/db/runner-instances.ts

@@ -29,6 +29,7 @@ import {releaseReservationUnits} from './reservations.js';
 import {ephemeralRegistrationTokens} from './schema/ephemeral-registration-tokens.js';
 import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {reservations} from './schema/reservations.js';
+import {runnerControlSessions} from './schema/runner-control-sessions.js';
 import {providerRunners, toRunnerInstance} from './schema/runner-instances.js';
 import {runnerSessions} from './schema/runner-sessions.js';
 import {runningJobExecutions} from './schema/running-job-executions.js';
@@ -190,7 +191,7 @@ export async function reportRunnerInstances(params: ReportRunnerInstancesParams)
         target: [providerRunners.provisionerId, providerRunners.providerRunnerId],
         targetWhere: isNotNull(providerRunners.providerRunnerId),
         set: {
-          reservationId: sql`CASE WHEN ${providerRunnerProjectionUpdateCondition()} THEN coalesce(excluded.reservation_id, ${providerRunners.reservationId}) ELSE ${providerRunners.reservationId} END`,
+          reservationId: sql`CASE WHEN ${providerRunnerProjectionUpdateCondition()} THEN coalesce(${providerRunners.reservationId}, excluded.reservation_id) ELSE ${providerRunners.reservationId} END`,
           templateKey: sql`CASE WHEN ${providerRunnerProjectionUpdateCondition()} THEN coalesce(excluded.template_key, ${providerRunners.templateKey}) ELSE ${providerRunners.templateKey} END`,
           labels: sql`CASE WHEN ${providerRunnerProjectionUpdateCondition()} THEN excluded.labels ELSE ${providerRunners.labels} END`,
           state: sql`CASE WHEN ${providerRunnerProjectionUpdateCondition()} THEN excluded.state ELSE ${providerRunners.state} END`,
@@ -247,6 +248,36 @@ export async function listActiveRunnerInstanceCountsByTemplateTx(
   );
 }
 
+export async function countStaleEnrolledRunnerInstances(params: {
+  graceSeconds: number;
+}): Promise<number> {
+  const [row] = await db()
+    .select({count: sql<number>`count(*)::int`})
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.state, 'running'),
+        isNull(providerRunners.workspaceId),
+        isNull(providerRunners.runnerSessionId),
+        lt(providerRunners.reportedAt, staleRunnerInstanceCutoff(params.graceSeconds)),
+        exists(
+          db()
+            .select({id: runnerControlSessions.id})
+            .from(runnerControlSessions)
+            .where(
+              and(
+                eq(runnerControlSessions.runnerInstanceId, providerRunners.id),
+                isNull(runnerControlSessions.closedAt),
+                gt(runnerControlSessions.expiresAt, sql`now()`),
+              ),
+            ),
+        ),
+      ),
+    );
+
+  return row?.count ?? 0;
+}
+
 export async function listActiveRunnerInstances(params: {
   workspaceId: string;
   windowSeconds: number;

package/src/db/runner-sessions.ts

@@ -59,33 +59,44 @@ 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, token.provisionerId))
+      .where(eq(provisionerTokens.id, runner.provisionerId))
       .limit(1);
     const labels = sanitizeRunnerLabelsOrThrow(params.labels, {
       scope: provisioner?.scope ?? 'workspace',
@@ -95,13 +106,13 @@ export async function createRunnerSessionConsumingActivationToken(params: {
     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,
+        runnerInstanceId: runner.runnerInstanceId,
+        provisionerId: runner.provisionerId,
+        providerRunnerId: runner.providerRunnerId,
         labels,
         toolCapabilities: params.toolCapabilities ?? null,
         toolCapabilitiesReportedAt: params.toolCapabilities ? sql`now()` : null,
@@ -113,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),
         ),
       );
@@ -128,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/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],
   );
 }