@shipfox/api-runners 9.2.0 → 10.0.0

package/src/index.ts

@@ -30,6 +30,7 @@ export {
   cancelRunnerJobs,
   type EnqueueJobExecutionParams,
   enqueueJobExecution,
+  getWorkspaceJobCounts,
   isJobLeaseActive,
   releaseJobExecution,
 } from '#db/index.js';

package/src/presentation/inter-module.ts

@@ -9,6 +9,7 @@ import {getEffectiveRunnerToolCapabilities} from '#core/runner-tool-capabilities
 import {
   cancelRunnerJobs,
   enqueueJobExecution,
+  getWorkspaceJobCounts,
   isJobLeaseActive,
   releaseJobExecution,
 } from '#db/job-executions.js';
@@ -38,6 +39,9 @@ export function createRunnersInterModulePresentation(): InterModulePresentation<
       const result = await getEffectiveRunnerToolCapabilities(input);
       return {capabilities: result.capabilities, reportFresh: result.reportFresh};
     },
+    getWorkspaceJobCounts: async ({workspaceIds}) => ({
+      counts: await getWorkspaceJobCounts({workspaceIds}),
+    }),
   });
 }
 

package/src/presentation/routes/poll-demand.test.ts

@@ -290,8 +290,8 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.statusCode).toBe(400);
   });
 
-  it('returns 400 for too many templates', async () => {
-    const templates = Array.from({length: 101}, (_, index) => ({
+  it('returns 400 for more than 1000 templates', async () => {
+    const templates = Array.from({length: 1001}, (_, index) => ({
       template_key: `linux-${index}`,
       labels: ['linux'],
       available_slots: 1,

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

@@ -23,6 +23,7 @@ import {
   runnersTestAuthClient,
 } from '#test/index.js';
 import {createRunnerRoutes} from './index.js';
+import {pollForRunnerActivationToken} from './runner-enrollment.js';
 
 const token = 'provisioner-test-token';
 const fakeUserAuth: AuthMethod = {name: AUTH_USER, authenticate: () => Promise.resolve()};
@@ -102,7 +103,8 @@ describe('runner enrollment control plane', () => {
       headers: {authorization: `Bearer ${controlToken}`},
       payload: {labels: ['Linux', 'linux'], provider_kind: 'docker', protocol_version: '1'},
     });
-    expect(enrolled.statusCode).toBe(204);
+    expect(enrolled.statusCode).toBe(200);
+    expect(enrolled.json()).toEqual({activation_token: null});
 
     const attached = await app.inject({
       method: 'POST',
@@ -137,6 +139,122 @@ describe('runner enrollment control plane', () => {
     expect(jobs.statusCode).toBe(401);
   });
 
+  it('promotes an intended reservation during enrollment and returns its activation token', async () => {
+    const workspaceId = crypto.randomUUID();
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Reservation insert returned no row');
+
+    const created = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/batch',
+      headers: {authorization: `Bearer ${token}`},
+      payload: {
+        runner_instances: [{template_key: 'linux', reservation_id: reservation.id}],
+      },
+    });
+    const runner = created.json().runner_instances[0];
+    const exchanged = await app.inject({
+      method: 'POST',
+      url: '/runner-enrollment/exchange',
+      payload: {bootstrap_token: runner.bootstrap_token},
+    });
+    const controlToken = exchanged.json().control_session_token;
+
+    const attached = await app.inject({
+      method: 'POST',
+      url: `/provisioners/runner-instances/${runner.runner_instance_id}/provider-runner`,
+      headers: {authorization: `Bearer ${token}`},
+      payload: {provider_runner_id: 'container-intended-assignment'},
+    });
+    expect(attached.json()).toEqual({attached: true});
+
+    const enrolled = await app.inject({
+      method: 'POST',
+      url: '/runner-control/enrollment',
+      headers: {authorization: `Bearer ${controlToken}`},
+      payload: {labels: ['linux'], provider_kind: 'docker', protocol_version: '1'},
+    });
+
+    expect(enrolled.statusCode).toBe(200);
+    expect(enrolled.json()).toEqual({activation_token: expect.any(String)});
+    const [instance] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.runner_instance_id));
+    expect(instance).toMatchObject({
+      workspaceId,
+      reservationId: reservation.id,
+      intendedReservationId: null,
+      state: 'running',
+      providerRunnerId: 'container-intended-assignment',
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('keeps enrollment successful when an intended reservation cannot be promoted', async () => {
+    const workspaceId = crypto.randomUUID();
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['gpu'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Reservation insert returned no row');
+
+    const created = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/batch',
+      headers: {authorization: `Bearer ${token}`},
+      payload: {runner_instances: [{reservation_id: reservation.id}]},
+    });
+    const runner = created.json().runner_instances[0];
+    const exchanged = await app.inject({
+      method: 'POST',
+      url: '/runner-enrollment/exchange',
+      payload: {bootstrap_token: runner.bootstrap_token},
+    });
+    const controlToken = exchanged.json().control_session_token;
+    await app.inject({
+      method: 'POST',
+      url: `/provisioners/runner-instances/${runner.runner_instance_id}/provider-runner`,
+      headers: {authorization: `Bearer ${token}`},
+      payload: {provider_runner_id: 'container-intended-mismatch'},
+    });
+
+    const enrolled = await app.inject({
+      method: 'POST',
+      url: '/runner-control/enrollment',
+      headers: {authorization: `Bearer ${controlToken}`},
+      payload: {labels: ['linux'], provider_kind: 'docker', protocol_version: '1'},
+    });
+
+    expect(enrolled.statusCode).toBe(200);
+    expect(enrolled.json()).toEqual({activation_token: null});
+    const [instance] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.runner_instance_id));
+    expect(instance).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: reservation.id,
+      state: 'running',
+    });
+  });
+
   it('rejects enrollment when the authenticated runner instance is terminal', async () => {
     const created = await app.inject({
       method: 'POST',
@@ -276,6 +394,31 @@ describe('runner enrollment control plane', () => {
     expect(closed.statusCode).toBe(401);
   });
 
+  it('keeps polling when an assignment cannot yet mint an activation token', async () => {
+    vi.useFakeTimers();
+    const abortController = new AbortController();
+    const getAssignment = vi.fn().mockResolvedValue({workspaceId: 'workspace-id'});
+    const issueActivationToken = vi
+      .fn()
+      .mockResolvedValueOnce(null)
+      .mockResolvedValueOnce('activation-token');
+    try {
+      const activationToken = pollForRunnerActivationToken({
+        deadline: Date.now() + 1_000,
+        intervalMs: 5,
+        signal: abortController.signal,
+        getAssignment,
+        issueActivationToken,
+      });
+      await Promise.resolve();
+      await Promise.resolve();
+      expect(issueActivationToken).toHaveBeenCalledTimes(1);
+      await vi.advanceTimersByTimeAsync(5);
+      await expect(activationToken).resolves.toBe('activation-token');
+    } finally {
+      vi.useRealTimers();
+    }
+  });
   it('rejects registration with an expired activation token', async () => {
     const created = await app.inject({
       method: 'POST',

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

@@ -11,6 +11,7 @@ import {
   runnerBootstrapExchangeResponseSchema,
   runnerControlHeartbeatResponseSchema,
   runnerEnrollmentBodySchema,
+  runnerEnrollmentResponseSchema,
 } from '@shipfox/api-runners-dto';
 import {ClientError, defineRoute} from '@shipfox/node-fastify';
 import {z} from 'zod';
@@ -43,9 +44,10 @@ export const createRunnerInstancesRoute = defineRoute({
     const results = await createRunnerInstancesWithBootstrapTokens({
       provisionerId: provisionerTokenId,
       ...(request.body.provider_kind ? {providerKind: request.body.provider_kind} : {}),
-      runnerInstances: request.body.runner_instances.map((runner) =>
-        runner.template_key ? {templateKey: runner.template_key} : {},
-      ),
+      runnerInstances: request.body.runner_instances.map((runner) => ({
+        ...(runner.template_key ? {templateKey: runner.template_key} : {}),
+        ...(runner.reservation_id ? {reservationId: runner.reservation_id} : {}),
+      })),
       ttlSeconds: config.RUNNER_BOOTSTRAP_TOKEN_TTL_SECONDS,
     });
     return {
@@ -111,16 +113,16 @@ export const enrollRunnerRoute = defineRoute({
   method: 'POST',
   path: '/enrollment',
   description: 'Declare the authenticated runner instance labels and protocol capabilities',
-  schema: {body: runnerEnrollmentBodySchema, response: {204: z.null()}},
+  schema: {body: runnerEnrollmentBodySchema, response: {200: runnerEnrollmentResponseSchema}},
   preHandler: authenticateRunnerControlSession,
   errorHandler: (error) => {
     if (error instanceof RunnerControlSessionInvalidError)
       throw new ClientError(error.message, 'runner-control-session-invalid', {status: 409});
     throw error;
   },
-  handler: async (request, reply) => {
+  handler: async (request) => {
     const session = requireRunnerControlSessionContext(request);
-    await enrollRunnerControlSession({
+    const activationToken = await enrollRunnerControlSession({
       runnerInstanceId: session.runnerInstanceId,
       provisionerId: session.provisionerId,
       labels: request.body.labels,
@@ -128,7 +130,7 @@ export const enrollRunnerRoute = defineRoute({
       providerKind: request.body.provider_kind,
       protocolVersion: request.body.protocol_version,
     });
-    return await reply.code(204).send();
+    return {activation_token: activationToken};
   },
 });
 
@@ -166,6 +168,27 @@ export const runnerControlHeartbeatRoute = defineRoute({
   },
 });
 
+type RunnerAssignmentPollParams = {
+  deadline: number;
+  intervalMs: number;
+  signal: AbortSignal;
+  getAssignment: () => Promise<unknown | null>;
+  issueActivationToken: () => Promise<string | null>;
+};
+export async function pollForRunnerActivationToken(
+  params: RunnerAssignmentPollParams,
+): Promise<string | null> {
+  while (Date.now() <= params.deadline) {
+    if (params.signal.aborted) return null;
+    const assignment = await params.getAssignment();
+    if (assignment) {
+      const activationToken = await params.issueActivationToken();
+      if (activationToken) return activationToken;
+    }
+    await new Promise((resolve) => setTimeout(resolve, params.intervalMs));
+  }
+  return null;
+}
 export const runnerAssignmentPollRoute = defineRoute({
   method: 'GET',
   path: '/assignment',
@@ -177,20 +200,17 @@ export const runnerAssignmentPollRoute = defineRoute({
     const deadline = Date.now() + config.RUNNER_ASSIGNMENT_POLL_MAX_WAIT_SECONDS * 1000;
     const abortController = new AbortController();
     reply.raw.on('close', () => abortController.abort());
-    while (Date.now() <= deadline) {
-      if (abortController.signal.aborted) return {activation_token: null};
-      const assignment = await getRunnerAssignment(session);
-      if (assignment) {
-        const activationToken = await issueRunnerActivationToken({
+    const activationToken = await pollForRunnerActivationToken({
+      deadline,
+      intervalMs: config.RUNNER_ASSIGNMENT_POLL_INTERVAL_MS,
+      signal: abortController.signal,
+      getAssignment: () => getRunnerAssignment(session),
+      issueActivationToken: () =>
+        issueRunnerActivationToken({
           ...session,
           ttlSeconds: config.RUNNER_ACTIVATION_TOKEN_TTL_SECONDS,
-        });
-        return {activation_token: activationToken};
-      }
-      await new Promise((resolve) =>
-        setTimeout(resolve, config.RUNNER_ASSIGNMENT_POLL_INTERVAL_MS),
-      );
-    }
-    return {activation_token: null};
+        }),
+    });
+    return {activation_token: activationToken};
   },
 });