fix: Fix dual GitHub App installation auth fallback for regional runners

Handle 404 installation token failures by resolving the installation from the active regional app and retrying auth. Add tests and avoid eager app-client creation in shared octokit auth.

Author: bogdankrasko

lambdas/functions/control-plane/src/github/octokit.test.ts

@@ -2,6 +2,7 @@ import { Octokit } from '@octokit/rest';
 import { ActionRequestMessage } from '../scale-runners/scale-up';
 import { getOctokit } from './octokit';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { createGithubAppAuth, createGithubInstallationAuth } from '../github/auth';
 
 const mockOctokit = {
   apps: {
@@ -11,7 +12,7 @@ const mockOctokit = {
 };
 
 vi.mock('../github/auth', async () => ({
-  createGithubInstallationAuth: vi.fn().mockImplementation(async (installationId) => {
+  createGithubInstallationAuth: vi.fn().mockImplementation(async (installationId: number) => {
     return { token: 'token', type: 'installation', installationId: installationId };
   }),
   createOctokitClient: vi.fn().mockImplementation(() => new Octokit()),
@@ -27,7 +28,11 @@ vi.mock('@octokit/rest', async () => ({
 // We've already mocked '../github/auth' above
 
 describe('Test getOctokit', () => {
-  const data = [
+  const data: Array<{
+    description: string;
+    input: { orgLevelRunner: boolean; installationId: number };
+    output: { callReposInstallation: boolean; callOrgInstallation: boolean };
+  }> = [
     {
       description: 'Should look-up org installation if installationId is 0.',
       input: { orgLevelRunner: false, installationId: 0 },
@@ -49,7 +54,7 @@ describe('Test getOctokit', () => {
     vi.clearAllMocks();
   });
 
-  it.each(data)(`$description`, async ({ input, output }) => {
+  it.each(data)(`$description`, async ({ input, output }: (typeof data)[number]) => {
     const payload = {
       eventType: 'workflow_job',
       id: 0,
@@ -74,6 +79,31 @@ describe('Test getOctokit', () => {
     } else if (output.callReposInstallation) {
       expect(mockOctokit.apps.getRepoInstallation).toHaveBeenCalled();
       expect(mockOctokit.apps.getOrgInstallation).not.toHaveBeenCalled();
+    } else {
+      expect(createGithubAppAuth).not.toHaveBeenCalled();
     }
   });
+
+  it('Should resolve installation again when event installation belongs to another app', async () => {
+    const payload = {
+      eventType: 'workflow_job',
+      id: 0,
+      installationId: 999,
+      repositoryOwner: 'owner',
+      repositoryName: 'repo',
+    } as ActionRequestMessage;
+
+    mockOctokit.apps.getOrgInstallation.mockResolvedValue({ data: { id: 123 } });
+
+    vi.mocked(createGithubInstallationAuth)
+      .mockRejectedValueOnce({ status: 404 })
+      .mockResolvedValueOnce({ token: 'token', type: 'installation', installationId: 123 });
+
+    await expect(getOctokit('', true, payload)).resolves.toBeDefined();
+
+    expect(createGithubAppAuth).toHaveBeenCalledTimes(1);
+    expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: 'owner' });
+    expect(createGithubInstallationAuth).toHaveBeenNthCalledWith(1, 999, '');
+    expect(createGithubInstallationAuth).toHaveBeenNthCalledWith(2, 123, '');
+  });
 });

lambdas/functions/control-plane/src/github/octokit.ts

@@ -2,17 +2,20 @@ import { Octokit } from '@octokit/rest';
 import { ActionRequestMessage } from '../scale-runners/scale-up';
 import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from './auth';
 
-export async function getInstallationId(
-  ghesApiUrl: string,
+function getErrorStatus(error: unknown): number | undefined {
+  if (typeof error !== 'object' || error === null) {
+    return undefined;
+  }
+
+  const errorWithStatus = error as { status?: number; response?: { status?: number } };
+  return errorWithStatus.status ?? errorWithStatus.response?.status;
+}
+
+async function resolveInstallationId(
+  githubClient: Octokit,
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
 ): Promise<number> {
-  if (payload.installationId !== 0) {
-    return payload.installationId;
-  }
-
-  const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl);
-  const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl);
   return enableOrgLevel
     ? (
         await githubClient.apps.getOrgInstallation({
@@ -27,6 +30,20 @@ export async function getInstallationId(
       ).data.id;
 }
 
+export async function getInstallationId(
+  ghesApiUrl: string,
+  enableOrgLevel: boolean,
+  payload: ActionRequestMessage,
+): Promise<number> {
+  if (payload.installationId !== 0) {
+    return payload.installationId;
+  }
+
+  const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl);
+  const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl);
+  return resolveInstallationId(githubClient, enableOrgLevel, payload);
+}
+
 /**
  *
  * Util method to get an octokit client based on provided installation id. This method should
@@ -40,7 +57,36 @@ export async function getOctokit(
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
 ): Promise<Octokit> {
-  const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload);
-  const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
-  return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  let githubAppClient: Octokit | undefined;
+  let installationId = payload.installationId !== 0 ? payload.installationId : undefined;
+
+  const getGithubAppClient = async (): Promise<Octokit> => {
+    if (githubAppClient === undefined) {
+      const appAuth = await createGithubAppAuth(undefined, ghesApiUrl);
+      githubAppClient = await createOctokitClient(appAuth.token, ghesApiUrl);
+    }
+
+    return githubAppClient;
+  };
+
+  try {
+    if (installationId === undefined) {
+      installationId = await resolveInstallationId(await getGithubAppClient(), enableOrgLevel, payload);
+    }
+
+    const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
+    return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  } catch (error) {
+    if (payload.installationId === 0 || getErrorStatus(error) !== 404) {
+      throw error;
+    }
+
+    const resolvedInstallationId = await resolveInstallationId(await getGithubAppClient(), enableOrgLevel, payload);
+    if (resolvedInstallationId === payload.installationId) {
+      throw error;
+    }
+
+    const ghAuth = await createGithubInstallationAuth(resolvedInstallationId, ghesApiUrl);
+    return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  }
 }

lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts

@@ -842,6 +842,34 @@ describe('scaleUp with GHES', () => {
       expect(mockedInstallationAuth).toHaveBeenCalledWith(200, 'https://github.enterprise.something/api/v3');
     });
 
+    it('Should resolve installation again when event installation belongs to another app', async () => {
+      mockOctokit.apps.getOrgInstallation.mockReset();
+      mockOctokit.apps.getOrgInstallation.mockImplementation(() => ({
+        data: {
+          id: 123,
+        },
+      }));
+
+      mockedInstallationAuth
+        .mockRejectedValueOnce({ status: 404 })
+        .mockResolvedValueOnce({
+          type: 'token',
+          tokenType: 'installation',
+          token: 'token',
+          createdAt: 'some-date',
+          expiresAt: 'some-date',
+          permissions: {},
+          repositorySelection: 'all',
+          installationId: 123,
+        });
+
+      await scaleUpModule.scaleUp(TEST_DATA);
+
+      expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: TEST_DATA_SINGLE.repositoryOwner });
+      expect(mockedInstallationAuth).toHaveBeenNthCalledWith(1, TEST_DATA_SINGLE.installationId, 'https://github.enterprise.something/api/v3');
+      expect(mockedInstallationAuth).toHaveBeenNthCalledWith(2, 123, 'https://github.enterprise.something/api/v3');
+    });
+
     it('Should reuse GitHub clients for same installation', async () => {
       const messages = createTestMessages(3, [
         { repositoryOwner: 'same-org' },

lambdas/functions/control-plane/src/scale-runners/scale-up.ts

@@ -149,15 +149,20 @@ function removeTokenFromLogging(config: string[]): string[] {
   return result;
 }
 
-export async function getInstallationId(
+function getErrorStatus(error: unknown): number | undefined {
+  if (typeof error !== 'object' || error === null) {
+    return undefined;
+  }
+
+  const errorWithStatus = error as { status?: number; response?: { status?: number } };
+  return errorWithStatus.status ?? errorWithStatus.response?.status;
+}
+
+async function resolveInstallationId(
   githubAppClient: Octokit,
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
 ): Promise<number> {
-  if (payload.installationId !== 0) {
-    return payload.installationId;
-  }
-
   return enableOrgLevel
     ? (
         await githubAppClient.apps.getOrgInstallation({
@@ -172,6 +177,18 @@ export async function getInstallationId(
       ).data.id;
 }
 
+export async function getInstallationId(
+  githubAppClient: Octokit,
+  enableOrgLevel: boolean,
+  payload: ActionRequestMessage,
+): Promise<number> {
+  if (payload.installationId !== 0) {
+    return payload.installationId;
+  }
+
+  return resolveInstallationId(githubAppClient, enableOrgLevel, payload);
+}
+
 export async function isJobQueued(githubInstallationClient: Octokit, payload: ActionRequestMessage): Promise<boolean> {
   let isQueued = false;
   if (payload.eventType === 'workflow_job') {
@@ -288,6 +305,39 @@ export async function createRunners(
   return instances;
 }
 
+async function createGithubInstallationClient(
+  githubAppClient: Octokit,
+  enableOrgLevel: boolean,
+  payload: ActionRequestMessage,
+  ghesApiUrl: string,
+): Promise<Octokit> {
+  let installationId = await getInstallationId(githubAppClient, enableOrgLevel, payload);
+
+  try {
+    const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
+    return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  } catch (error) {
+    if (payload.installationId === 0 || getErrorStatus(error) !== 404) {
+      throw error;
+    }
+
+    installationId = await resolveInstallationId(githubAppClient, enableOrgLevel, payload);
+    if (installationId === payload.installationId) {
+      throw error;
+    }
+
+    logger.warn('Retrying GitHub installation auth with installation resolved for current app', {
+      eventInstallationId: payload.installationId,
+      resolvedInstallationId: installationId,
+      repositoryOwner: payload.repositoryOwner,
+      repositoryName: payload.repositoryName,
+    });
+
+    const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
+    return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  }
+}
+
 export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<string[]> {
   logger.info('Received scale up requests', {
     n_requests: payloads.length,
@@ -373,9 +423,12 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
     // If we've not seen this owner/repo before, we'll need to create a GitHub
     // client for it.
     if (entry === undefined) {
-      const installationId = await getInstallationId(githubAppClient, enableOrgLevel, payload);
-      const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
-      const githubInstallationClient = await createOctokitClient(ghAuth.token, ghesApiUrl);
+      const githubInstallationClient = await createGithubInstallationClient(
+        githubAppClient,
+        enableOrgLevel,
+        payload,
+        ghesApiUrl,
+      );
 
       entry = {
         messages: [],