fix: retry installation auth for regional GitHub App events

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

@@ -824,6 +824,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: [],