fix(cloud-security): retry transient AssumeRole failures (first-scan "could not assume role")

apps/api/src/cloud-security/providers/aws-security.service.ts

@@ -1,5 +1,6 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { AssumeRoleCommand, STSClient } from '@aws-sdk/client-sts';
+import { retryAssume } from '@trycompai/integration-platform';
 import {
   CostExplorerClient,
   GetCostAndUsageCommand,
@@ -514,12 +515,14 @@ export class AWSSecurityService {
       region,
       credentials: getAwsBaseCredentials(partition),
     });
-    const roleAssumerResp = await baseSts.send(
-      new AssumeRoleCommand({
-        RoleArn: roleAssumerArn,
-        RoleSessionName: 'CompRoleAssumer',
-        DurationSeconds: 3600,
-      }),
+    const roleAssumerResp = await retryAssume(() =>
+      baseSts.send(
+        new AssumeRoleCommand({
+          RoleArn: roleAssumerArn,
+          RoleSessionName: 'CompRoleAssumer',
+          DurationSeconds: 3600,
+        }),
+      ),
     );
 
     const roleAssumerCreds = roleAssumerResp.Credentials;
@@ -541,13 +544,15 @@ export class AWSSecurityService {
 
     this.logger.log(`Assuming customer role ${roleArn} in region ${region}`);
 
-    const customerResp = await roleAssumerSts.send(
-      new AssumeRoleCommand({
-        RoleArn: roleArn,
-        ExternalId: externalId,
-        RoleSessionName: sessionName ?? 'CompSecurityAudit',
-        DurationSeconds: 3600,
-      }),
+    const customerResp = await retryAssume(() =>
+      roleAssumerSts.send(
+        new AssumeRoleCommand({
+          RoleArn: roleArn,
+          ExternalId: externalId,
+          RoleSessionName: sessionName ?? 'CompSecurityAudit',
+          DurationSeconds: 3600,
+        }),
+      ),
     );
 
     const customerCreds = customerResp.Credentials;

packages/integration-platform/src/index.ts

@@ -145,6 +145,13 @@ export {
 } from './manifests/aws/credentials';
 export type { AwsEnvironment } from './manifests/aws/credentials';
 
+// Shared AWS STS AssumeRole retry (transient / IAM-eventual-consistency safe),
+// reused by the Cloud Tests scanner in apps/api.
+export {
+  retryAssume,
+  isRetryableAssumeError,
+} from './manifests/aws/checks/assume-retry';
+
 
 // API Response types (for frontend and API type sharing)
 export type {

packages/integration-platform/src/manifests/aws/checks/__tests__/assume-retry.test.ts

@@ -0,0 +1,89 @@
+import { describe, expect, it } from 'bun:test';
+import { isRetryableAssumeError, retryAssume } from '../assume-retry';
+
+// Fast, deterministic options for tests — never actually sleep.
+const fastOpts = {
+  sleep: async () => {},
+  random: () => 0,
+  baseDelayMs: 1,
+  maxDelayMs: 1,
+};
+
+function awsError(name: string, extra: Record<string, unknown> = {}): Error {
+  return Object.assign(new Error(`${name} occurred`), { name, ...extra });
+}
+
+describe('isRetryableAssumeError', () => {
+  it('retries IAM/STS eventual-consistency AccessDenied (the customer bug)', () => {
+    expect(isRetryableAssumeError(awsError('AccessDenied'))).toBe(true);
+    expect(isRetryableAssumeError(awsError('AccessDeniedException'))).toBe(true);
+  });
+
+  it('retries expired-token, throttling, 5xx and network errors', () => {
+    expect(isRetryableAssumeError(awsError('ExpiredToken'))).toBe(true);
+    expect(isRetryableAssumeError(awsError('ThrottlingException'))).toBe(true);
+    expect(
+      isRetryableAssumeError({ name: 'X', $metadata: { httpStatusCode: 503 } }),
+    ).toBe(true);
+    expect(isRetryableAssumeError(new Error('socket hang up'))).toBe(true);
+  });
+
+  it('does NOT retry hard configuration errors', () => {
+    expect(isRetryableAssumeError(awsError('ValidationError'))).toBe(false);
+    expect(
+      isRetryableAssumeError(new Error('Invalid IAM Role ARN format')),
+    ).toBe(false);
+    expect(isRetryableAssumeError(null)).toBe(false);
+    expect(isRetryableAssumeError(undefined)).toBe(false);
+  });
+});
+
+describe('retryAssume', () => {
+  it('retries a transient AccessDenied then succeeds (no AWS change needed)', async () => {
+    let calls = 0;
+    const result = await retryAssume(async () => {
+      calls++;
+      if (calls < 3) throw awsError('AccessDenied');
+      return 'creds';
+    }, fastOpts);
+    expect(result).toBe('creds');
+    expect(calls).toBe(3);
+  });
+
+  it('does not retry a non-transient error (single attempt)', async () => {
+    let calls = 0;
+    await expect(
+      retryAssume(async () => {
+        calls++;
+        throw awsError('ValidationError');
+      }, fastOpts),
+    ).rejects.toThrow();
+    expect(calls).toBe(1);
+  });
+
+  it('gives up after the configured attempts on a persistent transient error', async () => {
+    let calls = 0;
+    await expect(
+      retryAssume(
+        async () => {
+          calls++;
+          throw awsError('AccessDenied');
+        },
+        { ...fastOpts, attempts: 4 },
+      ),
+    ).rejects.toThrow();
+    expect(calls).toBe(4);
+  });
+
+  it('succeeds on the first attempt without sleeping', async () => {
+    let slept = false;
+    const result = await retryAssume(async () => 'ok', {
+      ...fastOpts,
+      sleep: async () => {
+        slept = true;
+      },
+    });
+    expect(result).toBe('ok');
+    expect(slept).toBe(false);
+  });
+});

packages/integration-platform/src/manifests/aws/checks/assume-retry.ts

@@ -0,0 +1,106 @@
+/**
+ * Bounded retry for AWS STS `AssumeRole` calls.
+ *
+ * Why this exists: assuming a customer's cross-account role is a two-hop STS
+ * flow, and the FIRST attempt right after a connection is set up (or its trust
+ * policy edited) can fail transiently with NO change in AWS — IAM/STS is
+ * eventually consistent, so a brand-new/edited role can return `AccessDenied`
+ * on the first assume and then succeed seconds later. The base/roleAssumer
+ * session can likewise briefly return `ExpiredToken`. None of these classes are
+ * retried by the AWS SDK's default strategy (it only retries throttling/5xx/
+ * network), so a single transient failure surfaced to customers as a sticky
+ * "Could not assume AWS role" finding that "fixed itself" on the next run.
+ *
+ * This helper retries only those transient/eventually-consistent classes with
+ * capped exponential backoff + jitter. Hard configuration errors (bad ARN,
+ * ValidationError) and a role that is *persistently* denied still propagate —
+ * the backoff is capped so a genuinely broken role fails within a few seconds.
+ */
+
+const RETRYABLE_ASSUME_ERROR_NAMES = new Set<string>([
+  // IAM/STS eventual consistency on a new or just-edited cross-account role.
+  'AccessDenied',
+  'AccessDeniedException',
+  // Base/roleAssumer session momentarily expired, or an STS hiccup.
+  'ExpiredToken',
+  'ExpiredTokenException',
+  'IDPCommunicationError',
+  'RegionDisabledException',
+  // Standard transient classes (belt-and-suspenders alongside the SDK).
+  'Throttling',
+  'ThrottlingException',
+  'TooManyRequestsException',
+  'RequestLimitExceeded',
+  'ServiceUnavailable',
+  'ServiceUnavailableException',
+  'InternalError',
+  'InternalFailure',
+]);
+
+export function isRetryableAssumeError(error: unknown): boolean {
+  if (!error || typeof error !== 'object') return false;
+  const e = error as {
+    name?: string;
+    Code?: string;
+    code?: string;
+    __type?: string;
+    message?: string;
+    $metadata?: { httpStatusCode?: number };
+  };
+  const name = e.name ?? e.Code ?? e.code ?? e.__type ?? '';
+  if (RETRYABLE_ASSUME_ERROR_NAMES.has(name)) return true;
+
+  const status = e.$metadata?.httpStatusCode;
+  if (typeof status === 'number' && status >= 500) return true;
+
+  return /ECONNRESET|ETIMEDOUT|ENOTFOUND|EAI_AGAIN|socket hang up|network|timeout/i.test(
+    e.message ?? '',
+  );
+}
+
+export interface RetryAssumeOptions {
+  /** Total attempts including the first (default 4). */
+  attempts?: number;
+  /** Base backoff in ms; grows exponentially, capped by maxDelayMs (default 500). */
+  baseDelayMs?: number;
+  /** Maximum backoff per attempt in ms (default 5000). */
+  maxDelayMs?: number;
+  /** Injectable sleep so tests don't actually wait. */
+  sleep?: (ms: number) => Promise<void>;
+  /** Injectable jitter source in [0, 1); defaults to Math.random. */
+  random?: () => number;
+}
+
+const defaultSleep = (ms: number): Promise<void> =>
+  new Promise((resolve) => setTimeout(resolve, ms));
+
+/**
+ * Run `fn` (an AssumeRole call) with bounded exponential backoff + jitter,
+ * retrying only transient/eventually-consistent failures. Non-transient errors
+ * propagate immediately; transient ones that never recover propagate after the
+ * attempts are exhausted.
+ */
+export async function retryAssume<T>(
+  fn: () => Promise<T>,
+  options: RetryAssumeOptions = {},
+): Promise<T> {
+  const attempts = Math.max(1, options.attempts ?? 4);
+  const baseDelayMs = options.baseDelayMs ?? 500;
+  const maxDelayMs = options.maxDelayMs ?? 5000;
+  const sleep = options.sleep ?? defaultSleep;
+  const random = options.random ?? Math.random;
+
+  let lastError: unknown;
+  for (let attempt = 0; attempt < attempts; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error;
+      const isLastAttempt = attempt === attempts - 1;
+      if (isLastAttempt || !isRetryableAssumeError(error)) throw error;
+      const backoff = Math.min(maxDelayMs, baseDelayMs * 2 ** attempt);
+      await sleep(backoff + Math.floor(backoff * 0.5 * random()));
+    }
+  }
+  throw lastError;
+}

packages/integration-platform/src/manifests/aws/checks/shared.ts

@@ -1,5 +1,6 @@
 import { AssumeRoleCommand, STSClient } from '@aws-sdk/client-sts';
 import type { CheckContext, FindingSeverity } from '../../../types';
+import { retryAssume } from './assume-retry';
 
 export interface AwsSession {
   credentials: {
@@ -113,12 +114,14 @@ export async function assumeAwsSession(
     region,
     credentials: awsBaseCredentials(partition),
   });
-  const assumerResp = await baseSts.send(
-    new AssumeRoleCommand({
-      RoleArn: roleAssumerArn,
-      RoleSessionName: 'CompRoleAssumer',
-      DurationSeconds: 3600,
-    }),
+  const assumerResp = await retryAssume(() =>
+    baseSts.send(
+      new AssumeRoleCommand({
+        RoleArn: roleAssumerArn,
+        RoleSessionName: 'CompRoleAssumer',
+        DurationSeconds: 3600,
+      }),
+    ),
   );
   const assumer = assumerResp.Credentials;
   if (
@@ -139,13 +142,15 @@ export async function assumeAwsSession(
       sessionToken: assumer.SessionToken,
     },
   });
-  const res = await assumerSts.send(
-    new AssumeRoleCommand({
-      RoleArn: roleArn,
-      ExternalId: externalId,
-      RoleSessionName: 'CompEvidenceCheck',
-      DurationSeconds: 3600,
-    }),
+  const res = await retryAssume(() =>
+    assumerSts.send(
+      new AssumeRoleCommand({
+        RoleArn: roleArn,
+        ExternalId: externalId,
+        RoleSessionName: 'CompEvidenceCheck',
+        DurationSeconds: 3600,
+      }),
+    ),
   );
   const c = res.Credentials;
   if (!c?.AccessKeyId || !c.SecretAccessKey || !c.SessionToken) return null;