@shipfox/api-runners 9.3.0 → 10.0.0

package/drizzle/meta/_journal.json

@@ -8,6 +8,13 @@
       "when": 1777238442655,
       "tag": "0000_initial",
       "breakpoints": true
+    },
+    {
+      "idx": 1,
+      "version": "7",
+      "when": 1785161220762,
+      "tag": "0001_intended_reservation",
+      "breakpoints": true
     }
   ]
 }

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@shipfox/api-runners",
   "license": "MIT",
-  "version": "9.3.0",
+  "version": "10.0.0",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/ShipfoxHQ/shipfox.git",
@@ -21,24 +21,24 @@
     "@temporalio/workflow": "1.18.1",
     "drizzle-orm": "^0.45.2",
     "zod": "^4.4.3",
-    "@shipfox/api-auth-context": "9.3.0",
-    "@shipfox/api-auth-dto": "9.3.0",
-    "@shipfox/api-runners-dto": "9.0.2",
-    "@shipfox/api-workflows-dto": "9.3.0",
+    "@shipfox/api-auth-context": "10.0.0",
+    "@shipfox/api-runners-dto": "10.0.0",
     "@shipfox/config": "1.2.4",
-    "@shipfox/inter-module": "0.2.2",
+    "@shipfox/api-workflows-dto": "10.0.0",
+    "@shipfox/api-auth-dto": "10.0.0",
     "@shipfox/node-drizzle": "0.3.4",
     "@shipfox/node-error-monitoring": "0.3.0",
     "@shipfox/node-auth-root-key": "0.2.3",
-    "@shipfox/node-fastify": "0.3.4",
-    "@shipfox/node-module": "1.0.3",
+    "@shipfox/node-module": "1.0.4",
     "@shipfox/node-opentelemetry": "0.6.3",
+    "@shipfox/node-fastify": "0.4.0",
+    "@shipfox/node-rate-limit": "0.3.2",
     "@shipfox/node-outbox": "0.2.6",
     "@shipfox/node-postgres": "0.4.4",
-    "@shipfox/node-rate-limit": "0.3.2",
     "@shipfox/node-temporal": "0.4.4",
     "@shipfox/node-tokens": "0.3.2",
-    "@shipfox/runner-labels": "0.1.3"
+    "@shipfox/runner-labels": "0.1.3",
+    "@shipfox/inter-module": "0.2.2"
   },
   "imports": {
     "#*": "./dist/*"

package/src/config.test.ts

@@ -27,6 +27,17 @@ describe('EPHEMERAL_REGISTRATION_TOKEN_TTL_SECONDS validation', () => {
     expect(config.EPHEMERAL_REGISTRATION_TOKEN_TTL_SECONDS).toBe(3600);
   });
 });
+describe('runner assignment polling defaults', () => {
+  afterEach(() => {
+    vi.unstubAllEnvs();
+    vi.resetModules();
+  });
+  it('uses a 250 millisecond interval between assignment reads', async () => {
+    vi.resetModules();
+    const {config} = await import('#config.js');
+    expect(config.RUNNER_ASSIGNMENT_POLL_INTERVAL_MS).toBe(250);
+  });
+});
 describe('REGISTRATION_TOKEN_BATCH_MAX validation', () => {
   afterEach(() => {
     vi.unstubAllEnvs();

package/src/config.ts

@@ -24,7 +24,7 @@ export const config = createConfig({
   }),
   RUNNER_ASSIGNMENT_POLL_INTERVAL_MS: num({
     desc: 'Delay between durable assignment reads while a runner-control session waits for an assignment, in milliseconds.',
-    default: 1000,
+    default: 250,
   }),
   EPHEMERAL_REGISTRATION_TOKEN_TTL_SECONDS: num({
     desc: `Lifetime of a provisioner-minted runner registration token, in seconds. The token can be exchanged once by a runner before this time passes. Set this between 1 and ${EPHEMERAL_REGISTRATION_TOKEN_TTL_HARD_MAX_SECONDS}.`,

package/src/core/runner-activation.ts

@@ -1,5 +1,6 @@
 import {extractDisplayPrefix, generateOpaqueToken, hashOpaqueToken} from '@shipfox/node-tokens';
 import {and, eq, isNull, sql} from 'drizzle-orm';
+import type {Tx} from '#db/db.js';
 import {db} from '#db/db.js';
 import {runnerActivationTokens} from '#db/schema/runner-activation-tokens.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
@@ -9,46 +10,55 @@ export async function issueRunnerActivationToken(params: {
   provisionerId: string;
   ttlSeconds: number;
 }): Promise<string | null> {
-  return await db().transaction(async (tx) => {
-    await tx.execute(
-      sql`select pg_advisory_xact_lock(hashtext(${`runners_activation:${params.runnerInstanceId}`}))`,
-    );
-    const [runner] = await tx
-      .select({
-        workspaceId: providerRunners.workspaceId,
-        runnerSessionId: providerRunners.runnerSessionId,
-        state: providerRunners.state,
-      })
-      .from(providerRunners)
-      .where(
-        and(
-          eq(providerRunners.id, params.runnerInstanceId),
-          eq(providerRunners.provisionerId, params.provisionerId),
-        ),
-      )
-      .limit(1)
-      .for('update');
-    if (!runner?.workspaceId || runner.runnerSessionId || runner.state !== 'running') return null;
+  return await db().transaction(async (tx) => issueRunnerActivationTokenTx(tx, params));
+}
+
+export async function issueRunnerActivationTokenTx(
+  tx: Tx,
+  params: {
+    runnerInstanceId: string;
+    provisionerId: string;
+    ttlSeconds: number;
+  },
+): Promise<string | null> {
+  await tx.execute(
+    sql`select pg_advisory_xact_lock(hashtext(${`runners_activation:${params.runnerInstanceId}`}))`,
+  );
+  const [runner] = await tx
+    .select({
+      workspaceId: providerRunners.workspaceId,
+      runnerSessionId: providerRunners.runnerSessionId,
+      state: providerRunners.state,
+    })
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.id, params.runnerInstanceId),
+        eq(providerRunners.provisionerId, params.provisionerId),
+      ),
+    )
+    .limit(1)
+    .for('update');
+  if (!runner?.workspaceId || runner.runnerSessionId || runner.state !== 'running') return null;
 
-    await tx
-      .update(runnerActivationTokens)
-      .set({revokedAt: sql`now()`})
-      .where(
-        and(
-          eq(runnerActivationTokens.runnerInstanceId, params.runnerInstanceId),
-          isNull(runnerActivationTokens.consumedAt),
-          isNull(runnerActivationTokens.revokedAt),
-        ),
-      );
-    const rawToken = generateOpaqueToken('runnerActivationToken');
-    await tx.insert(runnerActivationTokens).values({
-      runnerInstanceId: params.runnerInstanceId,
-      hashedToken: hashOpaqueToken(rawToken),
-      prefix: extractDisplayPrefix(rawToken),
-      expiresAt: new Date(Date.now() + params.ttlSeconds * 1000),
-    });
-    return rawToken;
+  await tx
+    .update(runnerActivationTokens)
+    .set({revokedAt: sql`now()`})
+    .where(
+      and(
+        eq(runnerActivationTokens.runnerInstanceId, params.runnerInstanceId),
+        isNull(runnerActivationTokens.consumedAt),
+        isNull(runnerActivationTokens.revokedAt),
+      ),
+    );
+  const rawToken = generateOpaqueToken('runnerActivationToken');
+  await tx.insert(runnerActivationTokens).values({
+    runnerInstanceId: params.runnerInstanceId,
+    hashedToken: hashOpaqueToken(rawToken),
+    prefix: extractDisplayPrefix(rawToken),
+    expiresAt: new Date(Date.now() + params.ttlSeconds * 1000),
   });
+  return rawToken;
 }
 
 export async function getRunnerAssignment(params: {

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

@@ -1,12 +1,23 @@
 import type {RunnerToolCapabilitiesDto} from '@shipfox/api-runners-dto';
+import type {NodePgDatabase} from '@shipfox/node-drizzle';
+import {logger} from '@shipfox/node-opentelemetry';
 import {extractDisplayPrefix, generateOpaqueToken, hashOpaqueToken} from '@shipfox/node-tokens';
 import {canonicalizeLabels} from '@shipfox/runner-labels';
 import {and, eq, gt, isNull, notInArray, or, sql} from 'drizzle-orm';
-import {db} from '#db/db.js';
+import {config} from '#config.js';
+import {
+  ReservationExpiredError,
+  ReservationNotFoundError,
+  RunnerInstanceAlreadyAssignedError,
+  RunnerInstanceNotAssignableError,
+} from '#core/errors.js';
+import {db, type schema, type Tx} from '#db/db.js';
+import {assignRunnerInstancesTx} from '#db/runner-assignments.js';
 import {terminalStates} from '#db/runner-instances.js';
 import {provisionerTokens} from '#db/schema/provisioner-tokens.js';
 import {runnerBootstrapTokens, runnerControlSessions} from '#db/schema/runner-control-sessions.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
+import {issueRunnerActivationTokenTx} from './runner-activation.js';
 
 export class RunnerBootstrapTokenInvalidError extends Error {
   constructor() {
@@ -25,7 +36,10 @@ export class RunnerControlSessionInvalidError extends Error {
 export async function createRunnerInstancesWithBootstrapTokens(params: {
   provisionerId: string;
   providerKind?: string | null;
-  runnerInstances: Array<{templateKey?: string | null}>;
+  runnerInstances: Array<{
+    templateKey?: string | null;
+    reservationId?: string | null;
+  }>;
   ttlSeconds: number;
 }): Promise<Array<{runnerInstanceId: string; bootstrapToken: string}>> {
   const expiresAt = new Date(Date.now() + params.ttlSeconds * 1000);
@@ -37,6 +51,7 @@ export async function createRunnerInstancesWithBootstrapTokens(params: {
           provisionerId: params.provisionerId,
           providerKind: params.providerKind ?? null,
           templateKey: runner.templateKey ?? null,
+          intendedReservationId: runner.reservationId ?? null,
           state: 'starting' as const,
           labels: [],
           reportedAt: new Date(),
@@ -123,28 +138,87 @@ export async function enrollRunnerControlSession(params: {
   capabilities?: RunnerToolCapabilitiesDto | null;
   providerKind: string;
   protocolVersion: string;
-}): Promise<void> {
-  const updated = await db()
-    .update(providerRunners)
-    .set({
-      labels: [...canonicalizeLabels(params.labels)],
-      providerKind: params.providerKind,
-      protocolVersion: params.protocolVersion,
-      capabilities: params.capabilities ?? null,
-      state: 'running',
-      reportedAt: sql`now()`,
-      updatedAt: sql`now()`,
-    })
-    .where(
-      and(
-        eq(providerRunners.id, params.runnerInstanceId),
-        eq(providerRunners.provisionerId, params.provisionerId),
-        notInArray(providerRunners.state, [...terminalStates]),
-      ),
-    )
-    .returning({id: providerRunners.id});
-  if (!updated[0]) throw new RunnerControlSessionInvalidError();
-  await touchRunnerControlSession(params.runnerInstanceId, params.provisionerId);
+}): Promise<string | null> {
+  return await db().transaction(async (tx) => {
+    const [current] = await tx
+      .select({intendedReservationId: providerRunners.intendedReservationId})
+      .from(providerRunners)
+      .where(
+        and(
+          eq(providerRunners.id, params.runnerInstanceId),
+          eq(providerRunners.provisionerId, params.provisionerId),
+        ),
+      );
+    if (!current) throw new RunnerControlSessionInvalidError();
+    if (current.intendedReservationId)
+      await tx.execute(
+        sql`select pg_advisory_xact_lock(hashtext(${`runners_assignment:${params.provisionerId}:${current.intendedReservationId}`}))`,
+      );
+
+    const [updated] = await tx
+      .update(providerRunners)
+      .set({
+        labels: [...canonicalizeLabels(params.labels)],
+        providerKind: params.providerKind,
+        protocolVersion: params.protocolVersion,
+        capabilities: params.capabilities ?? null,
+        state: 'running',
+        reportedAt: sql`now()`,
+        updatedAt: sql`now()`,
+      })
+      .where(
+        and(
+          eq(providerRunners.id, params.runnerInstanceId),
+          eq(providerRunners.provisionerId, params.provisionerId),
+          notInArray(providerRunners.state, [...terminalStates]),
+        ),
+      )
+      .returning({
+        id: providerRunners.id,
+        intendedReservationId: providerRunners.intendedReservationId,
+      });
+    if (!updated) throw new RunnerControlSessionInvalidError();
+
+    await updateRunnerControlSessionLastSeen(tx, params.runnerInstanceId, params.provisionerId);
+
+    const intendedReservationId = updated.intendedReservationId;
+    if (!intendedReservationId) return null;
+
+    try {
+      return await tx.transaction(async (promotionTx) => {
+        await assignRunnerInstancesTx(promotionTx, {
+          provisionerId: params.provisionerId,
+          reservationId: intendedReservationId,
+          runnerInstanceIds: [params.runnerInstanceId],
+        });
+        return await issueRunnerActivationTokenTx(promotionTx, {
+          runnerInstanceId: params.runnerInstanceId,
+          provisionerId: params.provisionerId,
+          ttlSeconds: config.RUNNER_ACTIVATION_TOKEN_TTL_SECONDS,
+        });
+      });
+    } catch (error) {
+      const expectedFailure =
+        error instanceof ReservationExpiredError ||
+        error instanceof ReservationNotFoundError ||
+        error instanceof RunnerInstanceAlreadyAssignedError ||
+        error instanceof RunnerInstanceNotAssignableError;
+      const details = {
+        err: error,
+        runnerInstanceId: params.runnerInstanceId,
+        provisionerId: params.provisionerId,
+        reservationId: intendedReservationId,
+      };
+      if (expectedFailure)
+        logger().debug(details, 'Runner reservation could not be promoted during enrollment');
+      else
+        logger().error(
+          details,
+          'Unexpected failure promoting runner reservation during enrollment',
+        );
+      return null;
+    }
+  });
 }
 
 export async function attachRunnerControlProviderId(params: {
@@ -168,7 +242,24 @@ export async function attachRunnerControlProviderId(params: {
 }
 
 export async function touchRunnerControlSession(runnerInstanceId: string, provisionerId: string) {
+  await updateRunnerControlSessionLastSeen(db(), runnerInstanceId, provisionerId);
   await db()
+    .update(providerRunners)
+    .set({reportedAt: sql`now()`, updatedAt: sql`now()`})
+    .where(
+      and(
+        eq(providerRunners.id, runnerInstanceId),
+        eq(providerRunners.provisionerId, provisionerId),
+      ),
+    );
+}
+
+async function updateRunnerControlSessionLastSeen(
+  tx: Tx | NodePgDatabase<typeof schema>,
+  runnerInstanceId: string,
+  provisionerId: string,
+) {
+  await tx
     .update(runnerControlSessions)
     .set({lastSeenAt: sql`now()`})
     .where(
@@ -178,13 +269,4 @@ export async function touchRunnerControlSession(runnerInstanceId: string, provis
         isNull(runnerControlSessions.closedAt),
       ),
     );
-  await db()
-    .update(providerRunners)
-    .set({reportedAt: sql`now()`, updatedAt: sql`now()`})
-    .where(
-      and(
-        eq(providerRunners.id, runnerInstanceId),
-        eq(providerRunners.provisionerId, provisionerId),
-      ),
-    );
 }

package/src/db/index.ts

@@ -22,6 +22,7 @@ export {
   claimPendingJobExecution,
   enqueueJobExecution,
   expireStuckJobExecutions,
+  getWorkspaceJobCounts,
   isJobLeaseActive,
   listActiveRunningJobExecutions,
   listRunningJobExecutionsByRunnerInstanceTx,

package/src/db/job-executions.ts

@@ -47,6 +47,34 @@ export interface EnqueueJobExecutionParams {
   requiredLabels: string[];
 }
 
+export async function getWorkspaceJobCounts(params: {
+  workspaceIds: string[];
+}): Promise<Array<{workspaceId: string; queued: number; running: number}>> {
+  const [pendingRows, runningRows] = await Promise.all([
+    db()
+      .select({workspaceId: pendingJobExecutions.workspaceId, count: count()})
+      .from(pendingJobExecutions)
+      .where(inArray(pendingJobExecutions.workspaceId, params.workspaceIds))
+      .groupBy(pendingJobExecutions.workspaceId),
+    db()
+      .select({workspaceId: runningJobExecutions.workspaceId, count: count()})
+      .from(runningJobExecutions)
+      .where(inArray(runningJobExecutions.workspaceId, params.workspaceIds))
+      .groupBy(runningJobExecutions.workspaceId),
+  ]);
+
+  const queuedByWorkspace = new Map(pendingRows.map((row) => [row.workspaceId, Number(row.count)]));
+  const runningByWorkspace = new Map(
+    runningRows.map((row) => [row.workspaceId, Number(row.count)]),
+  );
+
+  return params.workspaceIds.map((workspaceId) => ({
+    workspaceId,
+    queued: queuedByWorkspace.get(workspaceId) ?? 0,
+    running: runningByWorkspace.get(workspaceId) ?? 0,
+  }));
+}
+
 // Idempotent while the job execution is still pending: a duplicate jobExecutionId already in
 // `runners_pending_jobs` is a no-op. Temporal retries the enqueue activity
 // at-least-once, so a unique-violation throw on a retry-after-lost-result

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

@@ -61,6 +61,25 @@ describe('assignRunnerInstances', () => {
     expect(results).toEqual([[runner.id], [runner.id]]);
   });
 
+  it('is idempotent after an assigned reservation is deleted', async () => {
+    const reservation = await createReservation();
+    const runner = await createEnrolledRunner();
+    await assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+    await db().delete(reservations).where(eq(reservations.id, reservation.id));
+
+    const assigned = await assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    expect(assigned).toEqual([runner.id]);
+  });
+
   it('rejects expired reservations', async () => {
     const reservation = await createReservation({expiresAt: new Date(Date.now() - 1_000)});
     const runner = await createEnrolledRunner();
@@ -120,6 +139,68 @@ describe('assignRunnerInstances', () => {
     await expect(assignment).rejects.toThrow(RunnerInstanceNotAssignableError);
   });
 
+  it('does not charge capacity for an already released reservation intent', async () => {
+    const reservation = await createReservation();
+    await db().insert(providerRunners).values({
+      provisionerId,
+      intendedReservationId: reservation.id,
+      reservationReleasedAt: new Date(),
+      providerRunnerId: crypto.randomUUID(),
+      state: 'failed',
+      reportedAt: new Date(),
+    });
+    const runner = await createEnrolledRunner();
+
+    await expect(
+      assignRunnerInstances({
+        provisionerId,
+        reservationId: reservation.id,
+        runnerInstanceIds: [runner.id],
+      }),
+    ).resolves.toEqual([runner.id]);
+  });
+
+  it('does not charge capacity for an already released reservation assignment', async () => {
+    const reservation = await createReservation();
+    await db().insert(providerRunners).values({
+      provisionerId,
+      reservationId: reservation.id,
+      reservationReleasedAt: new Date(),
+      providerRunnerId: crypto.randomUUID(),
+      state: 'failed',
+      reportedAt: new Date(),
+    });
+    const runner = await createEnrolledRunner();
+
+    await expect(
+      assignRunnerInstances({
+        provisionerId,
+        reservationId: reservation.id,
+        runnerInstanceIds: [runner.id],
+      }),
+    ).resolves.toEqual([runner.id]);
+  });
+
+  it('does not charge capacity for a terminal reservation intent before release bookkeeping', async () => {
+    const reservation = await createReservation();
+    await db().insert(providerRunners).values({
+      provisionerId,
+      intendedReservationId: reservation.id,
+      providerRunnerId: crypto.randomUUID(),
+      state: 'failed',
+      reportedAt: new Date(),
+    });
+    const runner = await createEnrolledRunner();
+
+    await expect(
+      assignRunnerInstances({
+        provisionerId,
+        reservationId: reservation.id,
+        runnerInstanceIds: [runner.id],
+      }),
+    ).resolves.toEqual([runner.id]);
+  });
+
   async function createReservation(
     overrides: Partial<{expiresAt: Date; requiredLabels: string[]}> = {},
   ) {