@shipfox/api-runners 4.0.0 → 6.0.0

package/src/core/rate-limit.ts

@@ -1,5 +1,4 @@
-import {createHmac} from 'node:crypto';
-import {config as authConfig} from '@shipfox/api-auth/config';
+import {rateLimitIdentifierKey} from '@shipfox/node-auth-root-key';
 import {
   checkRateLimit,
   hashRateLimitIdentifier,
@@ -56,7 +55,6 @@ export class RunnersRateLimitUnavailableError extends RateLimitUnavailableError<
 }
 
 const IDENTIFIER_HASH_DOMAIN = 'shipfox.runners.rate-limit.identifier.v1';
-const IDENTIFIER_SECRET_DERIVATION_DOMAIN = 'shipfox.runners.rate-limit.identifier-secret.v1';
 
 export function hashRunnersRateLimitIdentifier(params: {
   action: RunnersRateLimitAction;
@@ -110,12 +108,6 @@ export async function checkRunnersRateLimit(params: CheckRunnersRateLimitParams)
   }
 }
 
-function effectiveIdentifierSecret(): Buffer | string {
-  if (config.RATE_LIMIT_IDENTIFIER_SECRET) {
-    return config.RATE_LIMIT_IDENTIFIER_SECRET;
-  }
-
-  return createHmac('sha256', authConfig.AUTH_JWT_SECRET)
-    .update(IDENTIFIER_SECRET_DERIVATION_DOMAIN)
-    .digest();
+function effectiveIdentifierSecret(): Uint8Array {
+  return rateLimitIdentifierKey();
 }

package/src/core/runner-activation.ts

@@ -0,0 +1,72 @@
+import {extractDisplayPrefix, generateOpaqueToken, hashOpaqueToken} from '@shipfox/node-tokens';
+import {and, eq, isNull, sql} from 'drizzle-orm';
+import {db} from '#db/db.js';
+import {runnerActivationTokens} from '#db/schema/runner-activation-tokens.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+
+export async function issueRunnerActivationToken(params: {
+  runnerInstanceId: string;
+  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;
+
+    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: {
+  runnerInstanceId: string;
+  provisionerId: string;
+}) {
+  const [runner] = await db()
+    .select({
+      workspaceId: providerRunners.workspaceId,
+      runnerSessionId: providerRunners.runnerSessionId,
+    })
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.id, params.runnerInstanceId),
+        eq(providerRunners.provisionerId, params.provisionerId),
+      ),
+    )
+    .limit(1);
+  return runner?.workspaceId && !runner.runnerSessionId ? runner : null;
+}

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

@@ -0,0 +1,190 @@
+import type {RunnerToolCapabilitiesDto} from '@shipfox/api-runners-dto';
+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 {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';
+
+export class RunnerBootstrapTokenInvalidError extends Error {
+  constructor() {
+    super('Runner bootstrap token is invalid, expired, or has already been used');
+    this.name = 'RunnerBootstrapTokenInvalidError';
+  }
+}
+
+export class RunnerControlSessionInvalidError extends Error {
+  constructor() {
+    super('Runner control session is invalid, expired, or closed');
+    this.name = 'RunnerControlSessionInvalidError';
+  }
+}
+
+export async function createRunnerInstancesWithBootstrapTokens(params: {
+  provisionerId: string;
+  providerKind?: string | null;
+  runnerInstances: Array<{templateKey?: string | null}>;
+  ttlSeconds: number;
+}): Promise<Array<{runnerInstanceId: string; bootstrapToken: string}>> {
+  const expiresAt = new Date(Date.now() + params.ttlSeconds * 1000);
+  return await db().transaction(async (tx) => {
+    const instances = await tx
+      .insert(providerRunners)
+      .values(
+        params.runnerInstances.map((runner) => ({
+          provisionerId: params.provisionerId,
+          providerKind: params.providerKind ?? null,
+          templateKey: runner.templateKey ?? null,
+          state: 'starting' as const,
+          labels: [],
+          reportedAt: new Date(),
+        })),
+      )
+      .returning({id: providerRunners.id});
+    const results = instances.map((instance) => ({
+      runnerInstanceId: instance.id,
+      bootstrapToken: generateOpaqueToken('runnerBootstrapToken'),
+    }));
+    await tx.insert(runnerBootstrapTokens).values(
+      results.map((result) => ({
+        runnerInstanceId: result.runnerInstanceId,
+        provisionerId: params.provisionerId,
+        hashedToken: hashOpaqueToken(result.bootstrapToken),
+        prefix: extractDisplayPrefix(result.bootstrapToken),
+        expiresAt,
+      })),
+    );
+    return results;
+  });
+}
+
+export async function exchangeRunnerBootstrapToken(params: {
+  rawToken: string;
+  ttlSeconds: number;
+}): Promise<{runnerInstanceId: string; controlSessionToken: string; expiresAt: Date}> {
+  const controlSessionToken = generateOpaqueToken('runnerControlSession');
+  const expiresAt = new Date(Date.now() + params.ttlSeconds * 1000);
+  return await db().transaction(async (tx) => {
+    const [bootstrap] = await tx
+      .update(runnerBootstrapTokens)
+      .set({consumedAt: sql`now()`})
+      .where(
+        and(
+          eq(runnerBootstrapTokens.hashedToken, hashOpaqueToken(params.rawToken)),
+          isNull(runnerBootstrapTokens.consumedAt),
+          isNull(runnerBootstrapTokens.revokedAt),
+          gt(runnerBootstrapTokens.expiresAt, sql`now()`),
+        ),
+      )
+      .returning();
+    if (!bootstrap) throw new RunnerBootstrapTokenInvalidError();
+    const [session] = await tx
+      .insert(runnerControlSessions)
+      .values({
+        runnerInstanceId: bootstrap.runnerInstanceId,
+        provisionerId: bootstrap.provisionerId,
+        hashedToken: hashOpaqueToken(controlSessionToken),
+        prefix: extractDisplayPrefix(controlSessionToken),
+        expiresAt,
+      })
+      .returning({id: runnerControlSessions.id});
+    if (!session) throw new Error('Runner control session insert returned no row');
+    return {runnerInstanceId: bootstrap.runnerInstanceId, controlSessionToken, expiresAt};
+  });
+}
+
+export async function resolveRunnerControlSession(rawToken: string) {
+  const [session] = await db()
+    .select({
+      id: runnerControlSessions.id,
+      runnerInstanceId: runnerControlSessions.runnerInstanceId,
+      provisionerId: runnerControlSessions.provisionerId,
+    })
+    .from(runnerControlSessions)
+    .leftJoin(provisionerTokens, eq(provisionerTokens.id, runnerControlSessions.provisionerId))
+    .where(
+      and(
+        eq(runnerControlSessions.hashedToken, hashOpaqueToken(rawToken)),
+        isNull(runnerControlSessions.closedAt),
+        or(isNull(provisionerTokens.id), isNull(provisionerTokens.revokedAt)),
+        gt(runnerControlSessions.expiresAt, sql`now()`),
+      ),
+    )
+    .limit(1);
+  return session;
+}
+
+export async function enrollRunnerControlSession(params: {
+  runnerInstanceId: string;
+  provisionerId: string;
+  labels: string[];
+  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);
+}
+
+export async function attachRunnerControlProviderId(params: {
+  runnerInstanceId: string;
+  provisionerId: string;
+  providerRunnerId: string;
+}): Promise<boolean> {
+  const rows = await db()
+    .update(providerRunners)
+    .set({providerRunnerId: params.providerRunnerId, updatedAt: sql`now()`})
+    .where(
+      and(
+        eq(providerRunners.id, params.runnerInstanceId),
+        eq(providerRunners.provisionerId, params.provisionerId),
+        isNull(providerRunners.providerRunnerId),
+        notInArray(providerRunners.state, [...terminalStates]),
+      ),
+    )
+    .returning({id: providerRunners.id});
+  return rows.length === 1;
+}
+
+export async function touchRunnerControlSession(runnerInstanceId: string, provisionerId: string) {
+  await db()
+    .update(runnerControlSessions)
+    .set({lastSeenAt: sql`now()`})
+    .where(
+      and(
+        eq(runnerControlSessions.runnerInstanceId, runnerInstanceId),
+        eq(runnerControlSessions.provisionerId, provisionerId),
+        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/core/{provisioned-runners.test.ts → runner-instances.test.ts}

@@ -1,51 +1,51 @@
-import {reconcileProvisionedRunnersFromDbResult} from './provisioned-runners.js';
+import {reconcileRunnerInstancesFromDbResult} from './runner-instances.js';
 
-describe('reconcileProvisionedRunnersFromDbResult', () => {
+describe('reconcileRunnerInstancesFromDbResult', () => {
   it('keeps active provisioned runners', () => {
-    const result = reconcileProvisionedRunnersFromDbResult({
-      observedProvisionedRunnerIds: ['provisioned-runner-1'],
-      observedRows: [provisionedRunner({provisionedRunnerId: 'provisioned-runner-1'})],
-      boundJobExecutionsByProvisionedRunnerId: new Map(),
+    const result = reconcileRunnerInstancesFromDbResult({
+      observedRunnerInstanceIds: ['provisioned-runner-1'],
+      observedRows: [providerRunner({providerRunnerId: 'provisioned-runner-1'})],
+      boundJobExecutionsByRunnerInstanceId: new Map(),
     });
 
     expect(result[0]?.desiredIntent).toBe('keep');
   });
 
   it('terminates terminal provisioned runners', () => {
-    const result = reconcileProvisionedRunnersFromDbResult({
-      observedProvisionedRunnerIds: ['provisioned-runner-1'],
+    const result = reconcileRunnerInstancesFromDbResult({
+      observedRunnerInstanceIds: ['provisioned-runner-1'],
       observedRows: [
-        provisionedRunner({provisionedRunnerId: 'provisioned-runner-1', state: 'terminated'}),
+        providerRunner({providerRunnerId: 'provisioned-runner-1', state: 'terminated'}),
       ],
-      boundJobExecutionsByProvisionedRunnerId: new Map(),
+      boundJobExecutionsByRunnerInstanceId: new Map(),
     });
 
     expect(result[0]?.desiredIntent).toBe('terminate');
   });
 
   it('keeps orphan observed provisioned runners', () => {
-    const result = reconcileProvisionedRunnersFromDbResult({
-      observedProvisionedRunnerIds: ['provisioned-runner-1'],
+    const result = reconcileRunnerInstancesFromDbResult({
+      observedRunnerInstanceIds: ['provisioned-runner-1'],
       observedRows: [],
-      boundJobExecutionsByProvisionedRunnerId: new Map(),
+      boundJobExecutionsByRunnerInstanceId: new Map(),
     });
 
     expect(result[0]).toMatchObject({
-      provisionedRunnerId: 'provisioned-runner-1',
+      providerRunnerId: 'provisioned-runner-1',
       state: null,
       desiredIntent: 'keep',
     });
   });
 
   it('terminates active provisioned runners with a cancelled latest bound job', () => {
-    const result = reconcileProvisionedRunnersFromDbResult({
-      observedProvisionedRunnerIds: ['provisioned-runner-1'],
-      observedRows: [provisionedRunner({provisionedRunnerId: 'provisioned-runner-1'})],
-      boundJobExecutionsByProvisionedRunnerId: new Map([
+    const result = reconcileRunnerInstancesFromDbResult({
+      observedRunnerInstanceIds: ['provisioned-runner-1'],
+      observedRows: [providerRunner({providerRunnerId: 'provisioned-runner-1'})],
+      boundJobExecutionsByRunnerInstanceId: new Map([
         [
           'provisioned-runner-1',
           boundJobExecution({
-            provisionedRunnerId: 'provisioned-runner-1',
+            providerRunnerId: 'provisioned-runner-1',
             cancellationRequestedAt: new Date('2025-01-01T00:01:00.000Z'),
           }),
         ],
@@ -56,15 +56,15 @@ describe('reconcileProvisionedRunnersFromDbResult', () => {
   });
 });
 
-function provisionedRunner(params: {
-  provisionedRunnerId: string;
+function providerRunner(params: {
+  providerRunnerId: string;
   state?: 'starting' | 'running' | 'stopping' | 'stopped' | 'failed' | 'terminated';
 }) {
   return {
     id: crypto.randomUUID(),
     workspaceId: crypto.randomUUID(),
     provisionerId: crypto.randomUUID(),
-    provisionedRunnerId: params.provisionedRunnerId,
+    providerRunnerId: params.providerRunnerId,
     reservationId: null,
     templateKey: 'linux',
     labels: ['linux'],
@@ -85,7 +85,7 @@ function provisionedRunner(params: {
 }
 
 function boundJobExecution(params: {
-  provisionedRunnerId: string;
+  providerRunnerId: string;
   cancellationRequestedAt?: Date | null;
 }) {
   return {
@@ -93,7 +93,7 @@ function boundJobExecution(params: {
     workflowRunAttemptId: crypto.randomUUID(),
     jobId: crypto.randomUUID(),
     jobExecutionId: crypto.randomUUID(),
-    provisionedRunnerId: params.provisionedRunnerId,
+    providerRunnerId: params.providerRunnerId,
     startedAt: new Date('2025-01-01T00:00:00.000Z'),
     lastHeartbeatAt: new Date('2025-01-01T00:00:00.000Z'),
     cancellationRequestedAt: params.cancellationRequestedAt ?? null,