@shipfox/api-runners 2.0.0

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

@@ -0,0 +1,323 @@
+import {
+  AUTH_LEASED_JOB,
+  AUTH_PROVISIONER_TOKEN,
+  AUTH_RUNNER_REGISTRATION_TOKEN,
+  AUTH_RUNNER_SESSION,
+  AUTH_USER,
+  setProvisionerContext,
+} from '@shipfox/api-auth-context';
+import {
+  type AuthMethod,
+  ClientError,
+  closeApp,
+  createApp,
+  extractBearerToken,
+} from '@shipfox/node-fastify';
+import {vi} from '@shipfox/vitest/vi';
+import type {FastifyInstance, FastifyRequest} from 'fastify';
+import {db} from '#db/db.js';
+import {runningJobExecutions} from '#db/schema/running-job-executions.js';
+import {
+  provisionedRunnerCountDivergenceCount,
+  provisionedRunnerTerminateIntentIssuedCount,
+} from '#metrics/instance.js';
+import {pendingJobFactory, provisionedRunnerFactory, runnerSessionFactory} from '#test/index.js';
+import {runnerRoutes} from './index.js';
+
+const VALID_PROVISIONER_TOKEN = 'valid-provisioner-token';
+
+const passthroughAuth = (name: string): AuthMethod => ({
+  name,
+  authenticate: () => Promise.resolve(),
+});
+
+describe('POST /provisioners/demand/poll', () => {
+  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 !== VALID_PROVISIONER_TOKEN) {
+        throw new ClientError('Invalid provisioner token', 'unauthorized', {status: 401});
+      }
+      setProvisionerContext(request, {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: runnerRoutes,
+      swagger: false,
+    });
+    await app.ready();
+  });
+
+  afterAll(async () => {
+    await closeApp();
+  });
+
+  beforeEach(() => {
+    workspaceId = crypto.randomUUID();
+    provisionerTokenId = crypto.randomUUID();
+  });
+
+  it('returns demand stats and reservations when matching demand exists', 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);
+    expect(res.json()).toMatchObject({
+      stats: [{labels: ['linux'], queued: 1, reserved: 1}],
+      reservations: [{labels: ['linux'], count: 1}],
+      terminate_provisioned_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('returns stats without reservations in observe-only mode', 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: {
+        wait_seconds: 0,
+        max_reservations: 0,
+        templates: [
+          {
+            template_key: 'linux',
+            labels: ['linux'],
+            available_slots: 1,
+            starting: 0,
+            running: 1,
+          },
+        ],
+      },
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toMatchObject({
+      stats: [{labels: ['linux'], queued: 1, reserved: 0}],
+      reservations: [],
+      terminate_provisioned_runner_ids: [],
+    });
+  });
+
+  it('returns terminate intent ids for active provisioned runners with cancelled latest jobs', async () => {
+    await provisionedRunnerFactory.create({
+      workspaceId,
+      provisionerId: provisionerTokenId,
+      provisionedRunnerId: 'provisioned-runner-1',
+      state: 'running',
+    });
+    await insertRunningJob({
+      provisionedRunnerId: 'provisioned-runner-1',
+      cancellationRequestedAt: new Date('2025-01-01T00:01:00.000Z'),
+    });
+
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: {
+        wait_seconds: 0,
+        max_reservations: 0,
+        templates: [
+          {
+            template_key: 'linux',
+            labels: ['linux'],
+            available_slots: 1,
+            starting: 0,
+            running: 1,
+          },
+        ],
+      },
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toMatchObject({
+      reservations: [],
+      terminate_provisioned_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({
+      workspaceId,
+      provisionerId: provisionerTokenId,
+      provisionedRunnerId: 'provisioned-runner-1',
+      templateKey: 'linux',
+      state: 'running',
+    });
+    await provisionedRunnerFactory.create({
+      workspaceId,
+      provisionerId: provisionerTokenId,
+      provisionedRunnerId: 'provisioned-runner-2',
+      templateKey: 'linux',
+      state: 'running',
+    });
+    await insertRunningJob({
+      provisionedRunnerId: 'provisioned-runner-1',
+      cancellationRequestedAt: new Date('2025-01-01T00:01:00.000Z'),
+    });
+    const divergenceCallsBefore = divergenceSpy.mock.calls.length;
+    const intentCallsBefore = intentSpy.mock.calls.length;
+
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: {
+        wait_seconds: 0,
+        max_reservations: 0,
+        templates: [
+          {
+            template_key: 'linux',
+            labels: ['linux'],
+            available_slots: 1,
+            starting: 0,
+            running: 1,
+          },
+        ],
+      },
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toMatchObject({
+      reservations: [],
+      terminate_provisioned_runner_ids: ['provisioned-runner-1'],
+    });
+    const divergenceCalls = divergenceSpy.mock.calls
+      .slice(divergenceCallsBefore)
+      .filter(
+        ([value, attributes]) =>
+          value === 1 &&
+          JSON.stringify(attributes) ===
+            JSON.stringify({state: 'running', direction: 'backend-higher'}),
+      );
+    const intentCalls = intentSpy.mock.calls
+      .slice(intentCallsBefore)
+      .filter(
+        ([value, attributes]) =>
+          value === 1 &&
+          JSON.stringify(attributes) ===
+            JSON.stringify({surface: 'poll-demand', reason: 'job-cancelled'}),
+      );
+    expect(divergenceCalls).toHaveLength(1);
+    expect(intentCalls).toHaveLength(1);
+  });
+
+  it('returns 400 for max reservations above the request bound', async () => {
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: body({max_reservations: 1001}),
+    });
+
+    expect(res.statusCode).toBe(400);
+  });
+
+  it('returns 400 for too many templates', async () => {
+    const templates = Array.from({length: 101}, (_, index) => ({
+      template_key: `linux-${index}`,
+      labels: ['linux'],
+      available_slots: 1,
+      starting: 0,
+      running: 0,
+    }));
+
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: {wait_seconds: 0, max_reservations: 1, templates},
+    });
+
+    expect(res.statusCode).toBe(400);
+  });
+
+  it('returns 401 without authorization', async () => {
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      payload: body({max_reservations: 1}),
+    });
+
+    expect(res.statusCode).toBe(401);
+  });
+
+  it('returns 401 with invalid authorization', async () => {
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: 'Bearer invalid'},
+      payload: body({max_reservations: 1}),
+    });
+
+    expect(res.statusCode).toBe(401);
+  });
+
+  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,
+        },
+      ],
+    };
+  }
+
+  async function insertRunningJob(params: {
+    provisionedRunnerId: string;
+    cancellationRequestedAt?: Date | null;
+  }) {
+    const runnerSession = await runnerSessionFactory.create({workspaceId});
+
+    await db()
+      .insert(runningJobExecutions)
+      .values({
+        workspaceId,
+        workflowRunId: crypto.randomUUID(),
+        workflowRunAttemptId: crypto.randomUUID(),
+        jobId: crypto.randomUUID(),
+        jobExecutionId: crypto.randomUUID(),
+        projectId: crypto.randomUUID(),
+        runnerSessionId: runnerSession.id,
+        provisionerId: provisionerTokenId,
+        provisionedRunnerId: params.provisionedRunnerId,
+        requiredLabels: ['linux'],
+        runnerLabels: ['linux'],
+        startedAt: new Date('2025-01-01T00:00:00.000Z'),
+        lastHeartbeatAt: new Date('2025-01-01T00:00:00.000Z'),
+        cancellationRequestedAt: params.cancellationRequestedAt ?? null,
+      });
+  }
+});

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

@@ -0,0 +1,56 @@
+import {requireProvisionerContext} from '@shipfox/api-auth-context';
+import {pollDemandBodySchema, pollDemandResponseSchema} from '@shipfox/api-runners-dto';
+import {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 {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,
+    },
+  },
+  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);
+      }
+    });
+
+    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;
+
+    return toPollDemandResponseDto(result);
+  },
+});

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

@@ -0,0 +1,172 @@
+import {AUTH_PROVISIONER_TOKEN, AUTH_USER} from '@shipfox/api-auth-context';
+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 * as provisionerTokenDb from '#db/provisioner-tokens.js';
+import {revokeProvisionerToken} from '#db/provisioner-tokens.js';
+import {createProvisionerTokenAuthMethod} from '#presentation/auth/index.js';
+import {provisionerTokenFactory} from '#test/index.js';
+import {provisionerRoutes} from './index.js';
+
+const fakeUserAuth: AuthMethod = {
+  name: AUTH_USER,
+  authenticate: () => Promise.resolve(),
+};
+
+describe('provisioner me route', () => {
+  let app: FastifyInstance;
+
+  beforeEach(async () => {
+    await closeApp();
+    app = await createApp({
+      auth: [fakeUserAuth, createProvisionerTokenAuthMethod()],
+      routes: provisionerRoutes,
+      swagger: false,
+    });
+    await app.ready();
+  });
+
+  afterEach(async () => {
+    await closeApp();
+  });
+
+  test('uses provisioner auth for provisioner routes', () => {
+    expect(provisionerRoutes[2]?.auth).toBe(AUTH_PROVISIONER_TOKEN);
+  });
+
+  it('returns the authenticated provisioner identity', async () => {
+    const workspaceId = crypto.randomUUID();
+    const rawToken = generateOpaqueToken('provisionerToken');
+    const token = await provisionerTokenFactory.create({workspaceId}, {transient: {rawToken}});
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toEqual({id: token.id, workspace_id: workspaceId});
+    const touched = await provisionerTokenDb.resolveProvisionerTokenByHash(token.hashedToken);
+    expect(touched?.lastSeenAt).toBeInstanceOf(Date);
+  });
+
+  it('does not reject valid auth when the last-seen write fails', async () => {
+    await closeApp();
+    const workspaceId = crypto.randomUUID();
+    const rawToken = generateOpaqueToken('provisionerToken');
+    const token = await provisionerTokenFactory.create({workspaceId}, {transient: {rawToken}});
+    vi.spyOn(provisionerTokenDb, 'touchProvisionerLastSeen').mockRejectedValueOnce(
+      new Error('database unavailable'),
+    );
+    app = await createApp({
+      auth: [fakeUserAuth, createProvisionerTokenAuthMethod()],
+      routes: provisionerRoutes,
+      swagger: false,
+    });
+    await app.ready();
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(res.json()).toEqual({id: token.id, workspace_id: workspaceId});
+  });
+
+  it('returns 401 without an authorization header', async () => {
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('unauthorized');
+  });
+
+  it('returns 401 for a non-provisioner token type before lookup', async () => {
+    const rawToken = generateOpaqueToken('manualRegistrationToken');
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('unauthorized');
+  });
+
+  it('returns 401 for an unknown provisioner token', async () => {
+    const rawToken = generateOpaqueToken('provisionerToken');
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('unauthorized');
+  });
+
+  it('returns 401 for a revoked provisioner token', async () => {
+    const workspaceId = crypto.randomUUID();
+    const rawToken = generateOpaqueToken('provisionerToken');
+    const token = await provisionerTokenFactory.create({workspaceId}, {transient: {rawToken}});
+    await revokeProvisionerToken({
+      tokenId: token.id,
+      workspaceId,
+      revokedByUserId: crypto.randomUUID(),
+    });
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('provisioner-token-revoked');
+  });
+
+  it('returns 401 for an expired provisioner token', async () => {
+    const workspaceId = crypto.randomUUID();
+    const rawToken = generateOpaqueToken('provisionerToken');
+    await provisionerTokenFactory.create(
+      {workspaceId, expiresAt: new Date(Date.now() - 60_000)},
+      {transient: {rawToken}},
+    );
+
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('provisioner-token-expired');
+  });
+
+  it('returns 401 for a provisioner token expiring at the current instant', async () => {
+    const now = new Date('2026-01-01T00:00:00.000Z');
+    const workspaceId = crypto.randomUUID();
+    const rawToken = generateOpaqueToken('provisionerToken');
+    await provisionerTokenFactory.create({workspaceId, expiresAt: now}, {transient: {rawToken}});
+
+    vi.useFakeTimers();
+    vi.setSystemTime(now);
+    const res = await app.inject({
+      method: 'GET',
+      url: '/provisioners/me',
+      headers: {authorization: `Bearer ${rawToken}`},
+    });
+    vi.useRealTimers();
+
+    expect(res.statusCode).toBe(401);
+    expect(res.json().code).toBe('provisioner-token-expired');
+  });
+});