fix(bedrock): build role-session Bedrock clients per request instead of filling the shared client cache

An assumed or keyless role resolves to a fresh STS session on every call,
so a session-keyed cache entry was never read again and pushed other
providers' warm clients out of the shared LRU. Long-lived saved keys stay
cached. The session type now derives from AwsSessionCredentials, and
authorization refusals carry a reason.

Author: waleedlatif1

apps/sim/lib/internal/llm/bedrock-credential.ts

@@ -6,11 +6,21 @@ import { recordCredentialAccess } from '@/lib/oauth/token-resolution'
 import { AWS_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types'
 import type { BedrockCredentialSession } from '@/providers/types'
 
-/** A saved AWS credential could not be used for a Bedrock request. */
+/**
+ * A saved AWS credential could not be used for a Bedrock request: the caller is
+ * not allowed to use it (`unauthorized`), or it is not a live AWS credential or
+ * AWS refused it (`unusable`).
+ */
 export class BedrockCredentialError extends Error {
-  constructor(message: string, options?: { cause?: unknown }) {
-    super(message, options)
+  readonly reason: 'unauthorized' | 'unusable'
+
+  constructor(
+    message: string,
+    options?: { reason?: 'unauthorized' | 'unusable'; cause?: unknown }
+  ) {
+    super(message, { cause: options?.cause })
     this.name = 'BedrockCredentialError'
+    this.reason = options?.reason ?? 'unusable'
   }
 }
 
@@ -50,15 +60,19 @@ export async function resolveBedrockCredentialSession({
   signal,
 }: ResolveBedrockCredentialParams): Promise<BedrockCredentialSession> {
   if (!actingUserId) {
-    throw new BedrockCredentialError('Using an AWS credential requires an authenticated user')
+    throw new BedrockCredentialError('Using an AWS credential requires an authenticated user', {
+      reason: 'unauthorized',
+    })
   }
 
   const access = await authorizeCredentialUseForAuth(
     { success: true, userId: actingUserId, authType: AuthType.INTERNAL_JWT },
     { credentialId, workflowId, workspaceId, callerUserId: actingUserId }
   )
   if (!access.ok) {
-    throw new BedrockCredentialError(access.error || 'Not authorized to use this AWS credential')
+    throw new BedrockCredentialError(access.error || 'Not authorized to use this AWS credential', {
+      reason: 'unauthorized',
+    })
   }
   if (
     access.credentialType !== 'service_account' ||
@@ -99,6 +113,5 @@ export async function resolveBedrockCredentialSession({
     accessKeyId: session.accessKeyId,
     secretAccessKey: session.secretAccessKey,
     ...(session.sessionToken ? { sessionToken: session.sessionToken } : {}),
-    ...(session.expiration ? { expiration: session.expiration } : {}),
   }
 }

apps/sim/providers/bedrock/index.test.ts

@@ -169,7 +169,6 @@ describe('bedrockProvider credential handling', () => {
         accessKeyId: 'ASIASESSIONKEY000001',
         secretAccessKey: 'session-secret',
         sessionToken: 'session-token',
-        expiration: new Date('2030-01-01T00:00:00Z'),
       },
     })
 
@@ -202,30 +201,29 @@ describe('bedrockProvider credential handling', () => {
     expect(mockGetAwsCredentials).not.toHaveBeenCalled()
   })
 
-  it('reuses one client per session and builds a new one when the session changes', async () => {
-    const session = {
-      credentialId: 'cred-1',
-      accessKeyId: 'ASIASESSIONKEY000001',
-      secretAccessKey: 'session-secret',
-      sessionToken: 'session-token',
-      expiration: new Date('2030-01-01T00:00:00Z'),
+  it("keeps other clients warm while a role credential's STS sessions come and go", async () => {
+    const inlineKeys = {
+      bedrockAccessKeyId: 'AKIAIOSFODNN7EXAMPLE',
+      bedrockSecretKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
     }
+    await bedrockProvider.executeRequest({ ...baseRequest, ...inlineKeys })
     vi.mocked(BedrockRuntimeClient).mockClear()
 
-    await bedrockProvider.executeRequest({ ...baseRequest, bedrockSession: session })
-    await bedrockProvider.executeRequest({ ...baseRequest, bedrockSession: { ...session } })
-    expect(BedrockRuntimeClient).toHaveBeenCalledTimes(1)
+    for (let call = 0; call < 1_000; call++) {
+      await bedrockProvider.executeRequest({
+        ...baseRequest,
+        bedrockSession: {
+          credentialId: 'cred-role',
+          accessKeyId: `ASIASESSIONKEY${call}`,
+          secretAccessKey: `session-secret-${call}`,
+          sessionToken: `session-token-${call}`,
+        },
+      })
+    }
+    vi.mocked(BedrockRuntimeClient).mockClear()
 
-    await bedrockProvider.executeRequest({
-      ...baseRequest,
-      bedrockSession: {
-        ...session,
-        accessKeyId: 'ASIASESSIONKEY000002',
-        sessionToken: 'next-token',
-        expiration: new Date('2030-01-01T01:00:00Z'),
-      },
-    })
-    expect(BedrockRuntimeClient).toHaveBeenCalledTimes(2)
+    await bedrockProvider.executeRequest({ ...baseRequest, ...inlineKeys })
+    expect(BedrockRuntimeClient).not.toHaveBeenCalled()
   })
 
   it('builds a new client when a long-lived saved key pair is corrected', async () => {

apps/sim/providers/bedrock/index.ts

@@ -125,28 +125,31 @@ function enrichLastModelSegmentFromBedrockResponse(
 
 /**
  * The credentials a Bedrock client signs with, and the part of its cache key
- * that names them. A saved credential's session is used alone, so selecting one
- * never falls back to inline keys or the deployment's own AWS credentials. Its
- * key names the credential and session expiry plus a digest of the key material,
- * so a new session or a corrected key pair gets a fresh client without the
- * cache key holding a secret.
+ * that names them, or `null` when the client must not be cached. A saved
+ * credential's session is used alone, so selecting one never falls back to
+ * inline keys or the deployment's own AWS credentials. An STS session (assumed
+ * or keyless role) is fresh on every resolve, so its client is built for the one
+ * request rather than filling the shared client cache with entries never read
+ * again. A long-lived saved key pair is keyed by credential plus a digest of the
+ * key material, so a corrected pair gets a fresh client without the cache key
+ * holding a secret.
  */
 function resolveBedrockClientCredentials(request: ProviderRequest): {
   credentials: BedrockRuntimeClientConfig['credentials']
-  credentialKey: string
+  credentialKey: string | null
 } {
   const session = request.bedrockSession
   if (session) {
     const { accessKeyId, secretAccessKey, sessionToken } = session
+    if (sessionToken) {
+      return { credentials: { accessKeyId, secretAccessKey, sessionToken }, credentialKey: null }
+    }
     return {
-      credentials: sessionToken
-        ? { accessKeyId, secretAccessKey, sessionToken }
-        : { accessKeyId, secretAccessKey },
+      credentials: { accessKeyId, secretAccessKey },
       credentialKey: [
         'credential',
         session.credentialId,
-        session.expiration?.getTime() ?? 'long-lived',
-        sha256Hex([accessKeyId, secretAccessKey, sessionToken ?? ''].join('\0')),
+        sha256Hex([accessKeyId, secretAccessKey].join('\0')),
       ].join(':'),
     }
   }
@@ -204,10 +207,12 @@ export const bedrockProvider: ProviderConfig = {
     })
 
     const { credentials, credentialKey } = resolveBedrockClientCredentials(request)
-    const client = getCachedProviderClient(
-      `bedrock::${region}::${credentialKey}`,
-      () => new BedrockRuntimeClient({ region, credentials, requestHandler: awsRequestHandler() })
-    )
+    const buildClient = () =>
+      new BedrockRuntimeClient({ region, credentials, requestHandler: awsRequestHandler() })
+    const client =
+      credentialKey === null
+        ? buildClient()
+        : getCachedProviderClient(`bedrock::${region}::${credentialKey}`, buildClient)
 
     const { messages, systemContent } = convertBedrockRequestHistory(request)
     const canonicalModelId = `bedrock/${getBedrockBaseModelId(request.model)}`

apps/sim/providers/types.ts

@@ -1,5 +1,6 @@
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution'
 import { UserFailure } from '@/lib/core/errors/user-failure'
+import type { AwsSessionCredentials } from '@/lib/credentials/aws-service-account.server'
 import type { CustomBlockInputFieldType } from '@/blocks/custom/build-config'
 import type { ProviderTimingSegment, StreamingExecution, UserFile } from '@/executor/types'
 
@@ -193,17 +194,16 @@ export interface EvaluationInput {
   questions: unknown
 }
 
-/** Signing credentials resolved from a saved AWS credential for one Bedrock request. */
-export interface BedrockCredentialSession {
-  /** The saved credential the session came from; scopes the provider client cache. */
-  credentialId: string
-  accessKeyId: string
-  secretAccessKey: string
-  /** Present for STS sessions (assumed and keyless roles); absent for long-lived keys. */
-  sessionToken?: string
-  /** When an STS session ends; absent for long-lived keys. */
-  expiration?: Date
-}
+/**
+ * Signing credentials resolved from a saved AWS credential for one Bedrock
+ * request. `sessionToken` is present for STS sessions (assumed and keyless
+ * roles) and absent for long-lived keys; `credentialId` names the saved
+ * credential the session came from.
+ */
+export type BedrockCredentialSession = Pick<
+  AwsSessionCredentials,
+  'accessKeyId' | 'secretAccessKey' | 'sessionToken'
+> & { credentialId: string }
 
 export interface ProviderRequest {
   evaluation?: EvaluationInput