@shipfox/api-runners 5.0.0 → 7.0.1

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

@@ -14,17 +14,25 @@ import {
   extractBearerToken,
 } from '@shipfox/node-fastify';
 import {vi} from '@shipfox/vitest/vi';
+import {eq} from 'drizzle-orm';
 import type {FastifyInstance, FastifyRequest} from 'fastify';
 import {db} from '#db/db.js';
+import {provisionerCapabilitySnapshots} from '#db/schema/provisioner-capability-snapshots.js';
 import {runningJobExecutions} from '#db/schema/running-job-executions.js';
 import {
-  provisionedRunnerCountDivergenceCount,
-  provisionedRunnerTerminateIntentIssuedCount,
+  providerRunnerCountDivergenceCount,
+  providerRunnerTerminateIntentIssuedCount,
 } from '#metrics/instance.js';
-import {pendingJobFactory, provisionedRunnerFactory, runnerSessionFactory} from '#test/index.js';
-import {runnerRoutes} from './index.js';
+import {
+  pendingJobFactory,
+  providerRunnerFactory,
+  runnerSessionFactory,
+  runnersTestAuthClient,
+} from '#test/index.js';
+import {createRunnerRoutes} from './index.js';
 
 const VALID_PROVISIONER_TOKEN = 'valid-provisioner-token';
+const INSTALLATION_PROVISIONER_TOKEN = 'installation-provisioner-token';
 
 const passthroughAuth = (name: string): AuthMethod => ({
   name,
@@ -41,9 +49,13 @@ describe('POST /provisioners/demand/poll', () => {
     authenticate: (request: FastifyRequest) => {
       const rawToken = extractBearerToken(request.headers.authorization);
       if (rawToken !== VALID_PROVISIONER_TOKEN) {
+        if (rawToken === INSTALLATION_PROVISIONER_TOKEN) {
+          setProvisionerContext(request, {scope: 'installation', provisionerTokenId});
+          return Promise.resolve();
+        }
         throw new ClientError('Invalid provisioner token', 'unauthorized', {status: 401});
       }
-      setProvisionerContext(request, {workspaceId, provisionerTokenId});
+      setProvisionerContext(request, {scope: 'workspace', workspaceId, provisionerTokenId});
       return Promise.resolve();
     },
   };
@@ -57,7 +69,7 @@ describe('POST /provisioners/demand/poll', () => {
         passthroughAuth(AUTH_LEASED_JOB),
         fakeProvisionerAuth,
       ],
-      routes: runnerRoutes,
+      routes: createRunnerRoutes(runnersTestAuthClient),
       swagger: false,
     });
     await app.ready();
@@ -86,12 +98,24 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.json()).toMatchObject({
       stats: [{labels: ['linux'], queued: 1, reserved: 1}],
       reservations: [{labels: ['linux'], count: 1}],
-      terminate_provisioned_runner_ids: [],
+      terminate_provider_runner_ids: [],
     });
     expect(res.json().reservations[0].reservation_id).toEqual(expect.any(String));
     expect(res.json().reservations[0].expires_at).toEqual(expect.any(String));
   });
 
+  it('rejects installation provisioner credentials from workspace demand polling', async () => {
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${INSTALLATION_PROVISIONER_TOKEN}`},
+      payload: body({max_reservations: 1}),
+    });
+
+    expect(res.statusCode).toBe(403);
+    expect(res.json().code).toBe('forbidden');
+  });
+
   it('returns stats without reservations in observe-only mode', async () => {
     await pendingJobFactory.create({workspaceId, requiredLabels: ['linux']});
 
@@ -118,19 +142,46 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.json()).toMatchObject({
       stats: [{labels: ['linux'], queued: 1, reserved: 0}],
       reservations: [],
-      terminate_provisioned_runner_ids: [],
+      terminate_provider_runner_ids: [],
     });
+    const snapshots = await db()
+      .select()
+      .from(provisionerCapabilitySnapshots)
+      .where(eq(provisionerCapabilitySnapshots.provisionerId, provisionerTokenId));
+    expect(snapshots).toEqual([
+      expect.objectContaining({
+        workspaceId,
+        labels: ['linux'],
+        availableSlots: 1,
+        starting: 0,
+        running: 1,
+      }),
+    ]);
+
+    const withdrawal = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: {wait_seconds: 0, max_reservations: 0, templates: []},
+    });
+    const clearedSnapshots = await db()
+      .select()
+      .from(provisionerCapabilitySnapshots)
+      .where(eq(provisionerCapabilitySnapshots.provisionerId, provisionerTokenId));
+
+    expect(withdrawal.statusCode).toBe(200);
+    expect(clearedSnapshots).toEqual([]);
   });
 
   it('returns terminate intent ids for active provisioned runners with cancelled latest jobs', async () => {
-    await provisionedRunnerFactory.create({
+    await providerRunnerFactory.create({
       workspaceId,
       provisionerId: provisionerTokenId,
-      provisionedRunnerId: 'provisioned-runner-1',
+      providerRunnerId: 'provisioned-runner-1',
       state: 'running',
     });
     await insertRunningJob({
-      provisionedRunnerId: 'provisioned-runner-1',
+      providerRunnerId: 'provisioned-runner-1',
       cancellationRequestedAt: new Date('2025-01-01T00:01:00.000Z'),
     });
 
@@ -156,29 +207,29 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.statusCode).toBe(200);
     expect(res.json()).toMatchObject({
       reservations: [],
-      terminate_provisioned_runner_ids: ['provisioned-runner-1'],
+      terminate_provider_runner_ids: ['provisioned-runner-1'],
     });
   });
 
   it('records count divergence and terminate-intent metrics for the returned poll result', async () => {
-    const divergenceSpy = vi.spyOn(provisionedRunnerCountDivergenceCount, 'add');
-    const intentSpy = vi.spyOn(provisionedRunnerTerminateIntentIssuedCount, 'add');
-    await provisionedRunnerFactory.create({
+    const divergenceSpy = vi.spyOn(providerRunnerCountDivergenceCount, 'add');
+    const intentSpy = vi.spyOn(providerRunnerTerminateIntentIssuedCount, 'add');
+    await providerRunnerFactory.create({
       workspaceId,
       provisionerId: provisionerTokenId,
-      provisionedRunnerId: 'provisioned-runner-1',
+      providerRunnerId: 'provisioned-runner-1',
       templateKey: 'linux',
       state: 'running',
     });
-    await provisionedRunnerFactory.create({
+    await providerRunnerFactory.create({
       workspaceId,
       provisionerId: provisionerTokenId,
-      provisionedRunnerId: 'provisioned-runner-2',
+      providerRunnerId: 'provisioned-runner-2',
       templateKey: 'linux',
       state: 'running',
     });
     await insertRunningJob({
-      provisionedRunnerId: 'provisioned-runner-1',
+      providerRunnerId: 'provisioned-runner-1',
       cancellationRequestedAt: new Date('2025-01-01T00:01:00.000Z'),
     });
     const divergenceCallsBefore = divergenceSpy.mock.calls.length;
@@ -206,7 +257,7 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.statusCode).toBe(200);
     expect(res.json()).toMatchObject({
       reservations: [],
-      terminate_provisioned_runner_ids: ['provisioned-runner-1'],
+      terminate_provider_runner_ids: ['provisioned-runner-1'],
     });
     const divergenceCalls = divergenceSpy.mock.calls
       .slice(divergenceCallsBefore)
@@ -296,7 +347,7 @@ describe('POST /provisioners/demand/poll', () => {
   }
 
   async function insertRunningJob(params: {
-    provisionedRunnerId: string;
+    providerRunnerId: string;
     cancellationRequestedAt?: Date | null;
   }) {
     const runnerSession = await runnerSessionFactory.create({workspaceId});
@@ -312,7 +363,7 @@ describe('POST /provisioners/demand/poll', () => {
         projectId: crypto.randomUUID(),
         runnerSessionId: runnerSession.id,
         provisionerId: provisionerTokenId,
-        provisionedRunnerId: params.provisionedRunnerId,
+        providerRunnerId: params.providerRunnerId,
         requiredLabels: ['linux'],
         runnerLabels: ['linux'],
         startedAt: new Date('2025-01-01T00:00:00.000Z'),
@@ -321,3 +372,91 @@ describe('POST /provisioners/demand/poll', () => {
       });
   }
 });
+
+describe('POST /provisioners/demand/poll with installation provisioning configured', () => {
+  let app: FastifyInstance;
+  let workspaceId: string;
+  let provisionerTokenId: string;
+
+  const fakeProvisionerAuth: AuthMethod = {
+    name: AUTH_PROVISIONER_TOKEN,
+    authenticate: (request: FastifyRequest) => {
+      const rawToken = extractBearerToken(request.headers.authorization);
+      if (rawToken === INSTALLATION_PROVISIONER_TOKEN) {
+        setProvisionerContext(request, {scope: 'installation', provisionerTokenId});
+        return Promise.resolve();
+      }
+      setProvisionerContext(request, {scope: 'workspace', workspaceId, provisionerTokenId});
+      return Promise.resolve();
+    },
+  };
+
+  beforeAll(async () => {
+    app = await createApp({
+      auth: [
+        passthroughAuth(AUTH_USER),
+        passthroughAuth(AUTH_RUNNER_REGISTRATION_TOKEN),
+        passthroughAuth(AUTH_RUNNER_SESSION),
+        passthroughAuth(AUTH_LEASED_JOB),
+        fakeProvisionerAuth,
+      ],
+      routes: createRunnerRoutes(runnersTestAuthClient, {
+        installationProvisioning: {
+          policy: {filterEligibleWorkspaceIds: vi.fn().mockResolvedValue(new Set())},
+        },
+      }),
+      swagger: false,
+    });
+    await app.ready();
+  });
+
+  afterAll(async () => {
+    await closeApp();
+  });
+
+  beforeEach(() => {
+    workspaceId = crypto.randomUUID();
+    provisionerTokenId = crypto.randomUUID();
+  });
+
+  it('returns opaque empty demand when no workspace is eligible', async () => {
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${INSTALLATION_PROVISIONER_TOKEN}`},
+      payload: body({max_reservations: 1}),
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toEqual({stats: [], reservations: [], terminate_provider_runner_ids: []});
+  });
+
+  it('still serves workspace provisioner credentials', async () => {
+    await pendingJobFactory.create({workspaceId, requiredLabels: ['linux']});
+
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: body({max_reservations: 1}),
+    });
+
+    expect(res.statusCode).toBe(200);
+  });
+
+  function body(params: {max_reservations: number}) {
+    return {
+      wait_seconds: 0,
+      max_reservations: params.max_reservations,
+      templates: [
+        {
+          template_key: 'linux',
+          labels: ['linux'],
+          available_slots: 1,
+          starting: 0,
+          running: 0,
+        },
+      ],
+    };
+  }
+});

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

@@ -1,56 +1,116 @@
-import {requireProvisionerContext} from '@shipfox/api-auth-context';
+import {
+  requireProvisionerContext,
+  requireWorkspaceProvisionerContext,
+} from '@shipfox/api-auth-context';
+import type {PollDemandTemplateDto} from '@shipfox/api-runners-dto';
 import {pollDemandBodySchema, pollDemandResponseSchema} from '@shipfox/api-runners-dto';
-import {defineRoute} from '@shipfox/node-fastify';
+import {ClientError, defineRoute} from '@shipfox/node-fastify';
 import {config} from '#config.js';
 import {pollDemand, releaseReservationGrants} from '#core/demand.js';
-import type {ReservationGrant} from '#db/reservations.js';
+import {publishWorkspaceProvisionerCapabilitySnapshot} from '#db/provisioner-capability-snapshots.js';
+import {
+  listQueuedDemandWorkspaceIds,
+  pollInstallationDemandAndReserve,
+  type ReservationGrant,
+  type ReservationTemplate,
+} from '#db/reservations.js';
+import type {CreateRunnersModuleOptions} from '#installation-provisioning.js';
 import {toPollDemandResponseDto} from '#presentation/dto/index.js';
 
 const TERMINATE_INTENT_LIMIT = 1000;
 
-export const pollDemandRoute = defineRoute({
-  method: 'POST',
-  path: '/demand/poll',
-  description: 'Poll aggregate runner demand and reserve capacity for a provisioner',
-  schema: {
-    body: pollDemandBodySchema,
-    response: {
-      200: pollDemandResponseSchema,
+function toReservationTemplates(templates: PollDemandTemplateDto[]): ReservationTemplate[] {
+  return templates.map((template) => ({
+    templateKey: template.template_key,
+    labels: template.labels,
+    availableSlots: template.available_slots,
+    starting: template.starting,
+    running: template.running,
+  }));
+}
+
+export function createPollDemandRoute(options: CreateRunnersModuleOptions = {}) {
+  return defineRoute({
+    method: 'POST',
+    path: '/demand/poll',
+    description: 'Poll aggregate runner demand and reserve capacity for a provisioner',
+    schema: {
+      body: pollDemandBodySchema,
+      response: {
+        200: pollDemandResponseSchema,
+      },
     },
-  },
-  handler: async (request, reply) => {
-    const {provisionerTokenId, workspaceId} = requireProvisionerContext(request);
-    const abortController = new AbortController();
-    let responseFinished = false;
-    let responseReservations: ReservationGrant[] = [];
-    reply.raw.on('finish', () => {
-      responseFinished = true;
-    });
-    reply.raw.on('close', () => {
-      if (!responseFinished) {
-        abortController.abort();
-        void releaseReservationGrants(responseReservations).catch(() => undefined);
+    handler: async (request, reply) => {
+      const abortController = new AbortController();
+      let responseFinished = false;
+      let responseReservations: ReservationGrant[] = [];
+      reply.raw.on('finish', () => {
+        responseFinished = true;
+      });
+      reply.raw.on('close', () => {
+        if (!responseFinished) {
+          abortController.abort();
+          void releaseReservationGrants(responseReservations).catch(() => undefined);
+        }
+      });
+      const provisionerContext = requireProvisionerContext(request);
+      if (provisionerContext.scope === 'installation') {
+        if (!options.installationProvisioning) {
+          throw new ClientError(
+            'Installation provisioner credentials are not accepted on this host',
+            'forbidden',
+            {
+              status: 403,
+            },
+          );
+        }
+        const candidateWorkspaceIds = await listQueuedDemandWorkspaceIds();
+        const eligibleWorkspaceIds =
+          await options.installationProvisioning.policy.filterEligibleWorkspaceIds(
+            candidateWorkspaceIds,
+          );
+        const templates = toReservationTemplates(request.body.templates);
+        const result = await pollInstallationDemandAndReserve({
+          provisionerId: provisionerContext.provisionerTokenId,
+          maxReservations: request.body.max_reservations,
+          ttlSeconds: config.RESERVATION_TTL_SECONDS,
+          templates,
+          capabilityWindowSeconds: config.PROVISIONER_ACTIVE_WINDOW_SECONDS,
+          eligibleWorkspaceIds,
+          signal: abortController.signal,
+          onReservations: (reservations) => {
+            responseReservations.push(...reservations);
+          },
+        });
+        return toPollDemandResponseDto({
+          ...result,
+          terminateRunnerInstanceIds: [],
+        });
       }
-    });
+      const {provisionerTokenId, workspaceId} = requireWorkspaceProvisionerContext(request);
+
+      const templates = toReservationTemplates(request.body.templates);
+      await publishWorkspaceProvisionerCapabilitySnapshot({
+        workspaceId,
+        provisionerId: provisionerTokenId,
+        templates,
+      });
 
-    const result = await pollDemand({
-      workspaceId,
-      provisionerId: provisionerTokenId,
-      maxReservations: request.body.max_reservations,
-      waitSeconds: request.body.wait_seconds,
-      ttlSeconds: config.RESERVATION_TTL_SECONDS,
-      terminateIntentLimit: TERMINATE_INTENT_LIMIT,
-      templates: request.body.templates.map((template) => ({
-        templateKey: template.template_key,
-        labels: template.labels,
-        availableSlots: template.available_slots,
-        starting: template.starting,
-        running: template.running,
-      })),
-      signal: abortController.signal,
-    });
-    responseReservations = result.reservations;
+      const result = await pollDemand({
+        workspaceId,
+        provisionerId: provisionerTokenId,
+        maxReservations: request.body.max_reservations,
+        waitSeconds: request.body.wait_seconds,
+        ttlSeconds: config.RESERVATION_TTL_SECONDS,
+        terminateIntentLimit: TERMINATE_INTENT_LIMIT,
+        templates,
+        signal: abortController.signal,
+      });
+      responseReservations = result.reservations;
+
+      return toPollDemandResponseDto(result);
+    },
+  });
+}
 
-    return toPollDemandResponseDto(result);
-  },
-});
+export const pollDemandRoute = createPollDemandRoute();

package/src/presentation/routes/provisioner-me.test.ts

@@ -3,8 +3,12 @@ import type {AuthMethod} from '@shipfox/node-fastify';
 import {closeApp, createApp} from '@shipfox/node-fastify';
 import {generateOpaqueToken} from '@shipfox/node-tokens';
 import type {FastifyInstance} from 'fastify';
+import {createInstallationProvisionerToken} from '#core/provisioner-tokens.js';
 import * as provisionerTokenDb from '#db/provisioner-tokens.js';
-import {revokeProvisionerToken} from '#db/provisioner-tokens.js';
+import {
+  revokeInstallationProvisionerToken,
+  revokeProvisionerToken,
+} from '#db/provisioner-tokens.js';
 import {createProvisionerTokenAuthMethod} from '#presentation/auth/index.js';
 import {provisionerTokenFactory} from '#test/index.js';
 import {provisionerRoutes} from './index.js';
@@ -47,11 +51,41 @@ describe('provisioner me route', () => {
     });
 
     expect(res.statusCode).toBe(200);
-    expect(res.json()).toEqual({id: token.id, workspace_id: workspaceId});
+    expect(res.json()).toEqual({id: token.id, scope: 'workspace', workspace_id: workspaceId});
     const touched = await provisionerTokenDb.resolveProvisionerTokenByHash(token.hashedToken);
     expect(touched?.lastSeenAt).toBeInstanceOf(Date);
   });
 
+  it('returns an installation identity without a workspace', async () => {
+    const token = await createInstallationProvisionerToken({createdByUserId: crypto.randomUUID()});
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${token.rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toEqual({id: token.token.id, scope: 'installation'});
+  });
+
+  it('rejects a revoked installation provisioner token', async () => {
+    const token = await createInstallationProvisionerToken({createdByUserId: crypto.randomUUID()});
+    await revokeInstallationProvisionerToken({
+      tokenId: token.token.id,
+      revokedByUserId: crypto.randomUUID(),
+    });
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${token.rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('provisioner-token-revoked');
+  });
+
   it('does not reject valid auth when the last-seen write fails', async () => {
     await closeApp();
     const workspaceId = crypto.randomUUID();
@@ -74,7 +108,7 @@ describe('provisioner me route', () => {
     });
 
     expect(res.statusCode).toBe(200);
-    expect(res.json()).toEqual({id: token.id, workspace_id: workspaceId});
+    expect(res.json()).toEqual({id: token.id, scope: 'workspace', workspace_id: workspaceId});
   });
 
   it('returns 401 without an authorization header', async () => {