feat(storage): add runner config consumer

lambdas/libs/storage-providers/aws/ssm/runner-config-consumer-common.ts

@@ -0,0 +1,14 @@
+export {
+  canonicalSsmTokenPath,
+  composeSsmParameterName,
+  delay,
+  errorName,
+  isRetryableProviderError,
+  positiveIntegerOption,
+  resolvePollingOptions,
+  throwIfCancelled,
+  validateConsumeOptions,
+  withCallDeadline,
+  type ResolvedRunnerConfigPollingOptions,
+  type RunnerConfigPollingOptions,
+} from '../../runner-config-consumer-common';

lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.test.ts

@@ -0,0 +1,220 @@
+import { DeleteParameterCommand, GetParameterCommand, type SSMClient } from '@aws-sdk/client-ssm';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  AwsSdkSsmRunnerConfigApi,
+  createAwsSsmRunnerConfigConsumer,
+  type AwsSsmRunnerConfigApi,
+} from './runner-config-consumer';
+
+function namedError(name: string, message = 'provider detail'): Error {
+  const error = new Error(message);
+  error.name = name;
+  return error;
+}
+
+describe('AWS SDK SSM runner config API', () => {
+  it('decrypts the parameter and deletes it with the caller abort signal', async () => {
+    const send = vi
+      .fn()
+      .mockResolvedValueOnce({ Parameter: { Value: 'encoded-jit' } })
+      .mockResolvedValueOnce({});
+    const api = new AwsSdkSsmRunnerConfigApi({ send } as unknown as SSMClient);
+    const signal = new AbortController().signal;
+
+    await expect(api.getParameter('/runner/tokens/runner-123', signal)).resolves.toBe('encoded-jit');
+    await expect(api.deleteParameter('/runner/tokens/runner-123', signal)).resolves.toBeUndefined();
+
+    expect(send.mock.calls[0][0]).toBeInstanceOf(GetParameterCommand);
+    expect(send.mock.calls[0][0].input).toEqual({
+      Name: '/runner/tokens/runner-123',
+      WithDecryption: true,
+    });
+    expect(send.mock.calls[0][1]).toEqual({ abortSignal: signal });
+    expect(send.mock.calls[1][0]).toBeInstanceOf(DeleteParameterCommand);
+    expect(send.mock.calls[1][0].input).toEqual({ Name: '/runner/tokens/runner-123' });
+    expect(send.mock.calls[1][1]).toEqual({ abortSignal: signal });
+  });
+});
+
+describe('SSM runner config consumer', () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it('polls a missing parameter, reads it, and deletes it before returning', async () => {
+    const getParameter = vi
+      .fn<AwsSsmRunnerConfigApi['getParameter']>()
+      .mockRejectedValueOnce(namedError('ParameterNotFound'))
+      .mockResolvedValueOnce('encoded-jit');
+    const deleteParameter = vi.fn<AwsSsmRunnerConfigApi['deleteParameter']>().mockResolvedValue(undefined);
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      {
+        api: { getParameter, deleteParameter },
+        callTimeoutMs: 100,
+        configTimeoutMs: 500,
+        pollIntervalMs: 1,
+      },
+    );
+
+    await expect(
+      consumer.consume('runner-123', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).resolves.toBe('encoded-jit');
+    expect(getParameter).toHaveBeenCalledTimes(2);
+    expect(deleteParameter).toHaveBeenCalledOnce();
+    expect(deleteParameter).toHaveBeenCalledWith('/runner/tokens/runner-123', expect.any(AbortSignal));
+  });
+
+  it('retries a transient delete failure without returning the value early', async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z'));
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockResolvedValue('encoded-jit'),
+      deleteParameter: vi
+        .fn()
+        .mockRejectedValueOnce(namedError('ThrottlingException'))
+        .mockResolvedValueOnce(undefined),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 100, configTimeoutMs: 2_000, deleteAttempts: 2, pollIntervalMs: 1 },
+    );
+
+    const pending = consumer.consume('runner-123', {
+      deadlineMs: Date.now() + 3_000,
+      signal: new AbortController().signal,
+    });
+    await vi.runAllTimersAsync();
+
+    await expect(pending).resolves.toBe('encoded-jit');
+    expect(api.deleteParameter).toHaveBeenCalledTimes(2);
+  });
+
+  it('fails closed when another reader deletes the SSM parameter first', async () => {
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockResolvedValue('encoded-jit'),
+      deleteParameter: vi.fn().mockRejectedValue(namedError('ParameterNotFound')),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 100, configTimeoutMs: 100, deleteAttempts: 3, pollIntervalMs: 1 },
+    );
+
+    await expect(
+      consumer.consume('runner-123', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).rejects.toThrow('runner configuration could not be deleted from SSM');
+    expect(api.deleteParameter).toHaveBeenCalledOnce();
+  });
+
+  it('sanitizes non-retryable provider failures', async () => {
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockRejectedValue(namedError('AccessDeniedException', 'encoded-jit-secret')),
+      deleteParameter: vi.fn(),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 },
+    );
+
+    const pending = consumer.consume('runner-123', {
+      deadlineMs: Date.now() + 1_000,
+      signal: new AbortController().signal,
+    });
+    await expect(pending).rejects.toThrow('failed to read runner configuration from SSM');
+    await expect(pending).rejects.not.toThrow('encoded-jit-secret');
+    expect(api.deleteParameter).not.toHaveBeenCalled();
+  });
+
+  it('rejects an empty SSM parameter value without attempting deletion', async () => {
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockResolvedValue(''),
+      deleteParameter: vi.fn(),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 },
+    );
+
+    await expect(
+      consumer.consume('runner-123', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).rejects.toThrow('failed to read runner configuration from SSM');
+    expect(api.deleteParameter).not.toHaveBeenCalled();
+  });
+
+  it('validates the full parameter name before calling SSM', async () => {
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn(),
+      deleteParameter: vi.fn(),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: `/${'x'.repeat(890)}` },
+      { api, callTimeoutMs: 100, configTimeoutMs: 100, pollIntervalMs: 1 },
+    );
+
+    await expect(
+      consumer.consume('runner-1234567890', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).rejects.toThrow('aws_ssm runner configuration key is invalid');
+    expect(api.getParameter).not.toHaveBeenCalled();
+  });
+
+  it('stops a provider call immediately when the caller aborts', async () => {
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockReturnValue(new Promise(() => undefined)),
+      deleteParameter: vi.fn(),
+    };
+    const controller = new AbortController();
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 10_000, configTimeoutMs: 10_000, pollIntervalMs: 1 },
+    );
+    const pending = consumer.consume('runner-123', {
+      deadlineMs: Date.now() + 10_000,
+      signal: controller.signal,
+    });
+
+    controller.abort();
+
+    await expect(pending).rejects.toThrow('runner configuration consumption was cancelled');
+  });
+
+  it('reserves a bounded delete attempt when the value appears near the polling deadline', async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z'));
+    const startedAt = Date.now();
+    let deleteStartedAt: number | undefined;
+    const api: AwsSsmRunnerConfigApi = {
+      getParameter: vi.fn().mockResolvedValueOnce(undefined).mockResolvedValueOnce('encoded-jit'),
+      deleteParameter: vi.fn().mockImplementation(async () => {
+        deleteStartedAt = Date.now();
+      }),
+    };
+    const consumer = createAwsSsmRunnerConfigConsumer(
+      { RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_ssm', SSM_TOKEN_PATH: '/runner/tokens' },
+      { api, callTimeoutMs: 100, configTimeoutMs: 1_000, pollIntervalMs: 99 },
+    );
+
+    const pending = consumer.consume('runner-123', {
+      deadlineMs: startedAt + 200,
+      signal: new AbortController().signal,
+    });
+    await vi.runAllTimersAsync();
+
+    await expect(pending).resolves.toBe('encoded-jit');
+    expect(api.getParameter).toHaveBeenCalledTimes(2);
+    expect(deleteStartedAt).toBe(startedAt + 99);
+    expect(deleteStartedAt).toBeLessThanOrEqual(startedAt + 100);
+  });
+});

lambdas/libs/storage-providers/aws/ssm/runner-config-consumer.ts

@@ -0,0 +1,163 @@
+import { DeleteParameterCommand, GetParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
+
+import type { RunnerConfigConsumer, RunnerConfigConsumeOptions } from '../../core';
+import {
+  composeSsmParameterName,
+  delay,
+  errorName,
+  isRetryableProviderError,
+  positiveIntegerOption,
+  resolvePollingOptions,
+  throwIfCancelled,
+  validateConsumeOptions,
+  withCallDeadline,
+  type RunnerConfigPollingOptions,
+} from './runner-config-consumer-common';
+
+const DEFAULT_DELETE_ATTEMPTS = 3;
+
+export interface AwsSsmRunnerConfigApi {
+  getParameter(name: string, signal: AbortSignal): Promise<string | undefined>;
+  deleteParameter(name: string, signal: AbortSignal): Promise<void>;
+}
+
+export class AwsSdkSsmRunnerConfigApi implements AwsSsmRunnerConfigApi {
+  private client?: SSMClient;
+
+  public constructor(client?: SSMClient) {
+    this.client = client;
+  }
+
+  private getClient(): SSMClient {
+    // Lifecycle hooks can be snapshotted before their first request. Constructing
+    // the untraced client here avoids persisting connection state in that snapshot.
+    this.client ??= new SSMClient({ maxAttempts: 1 });
+    return this.client;
+  }
+
+  public async getParameter(name: string, signal: AbortSignal): Promise<string | undefined> {
+    const response = await this.getClient().send(new GetParameterCommand({ Name: name, WithDecryption: true }), {
+      abortSignal: signal,
+    });
+    return response.Parameter?.Value;
+  }
+
+  public async deleteParameter(name: string, signal: AbortSignal): Promise<void> {
+    await this.getClient().send(new DeleteParameterCommand({ Name: name }), { abortSignal: signal });
+  }
+}
+
+export interface AwsSsmRunnerConfigConsumerOptions extends RunnerConfigPollingOptions {
+  api?: AwsSsmRunnerConfigApi;
+  deleteAttempts?: number;
+}
+
+export function createAwsSsmRunnerConfigConsumer(
+  environment: { SSM_TOKEN_PATH: string },
+  options: AwsSsmRunnerConfigConsumerOptions = {},
+): RunnerConfigConsumer {
+  return new AwsSsmRunnerConfigConsumer(
+    environment.SSM_TOKEN_PATH,
+    options.api ?? new AwsSdkSsmRunnerConfigApi(),
+    options,
+  );
+}
+
+class AwsSsmRunnerConfigConsumer implements RunnerConfigConsumer {
+  private readonly callTimeoutMs: number;
+  private readonly configTimeoutMs: number;
+  private readonly deleteAttempts: number;
+  private readonly pollIntervalMs: number;
+
+  public constructor(
+    private readonly tokenPath: string,
+    private readonly api: AwsSsmRunnerConfigApi,
+    options: AwsSsmRunnerConfigConsumerOptions,
+  ) {
+    const polling = resolvePollingOptions(options);
+    this.callTimeoutMs = polling.callTimeoutMs;
+    this.configTimeoutMs = polling.configTimeoutMs;
+    this.pollIntervalMs = polling.pollIntervalMs;
+    this.deleteAttempts = positiveIntegerOption('deleteAttempts', options.deleteAttempts, DEFAULT_DELETE_ATTEMPTS);
+  }
+
+  public async consume(runnerId: string, options: RunnerConfigConsumeOptions): Promise<string> {
+    validateConsumeOptions(options);
+    const parameterName = composeSsmParameterName(this.tokenPath, runnerId);
+    const startedAt = Date.now();
+    const remainingMs = Math.max(0, options.deadlineMs - startedAt);
+    // Preserve enough of short hook budgets for at least one bounded delete
+    // attempt without reviving the old fixed reserve that could consume the
+    // entire polling window.
+    const deleteReserveMs = Math.min(this.callTimeoutMs, Math.max(1, Math.floor(remainingMs / 2)));
+    const pollDeadline = Math.min(startedAt + this.configTimeoutMs, options.deadlineMs - deleteReserveMs);
+    let runnerConfig: string | undefined;
+
+    while (Date.now() < pollDeadline) {
+      throwIfCancelled(options.signal);
+      try {
+        runnerConfig = await this.read(parameterName, pollDeadline, options.signal);
+        if (runnerConfig !== undefined) {
+          if (runnerConfig.length === 0) {
+            throw new Error('runner configuration record has an invalid value');
+          }
+          break;
+        }
+      } catch (error) {
+        if (options.signal.aborted) {
+          throw new Error('runner configuration consumption was cancelled');
+        }
+        if (!isSsmNotFound(error) && !isRetryableProviderError(error)) {
+          throw new Error('failed to read runner configuration from SSM');
+        }
+      }
+
+      const remaining = pollDeadline - Date.now();
+      if (remaining > 0) {
+        await delay(Math.min(this.pollIntervalMs, remaining), options.signal);
+      }
+    }
+
+    if (runnerConfig === undefined) {
+      throw new Error('runner configuration did not become available before the deadline');
+    }
+
+    await this.delete(parameterName, options);
+    return runnerConfig;
+  }
+
+  private async read(name: string, deadlineMs: number, signal: AbortSignal): Promise<string | undefined> {
+    return withCallDeadline(signal, deadlineMs, this.callTimeoutMs, (callSignal) =>
+      this.api.getParameter(name, callSignal),
+    );
+  }
+
+  private async delete(name: string, options: RunnerConfigConsumeOptions): Promise<void> {
+    for (let attempt = 1; attempt <= this.deleteAttempts; attempt += 1) {
+      try {
+        await withCallDeadline(options.signal, options.deadlineMs, this.callTimeoutMs, (callSignal) =>
+          this.api.deleteParameter(name, callSignal),
+        );
+        return;
+      } catch (error) {
+        if (options.signal.aborted) {
+          throw new Error('runner configuration consumption was cancelled');
+        }
+        if (!isRetryableProviderError(error) || attempt === this.deleteAttempts) {
+          break;
+        }
+
+        const remaining = options.deadlineMs - Date.now();
+        if (remaining <= 0) {
+          break;
+        }
+        await delay(Math.min(2 ** (attempt - 1) * 1_000, 5_000, remaining), options.signal);
+      }
+    }
+    throw new Error('runner configuration could not be deleted from SSM');
+  }
+}
+
+function isSsmNotFound(error: unknown): boolean {
+  return errorName(error) === 'ParameterNotFound';
+}