feat: add stateless registration cleanup with multi-App support

lambdas/functions/termination-watcher/package.json

@@ -32,7 +32,8 @@
     "@octokit/core": "7.0.6",
     "@octokit/plugin-throttling": "11.0.3",
     "@octokit/request": "^9.2.2",
-    "@octokit/rest": "22.0.1"
+    "@octokit/rest": "22.0.1",
+    "@aws-github-runner/compute-providers": "*"
   },
   "nx": {
     "includedScripts": [

lambdas/functions/termination-watcher/src/github-app-client.test.ts

@@ -0,0 +1,208 @@
+import { createCommonStorage, type GitHubAppCredentialsStore } from '@aws-github-runner/storage-providers';
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { createThrottleOptions, resetAppCredentialsCache, createRunnerInstallationClient } from './github-app-client';
+import type { EndpointDefaults } from '@octokit/types';
+
+vi.mock('@aws-github-runner/storage-providers', () => ({
+  createCommonStorage: vi.fn(),
+}));
+
+const mockedCreateCommonStorage = vi.mocked(createCommonStorage);
+const mockGetCredentials = vi.fn<GitHubAppCredentialsStore['get']>();
+const credentialsStore = { get: mockGetCredentials } satisfies GitHubAppCredentialsStore;
+
+const mockCreateAppAuth = vi.fn();
+vi.mock('@octokit/auth-app', () => ({
+  createAppAuth: (...args: unknown[]) => mockCreateAppAuth(...args),
+}));
+
+const mockPaginate = {
+  iterator: vi.fn(),
+};
+
+const mockActions = {
+  listSelfHostedRunnersForOrg: vi.fn(),
+  listSelfHostedRunnersForRepo: vi.fn(),
+  deleteSelfHostedRunnerFromOrg: vi.fn(),
+  deleteSelfHostedRunnerFromRepo: vi.fn(),
+};
+
+const mockApps = {
+  getOrgInstallation: vi.fn(),
+  getRepoInstallation: vi.fn(),
+};
+
+const mockHookAfter = vi.fn();
+
+function MockOctokit() {
+  return {
+    hook: { after: mockHookAfter },
+    actions: mockActions,
+    apps: mockApps,
+    paginate: mockPaginate,
+  };
+}
+MockOctokit.plugin = vi.fn().mockReturnValue(MockOctokit);
+
+vi.mock('@octokit/rest', () => ({
+  Octokit: MockOctokit,
+}));
+
+vi.mock('@octokit/plugin-throttling', () => ({
+  throttling: vi.fn(),
+}));
+
+vi.mock('@octokit/request', () => ({
+  request: {
+    defaults: vi.fn().mockReturnValue(vi.fn()),
+  },
+}));
+
+function setupAuthMocks() {
+  mockGetCredentials.mockResolvedValue([{ appId: 12345, privateKey: 'fake-private-key' }]);
+
+  // App auth returns app token
+  const mockAuth = vi.fn();
+  mockAuth.mockImplementation((opts: { type: string }) => {
+    if (opts.type === 'app') {
+      return Promise.resolve({ token: 'app-token' });
+    }
+    return Promise.resolve({ token: 'installation-token' });
+  });
+  mockCreateAppAuth.mockReturnValue(mockAuth);
+}
+
+describe('multi-App installation clients', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    resetAppCredentialsCache();
+    mockedCreateCommonStorage.mockReturnValue({ githubAppCredentials: credentialsStore });
+    setupAuthMocks();
+  });
+
+  it('creates an organization installation client for scheduled registration cleanup', async () => {
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+    const client = await createRunnerInstallationClient('test-org', 'Org', '');
+    expect(client.actions).toBe(mockActions);
+    expect(mockApps.getOrgInstallation).toHaveBeenCalledWith({ org: 'test-org' });
+    expect(mockCreateAppAuth).toHaveBeenCalledWith({
+      appId: 12345,
+      privateKey: 'fake-private-key',
+      installationId: 999,
+    });
+  });
+
+  it('uses additional credentials and keeps App and installation auth paired', async () => {
+    mockGetCredentials.mockResolvedValue([
+      { appId: 1, privateKey: 'one' },
+      { appId: 2, privateKey: 'two' },
+    ]);
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
+    const random = vi.spyOn(Math, 'random').mockReturnValue(0.99);
+    try {
+      await createRunnerInstallationClient('test-org', 'Org', '');
+      expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
+      expect(mockCreateAppAuth).toHaveBeenNthCalledWith(2, { appId: 2, privateKey: 'two', installationId: 222 });
+    } finally {
+      random.mockRestore();
+    }
+  });
+
+  it('selects another App after an exhausted installation response', async () => {
+    mockGetCredentials.mockResolvedValue([
+      { appId: 1, privateKey: 'one' },
+      { appId: 2, privateKey: 'two' },
+    ]);
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
+    const random = vi.spyOn(Math, 'random').mockReturnValue(0);
+    try {
+      await createRunnerInstallationClient('test-org', 'Org', '');
+      mockHookAfter.mock.calls[0][1]({ headers: { 'x-ratelimit-remaining': '0' } });
+      mockCreateAppAuth.mockClear();
+      await createRunnerInstallationClient('test-org', 'Org', '');
+      expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
+    } finally {
+      random.mockRestore();
+    }
+  });
+
+  it('skips an App in secondary-limit cooldown', async () => {
+    mockGetCredentials.mockResolvedValue([
+      { appId: 1, privateKey: 'one' },
+      { appId: 2, privateKey: 'two' },
+    ]);
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
+    createThrottleOptions(1).onSecondaryRateLimit(60, { method: 'GET', url: '/runners' } as Required<EndpointDefaults>);
+    const random = vi.spyOn(Math, 'random').mockReturnValue(0);
+    try {
+      await createRunnerInstallationClient('test-org', 'Org', '');
+      expect(mockCreateAppAuth).toHaveBeenNthCalledWith(1, { appId: 2, privateKey: 'two' });
+    } finally {
+      random.mockRestore();
+    }
+  });
+
+  it('tries another configured App when installation authentication fails', async () => {
+    mockGetCredentials.mockResolvedValue([
+      { appId: 1, privateKey: 'one' },
+      { appId: 2, privateKey: 'two' },
+    ]);
+    mockApps.getOrgInstallation
+      .mockRejectedValueOnce(new Error('not installed'))
+      .mockResolvedValue({ data: { id: 222 } });
+    const random = vi.spyOn(Math, 'random').mockReturnValue(0);
+    try {
+      await createRunnerInstallationClient('test-org', 'Org', '');
+      expect(mockCreateAppAuth).toHaveBeenLastCalledWith({ appId: 2, privateKey: 'two', installationId: 222 });
+    } finally {
+      random.mockRestore();
+    }
+  });
+  it('reuses a selected App client within one invocation, while still switching exhausted Apps', async () => {
+    mockGetCredentials.mockResolvedValue([
+      { appId: 1, privateKey: 'one' },
+      { appId: 2, privateKey: 'two' },
+    ]);
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
+    const random = vi.spyOn(Math, 'random').mockReturnValue(0);
+    const clients = new Map();
+    try {
+      const first = await createRunnerInstallationClient('test-org', 'Org', '', clients);
+      // Unknown budgets tie; with the same selection, the client is reused.
+      expect(await createRunnerInstallationClient('test-org', 'Org', '', clients)).toBe(first);
+      expect(mockCreateAppAuth).toHaveBeenCalledTimes(2);
+      mockHookAfter.mock.calls[0][1]({ headers: { 'x-ratelimit-remaining': '0' } });
+      const second = await createRunnerInstallationClient('test-org', 'Org', '', clients);
+      expect(await createRunnerInstallationClient('test-org', 'Org', '', clients)).toBe(second);
+      expect(mockCreateAppAuth).toHaveBeenCalledTimes(4);
+      expect(mockApps.getOrgInstallation).toHaveBeenCalledTimes(2);
+      await createRunnerInstallationClient('test-org', 'Org', '', new Map());
+      expect(mockCreateAppAuth).toHaveBeenCalledTimes(6);
+    } finally {
+      random.mockRestore();
+    }
+  });
+
+  it('separates cached clients by owner, runner type, and GHES endpoint', async () => {
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 222 } });
+    mockApps.getRepoInstallation.mockResolvedValue({ data: { id: 333 } });
+    const clients = new Map();
+    await createRunnerInstallationClient('owner', 'Org', '', clients);
+    await createRunnerInstallationClient('other', 'Org', '', clients);
+    await createRunnerInstallationClient('owner/repo', 'Repo', '', clients);
+    await createRunnerInstallationClient('owner', 'Org', 'https://ghe.example/api/v3', clients);
+    expect(clients.size).toBe(4);
+    expect(mockCreateAppAuth).toHaveBeenCalledTimes(8);
+  });
+
+  it('does not cache failed authentication', async () => {
+    const clients = new Map();
+    mockApps.getOrgInstallation
+      .mockRejectedValueOnce(new Error('unavailable'))
+      .mockResolvedValue({ data: { id: 222 } });
+    await expect(createRunnerInstallationClient('owner', 'Org', '', clients)).rejects.toThrow('unavailable');
+    expect(clients.size).toBe(0);
+    await createRunnerInstallationClient('owner', 'Org', '', clients);
+    expect(clients.size).toBe(1);
+  });
+});

lambdas/functions/termination-watcher/src/github-app-client.ts

@@ -0,0 +1,178 @@
+import { createAppAuth } from '@octokit/auth-app';
+import { Octokit } from '@octokit/rest';
+import { throttling } from '@octokit/plugin-throttling';
+import { request } from '@octokit/request';
+import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
+import { createCommonStorage, type GitHubAppCredential } from '@aws-github-runner/storage-providers';
+import type { EndpointDefaults } from '@octokit/types';
+
+const logger = createChildLogger('github-app-client');
+
+let appCredentialsPromise: Promise<GitHubAppCredential[]> | undefined;
+
+const appBudgets = new Map<number, { remaining: number; cooldownUntil: number }>();
+
+function coolDown(appId?: number): void {
+  if (appId !== undefined) appBudgets.set(appId, { remaining: 0, cooldownUntil: Date.now() + 60000 });
+}
+
+function selectCredential(credentials: GitHubAppCredential[]): GitHubAppCredential {
+  const offset = Math.floor(Math.random() * credentials.length);
+  const rotated = [...credentials.slice(offset), ...credentials.slice(0, offset)];
+  const available = rotated.filter(
+    (credential) => (appBudgets.get(credential.appId)?.cooldownUntil ?? 0) <= Date.now(),
+  );
+  return (available.length ? available : rotated).reduce((best, credential) =>
+    (appBudgets.get(credential.appId)?.remaining ?? Infinity) > (appBudgets.get(best.appId)?.remaining ?? Infinity)
+      ? credential
+      : best,
+  );
+}
+
+export function createThrottleOptions(appId?: number) {
+  return {
+    onRateLimit: (_retryAfter: number, options: Required<EndpointDefaults>) => {
+      coolDown(appId);
+      logger.warn(`Rate limit hit for ${options.method} ${options.url}`);
+      return false;
+    },
+    onSecondaryRateLimit: (_retryAfter: number, options: Required<EndpointDefaults>) => {
+      coolDown(appId);
+      logger.warn(`Secondary rate limit hit for ${options.method} ${options.url}`);
+      return false;
+    },
+  };
+}
+
+async function loadAppCredentials(): Promise<GitHubAppCredential[]> {
+  const credentials = await createCommonStorage().githubAppCredentials.get();
+  if (credentials.length === 0) {
+    throw new Error('No GitHub App credentials found');
+  }
+  return credentials;
+}
+
+function getAppCredentials(): Promise<GitHubAppCredential[]> {
+  if (!appCredentialsPromise) {
+    appCredentialsPromise = loadAppCredentials().catch((error: unknown) => {
+      appCredentialsPromise = undefined;
+      throw error;
+    });
+  }
+  return appCredentialsPromise;
+}
+
+export function resetAppCredentialsCache(): void {
+  appCredentialsPromise = undefined;
+  appBudgets.clear();
+}
+
+function createOctokitInstance(token: string, ghesApiUrl: string, appId?: number): Octokit {
+  const CustomOctokit = Octokit.plugin(throttling);
+  const octokitOptions: ConstructorParameters<typeof Octokit>[0] = {
+    auth: token,
+  };
+  if (ghesApiUrl) {
+    octokitOptions.baseUrl = ghesApiUrl;
+  }
+  const client = new CustomOctokit({
+    ...octokitOptions,
+    userAgent: 'github-aws-runners-termination-watcher',
+    throttle: createThrottleOptions(appId),
+  });
+  if (appId !== undefined)
+    client.hook.after('request', (response) => {
+      const remaining = Number.parseInt(String(response.headers['x-ratelimit-remaining']), 10);
+      if (Number.isFinite(remaining))
+        appBudgets.set(appId, {
+          remaining,
+          cooldownUntil: remaining === 0 ? Date.now() + 60000 : 0,
+        });
+    });
+  return client;
+}
+
+async function createAuthenticatedClient(ghesApiUrl: string, credential: GitHubAppCredential): Promise<Octokit> {
+  const { appId, privateKey } = credential;
+  const authOptions: { appId: number; privateKey: string; request?: typeof request } = {
+    appId,
+    privateKey,
+  };
+  if (ghesApiUrl) {
+    authOptions.request = request.defaults({ baseUrl: ghesApiUrl });
+  }
+  const auth = createAppAuth(authOptions);
+  const appAuth = await auth({ type: 'app' });
+  return createOctokitInstance(appAuth.token, ghesApiUrl);
+}
+
+async function getInstallationId(octokit: Octokit, owner: string): Promise<number> {
+  const { data: installation } = await octokit.apps.getOrgInstallation({ org: owner });
+  return installation.id;
+}
+
+async function getInstallationIdForRepo(octokit: Octokit, owner: string, repo: string): Promise<number> {
+  const { data: installation } = await octokit.apps.getRepoInstallation({ owner, repo });
+  return installation.id;
+}
+
+async function createInstallationClient(
+  appOctokit: Octokit,
+  owner: string,
+  runnerType: string,
+  ghesApiUrl: string,
+  credential: GitHubAppCredential,
+): Promise<Octokit> {
+  let installationId: number;
+  if (runnerType === 'Repo') {
+    const [repoOwner, repo] = owner.split('/');
+    installationId = await getInstallationIdForRepo(appOctokit, repoOwner, repo);
+  } else {
+    installationId = await getInstallationId(appOctokit, owner);
+  }
+
+  const { appId, privateKey } = credential;
+  const authOptions: { appId: number; privateKey: string; installationId: number; request?: typeof request } = {
+    appId,
+    privateKey,
+    installationId,
+  };
+  if (ghesApiUrl) {
+    authOptions.request = request.defaults({ baseUrl: ghesApiUrl });
+  }
+  const auth = createAppAuth(authOptions);
+  const installationAuth = await auth({ type: 'installation' });
+  return createOctokitInstance(installationAuth.token, ghesApiUrl, appId);
+}
+
+/** Kept by the caller for one invocation; never stores scan progress. */
+export type InstallationClientCache = Map<string, Octokit>;
+
+export async function createRunnerInstallationClient(
+  owner: string,
+  runnerType: string,
+  ghesApiUrl: string,
+  clients?: InstallationClientCache,
+): Promise<Octokit> {
+  const remaining = [...(await getAppCredentials())];
+  while (remaining.length) {
+    const credential = selectCredential(remaining);
+    remaining.splice(remaining.indexOf(credential), 1);
+    // Select against current quota before consulting the cache, so a reused
+    // client never pins a batch to an exhausted or cooling-down App.
+    const key = JSON.stringify([credential.appId, runnerType, owner, ghesApiUrl]);
+    const cached = clients?.get(key);
+    if (cached) return cached;
+    try {
+      const appClient = await createAuthenticatedClient(ghesApiUrl, credential);
+      const client = await createInstallationClient(appClient, owner, runnerType, ghesApiUrl, credential);
+      clients?.set(key, client);
+      return client;
+    } catch (error) {
+      coolDown(credential.appId);
+      if (!remaining.length) throw error;
+      logger.warn('GitHub App authentication failed; trying another configured app', { appId: credential.appId });
+    }
+  }
+  throw new Error('No GitHub App credentials found');
+}