@shipfox/api-runners 12.1.0 → 12.3.0

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

@@ -4,16 +4,18 @@ import {db} from '#db/db.js';
 import {createRunnerSessionConsumingEphemeralToken} from '#db/ephemeral-registration-tokens.js';
 import {
   attachRunnerInstanceProviderId,
+  countStaleEnrolledRunnerInstances,
   listActiveRunnerInstanceCountsByTemplateTx,
   listActiveRunnerInstances,
   listProvisionerTerminateIntentRowsTx,
   listProvisionerTerminateIntents,
   reapStaleRunnerInstances,
   reconcileRunnerInstances,
-  reportRunnerInstances,
+  reportRunnerInstances as reportRunnerInstancesDb,
 } 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';
@@ -25,6 +27,12 @@ import {
   runnerSessionFactory,
 } from '#test/index.js';
 
+type ReportRunnerInstancesTestParams = Parameters<typeof reportRunnerInstancesDb>[0];
+
+function reportRunnerInstances(params: ReportRunnerInstancesTestParams) {
+  return reportRunnerInstancesDb(params);
+}
+
 function providerRunnerRowsFor(params: {workspaceId: string; provisionerId: string}) {
   return db()
     .select()
@@ -93,6 +101,7 @@ describe('reportRunnerInstances', () => {
     const reportedAt = new Date();
 
     const report = await reportRunnerInstances({
+      scope: 'installation',
       workspaceId: null,
       provisionerId,
       events: [event({providerRunnerId: 'installation-runner', reportedAt})],
@@ -118,6 +127,7 @@ describe('reportRunnerInstances', () => {
     const reportedAt = new Date();
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -140,6 +150,7 @@ describe('reportRunnerInstances', () => {
     const reportedAt = new Date();
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -160,6 +171,7 @@ describe('reportRunnerInstances', () => {
   it('accepts delayed events that move the lifecycle forward', async () => {
     const newest = new Date();
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -168,6 +180,7 @@ describe('reportRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -188,6 +201,7 @@ describe('reportRunnerInstances', () => {
   it('rejects older out-of-order events in the same lifecycle state', async () => {
     const newest = new Date();
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -201,6 +215,7 @@ describe('reportRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -218,10 +233,144 @@ 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');
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -235,6 +384,7 @@ describe('reportRunnerInstances', () => {
     });
 
     const terminal = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -247,6 +397,7 @@ describe('reportRunnerInstances', () => {
       ],
     });
     const revived = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -270,6 +421,7 @@ describe('reportRunnerInstances', () => {
 
   it('clamps future reported times so they do not pin provisioned runner state', async () => {
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -282,6 +434,7 @@ describe('reportRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -305,6 +458,7 @@ describe('reportRunnerInstances', () => {
     const terminatedAt = new Date('2025-01-01T00:03:00.000Z');
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -346,6 +500,7 @@ describe('reportRunnerInstances', () => {
     const terminatedAt = new Date('2025-01-01T00:03:00.000Z');
     const startedAt = new Date('2025-01-01T00:00:00.000Z');
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -358,6 +513,7 @@ describe('reportRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -378,6 +534,7 @@ describe('reportRunnerInstances', () => {
 
   it('uses server update time for active provisioned runner windows', async () => {
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -401,6 +558,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(2);
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', reservationId, state: 'failed'})],
@@ -429,6 +587,7 @@ describe('reportRunnerInstances', () => {
     if (!runner.providerRunnerId) throw new Error('Runner instance provider id missing');
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: runner.providerRunnerId, state: 'failed'})],
@@ -479,6 +638,7 @@ describe('reportRunnerInstances', () => {
     });
 
     const result = await reportRunnerInstances({
+      scope: 'installation',
       workspaceId: null,
       provisionerId: installationProvisioner.id,
       events: [event({providerRunnerId: 'installation-pre-enrollment-runner', state: 'failed'})],
@@ -521,6 +681,7 @@ describe('reportRunnerInstances', () => {
     });
 
     const result = await reportRunnerInstances({
+      scope: 'installation',
       workspaceId: null,
       provisionerId: installationProvisioner.id,
       events: [
@@ -552,6 +713,7 @@ describe('reportRunnerInstances', () => {
     });
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -581,11 +743,13 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(2);
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', reservationId, state: 'failed'})],
     });
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', reservationId, state: 'failed'})],
@@ -600,6 +764,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(2);
     const failedAt = new Date();
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -613,6 +778,7 @@ describe('reportRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -636,6 +802,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(2);
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -656,6 +823,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(3);
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -673,6 +841,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(1);
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', reservationId, state: 'failed'})],
@@ -695,6 +864,7 @@ describe('reportRunnerInstances', () => {
     if (!reservation) throw new Error('Expected reservation');
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -715,6 +885,7 @@ describe('reportRunnerInstances', () => {
     const reservationId = await createReservation(1);
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -750,6 +921,7 @@ describe('reportRunnerInstances', () => {
     });
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -776,6 +948,7 @@ describe('reportRunnerInstances', () => {
     const reportedAt = new Date('2025-01-01T00:00:00.000Z');
 
     const result = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -819,11 +992,13 @@ describe('reportRunnerInstances', () => {
     });
 
     const first = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', state: 'terminated'})],
     });
     const second = await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [event({providerRunnerId: 'provisioned-runner-1', state: 'terminated'})],
@@ -1365,6 +1540,7 @@ describe('reapStaleRunnerInstances', () => {
 
     await lockHolderReady.promise;
     const report = reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -1845,6 +2021,7 @@ describe('reconcileRunnerInstances', () => {
     });
 
     await reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -1877,6 +2054,7 @@ describe('reconcileRunnerInstances', () => {
 
     await lockHolderReady.promise;
     const report = reportRunnerInstances({
+      scope: 'workspace',
       workspaceId,
       provisionerId,
       events: [
@@ -2003,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();
@@ -2026,6 +2300,7 @@ describe('runner instance provider attachment', () => {
       providerRunnerId,
     });
     await reportRunnerInstances({
+      scope: 'installation',
       workspaceId: null,
       provisionerId,
       events: [

package/src/db/runner-instances.ts

@@ -1,5 +1,4 @@
 import {logger} from '@shipfox/node-opentelemetry';
-import {canonicalizeLabels} from '@shipfox/runner-labels';
 import {
   and,
   asc,
@@ -19,6 +18,7 @@ import {
 } from 'drizzle-orm';
 import {alias} from 'drizzle-orm/pg-core';
 import type {RunnerInstance, RunnerInstanceState} from '#core/entities/runner-instance.js';
+import {sanitizeRunnerLabels} from '#core/runner-labels.js';
 import type {Tx} from './db.js';
 import {db} from './db.js';
 import {
@@ -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';
@@ -74,6 +75,7 @@ export interface RunnerInstanceReportEvent {
 }
 
 export interface ReportRunnerInstancesParams {
+  scope: 'installation' | 'workspace';
   workspaceId: string | null;
   provisionerId: string;
   events: RunnerInstanceReportEvent[];
@@ -165,7 +167,11 @@ export async function reportRunnerInstances(params: ReportRunnerInstancesParams)
       providerRunnerId: event.providerRunnerId,
       reservationId: event.reservationId,
       templateKey: event.templateKey,
-      labels: [...canonicalizeLabels(event.labels)],
+      labels: sanitizeRunnerLabels(event.labels, {
+        scope: params.scope,
+        logLevel: 'debug',
+        source: 'runner instance report',
+      }),
       state: event.state,
       reason: event.reason,
       runnerSessionId: event.runnerSessionId,
@@ -185,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`,
@@ -242,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;