From a47209284026d3d59c8cf18b3e6635e6b1fd7d7b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 12 Nov 2026 16:52:10 -0800 Subject: [PATCH] fix(executor): resolve saved model credentials inside each fallback attempt in Router and Evaluator --- apps/docs/openapi-v2-knowledge.json | 4 +- apps/sim/executor/handlers/agent/handler.ts | 33 ++++------ .../handlers/evaluator/handler.test.ts | 45 +++++++++++++ .../executor/handlers/evaluator/handler.ts | 19 +----- .../executor/handlers/router/handler.test.ts | 48 ++++++++++++++ apps/sim/executor/handlers/router/handler.ts | 36 ++--------- apps/sim/executor/utils/bedrock-credential.ts | 29 --------- apps/sim/executor/utils/model-credentials.ts | 64 +++++++++++++++++++ .../utils/model-fallback-request.test.ts | 63 ++++++++++++++++++ .../executor/utils/model-fallback-request.ts | 21 ++++++ apps/sim/lib/api/contracts/v2/knowledge.ts | 2 +- 11 files changed, 263 insertions(+), 101 deletions(-) delete mode 100644 apps/sim/executor/utils/bedrock-credential.ts create mode 100644 apps/sim/executor/utils/model-credentials.ts diff --git a/apps/docs/openapi-v2-knowledge.json b/apps/docs/openapi-v2-knowledge.json index c98a9c09f81..b490259f54c 100644 --- a/apps/docs/openapi-v2-knowledge.json +++ b/apps/docs/openapi-v2-knowledge.json @@ -6933,7 +6933,7 @@ "type": "null" } ], - "description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for API-key and unauthenticated sources." + "description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for unauthenticated sources and for `apiKey` sources that use an inline API key." }, "sourceConfig": { "type": "object", @@ -7331,7 +7331,7 @@ "type": "null" } ], - "description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for API-key and unauthenticated sources." + "description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for unauthenticated sources and for `apiKey` sources that use an inline API key." }, "sourceConfig": { "type": "object", diff --git a/apps/sim/executor/handlers/agent/handler.ts b/apps/sim/executor/handlers/agent/handler.ts index 87b57561bf2..03da301096e 100644 --- a/apps/sim/executor/handlers/agent/handler.ts +++ b/apps/sim/executor/handlers/agent/handler.ts @@ -91,9 +91,12 @@ import type { StreamingExecution, UserFile, } from '@/executor/types' -import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential' import { collectBlockData } from '@/executor/utils/block-data' import { stringifyJSON } from '@/executor/utils/json' +import { + BLOCK_MODEL_CREDENTIAL_FIELDS, + resolveBlockModelCredentials, +} from '@/executor/utils/model-credentials' import { assertFallbackFamilyCredentials, getModelFallbacks, @@ -109,7 +112,6 @@ import type { ResolvedSecretTraceRegistry, } from '@/executor/utils/resolved-secret-trace-registry' import { annotateDuplicateToolBindings } from '@/executor/utils/tool-binding-labels' -import { resolveVertexCredential } from '@/executor/utils/vertex-credential' import { executeProviderRequest } from '@/providers' import { formatAttachmentSizes, @@ -2710,9 +2712,7 @@ export class AgentBlockHandler implements BlockHandler { ? config.modelInputs : omit(config.modelInputs, [ ...PROVIDER_FAMILY_CREDENTIAL_FIELDS, - 'vertexCredential', - 'bedrockCredential', - 'bedrockAuthMode', + ...BLOCK_MODEL_CREDENTIAL_FIELDS, ])), apiKey: resolveFallbackApiKey({ candidate, @@ -2989,19 +2989,12 @@ export class AgentBlockHandler implements BlockHandler { const model = providerRequest.model try { - let finalApiKey: string | undefined = providerRequest.apiKey - - if (providerId === 'vertex' && providerRequest.vertexCredential) { - finalApiKey = await resolveVertexCredential({ - credentialId: providerRequest.vertexCredential, - actingUserId: ctx.userId, - workspaceId: ctx.workspaceId, - workflowId: ctx.workflowId, - executionId: ctx.executionId, - callerLabel: 'vertex-agent', - }) - } - const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, providerRequest) + const credentials = await resolveBlockModelCredentials( + ctx, + providerId, + providerRequest, + 'vertex-agent' + ) const { blockData, blockNameMapping } = collectBlockData(ctx) const agentMemoryRetrieval = agentConversation?.memoryId @@ -3024,7 +3017,7 @@ export class AgentBlockHandler implements BlockHandler { : providerRequest.tools, temperature: providerRequest.temperature, maxTokens: providerRequest.maxTokens, - apiKey: finalApiKey, + apiKey: credentials.apiKey ?? providerRequest.apiKey, azureEndpoint: providerRequest.azureEndpoint, azureApiVersion: providerRequest.azureApiVersion, vertexProject: providerRequest.vertexProject, @@ -3032,7 +3025,7 @@ export class AgentBlockHandler implements BlockHandler { bedrockAccessKeyId: providerRequest.bedrockAccessKeyId, bedrockSecretKey: providerRequest.bedrockSecretKey, bedrockRegion: providerRequest.bedrockRegion, - bedrockSession, + bedrockSession: credentials.bedrockSession, responseFormat: providerRequest.responseFormat, workflowId: providerRequest.workflowId, workspaceId: ctx.workspaceId, diff --git a/apps/sim/executor/handlers/evaluator/handler.test.ts b/apps/sim/executor/handlers/evaluator/handler.test.ts index 68c2c4b3e19..c443c1efcab 100644 --- a/apps/sim/executor/handlers/evaluator/handler.test.ts +++ b/apps/sim/executor/handlers/evaluator/handler.test.ts @@ -5,6 +5,7 @@ import { import { llmBedrockCredentialMock, llmBedrockCredentialMockFns, + MockBedrockCredentialError, } from '@sim/testing/mocks/llm-bedrock-credential.mock' import { permissionsMock, permissionsMockFns } from '@sim/testing/mocks/permissions.mock' import '@sim/testing/mocks/executor' @@ -252,6 +253,50 @@ describe('EvaluatorBlockHandler', () => { expect(JSON.stringify(fallbackRequest)).not.toContain('saved-secret') }) + it('scores with a fallback model when the saved AWS credential is unavailable on the final try', async () => { + mockGetProviderFromModel.mockImplementation((model: string) => + model.startsWith('claude') ? 'anthropic' : 'bedrock' + ) + mockResolveBedrockCredentialSession.mockRejectedValueOnce( + new MockBedrockCredentialError('AWS STS is unavailable', { reason: 'unavailable' }) + ) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: '{"score1":7}', + model: 'claude-sonnet-5', + tokens: { input: 12, output: 3, total: 15 }, + cost: 0, + }) + mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry() + mockContext.blockLogs = [ + { + blockId: mockBlock.id, + startedAt: '', + endedAt: '', + durationMs: 0, + success: false, + executionOrder: 1, + }, + ] + + const output = await handler.execute( + mockContext, + mockBlock, + { + ...admissionInputs, + model: 'bedrock/claude', + bedrockAuthMode: 'credential', + bedrockCredential: 'aws-cred', + fallbackModels: [{ model: 'claude-sonnet-5' }], + }, + { nodeId: mockBlock.id, retry: { attempt: 2, maxTries: 2, isFinalTry: true } } + ) + + expect(output).toMatchObject({ model: 'claude-sonnet-5', score1: 7 }) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) + expect(mockExecuteProviderRequest.mock.calls[0][0]).toBe('anthropic') + expect(mockContext.blockLogs[0].modelFallbacks).toEqual(['bedrock/claude']) + }) + it('fails a Bedrock evaluation whose AWS credential cannot be used instead of sending it', async () => { mockGetProviderFromModel.mockReturnValue('bedrock') mockResolveBedrockCredentialSession.mockRejectedValueOnce( diff --git a/apps/sim/executor/handlers/evaluator/handler.ts b/apps/sim/executor/handlers/evaluator/handler.ts index 66d2da8be0f..7eddaf327f5 100644 --- a/apps/sim/executor/handlers/evaluator/handler.ts +++ b/apps/sim/executor/handlers/evaluator/handler.ts @@ -10,7 +10,6 @@ import type { BlockOutput } from '@/blocks/types' import { validateModelProvider } from '@/ee/access-control/utils/permission-check' import { BlockType, DEFAULTS, EVALUATOR } from '@/executor/constants' import type { BlockHandler, BlockNodeMetadata, ExecutionContext } from '@/executor/types' -import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential' import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json' import { executeModelRequestWithFallbacks } from '@/executor/utils/model-fallback-request' import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection' @@ -19,7 +18,6 @@ import type { ResolvedSecretInputPath, ResolvedSecretTraceRegistry, } from '@/executor/utils/resolved-secret-trace-registry' -import { resolveVertexCredential } from '@/executor/utils/vertex-credential' import { resolveProxiedModelCost } from '@/providers/cost-policy' import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models' import type { ProviderRequest } from '@/providers/types' @@ -181,19 +179,6 @@ export class EvaluatorBlockHandler implements BlockHandler { await validateModelProvider(ctx.userId, ctx.workspaceId, model, ctx) const providerId = getProviderFromModel(model) - let finalApiKey: string | undefined = evaluatorConfig.apiKey - if (providerId === 'vertex' && evaluatorConfig.vertexCredential) { - finalApiKey = await resolveVertexCredential({ - credentialId: evaluatorConfig.vertexCredential, - actingUserId: ctx.userId, - workspaceId: ctx.workspaceId, - workflowId: ctx.workflowId, - executionId: ctx.executionId, - callerLabel: 'vertex-evaluator', - }) - } - const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, evaluatorConfig) - try { const providerRequest: ProviderRequest = { model, @@ -208,7 +193,7 @@ export class EvaluatorBlockHandler implements BlockHandler { ]), temperature: EVALUATOR.DEFAULT_TEMPERATURE, - apiKey: finalApiKey, + apiKey: evaluatorConfig.apiKey, azureEndpoint: inputs.azureEndpoint, azureApiVersion: inputs.azureApiVersion, vertexProject: evaluatorConfig.vertexProject, @@ -216,7 +201,6 @@ export class EvaluatorBlockHandler implements BlockHandler { bedrockAccessKeyId: evaluatorConfig.bedrockAccessKeyId, bedrockSecretKey: evaluatorConfig.bedrockSecretKey, bedrockRegion: evaluatorConfig.bedrockRegion, - bedrockSession, workflowId: ctx.workflowId, workspaceId: ctx.workspaceId, } @@ -231,6 +215,7 @@ export class EvaluatorBlockHandler implements BlockHandler { providerId, request: providerRequest, resolvedSecretTraceRegistry: modelInputProjection.registry, + credentials: { inputs: evaluatorConfig, callerLabel: 'vertex-evaluator' }, }) const parsedContent = this.extractJSONFromResponse( diff --git a/apps/sim/executor/handlers/router/handler.test.ts b/apps/sim/executor/handlers/router/handler.test.ts index c095c2be3d6..2230532f119 100644 --- a/apps/sim/executor/handlers/router/handler.test.ts +++ b/apps/sim/executor/handlers/router/handler.test.ts @@ -5,6 +5,7 @@ import { import { llmBedrockCredentialMock, llmBedrockCredentialMockFns, + MockBedrockCredentialError, } from '@sim/testing/mocks/llm-bedrock-credential.mock' import { permissionsMock, permissionsMockFns } from '@sim/testing/mocks/permissions.mock' import '@sim/testing/mocks/executor' @@ -223,6 +224,53 @@ describe('RouterBlockHandler', () => { expect(mockExecuteProviderRequest).not.toHaveBeenCalled() }) + it('routes with a fallback model when the saved AWS credential is unavailable on the final try', async () => { + mockGetProviderFromModel.mockImplementation((model: string) => + model.startsWith('claude') ? 'anthropic' : 'bedrock' + ) + mockResolveBedrockCredentialSession.mockRejectedValueOnce( + new MockBedrockCredentialError('AWS STS is unavailable', { reason: 'unavailable' }) + ) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'target-block-1', + model: 'claude-sonnet-5', + tokens: { input: 10, output: 2, total: 12 }, + cost: 0, + }) + mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry() + mockContext.blockLogs = [ + { + blockId: mockBlock.id, + startedAt: '', + endedAt: '', + durationMs: 0, + success: false, + executionOrder: 1, + }, + ] + + const output = await handler.execute( + mockContext, + mockBlock, + { + prompt: 'Pick a destination', + model: 'bedrock/claude', + bedrockAuthMode: 'credential', + bedrockCredential: 'aws-cred', + fallbackModels: [{ model: 'claude-sonnet-5' }], + }, + { nodeId: mockBlock.id, retry: { attempt: 2, maxTries: 2, isFinalTry: true } } + ) + + expect(output).toMatchObject({ + model: 'claude-sonnet-5', + selectedPath: { blockId: 'target-block-1' }, + }) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) + expect(mockExecuteProviderRequest.mock.calls[0][0]).toBe('anthropic') + expect(mockContext.blockLogs[0].modelFallbacks).toEqual(['bedrock/claude']) + }) + it('selects the same legacy destination when a fallback provider answers', async () => { mockExecuteProviderRequest .mockRejectedValueOnce(new Error('overloaded')) diff --git a/apps/sim/executor/handlers/router/handler.ts b/apps/sim/executor/handlers/router/handler.ts index ab13af804fe..784bc35620d 100644 --- a/apps/sim/executor/handlers/router/handler.ts +++ b/apps/sim/executor/handlers/router/handler.ts @@ -17,11 +17,9 @@ import { ROUTER, } from '@/executor/constants' import type { BlockHandler, BlockNodeMetadata, ExecutionContext } from '@/executor/types' -import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential' import { executeModelRequestWithFallbacks } from '@/executor/utils/model-fallback-request' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry' -import { resolveVertexCredential } from '@/executor/utils/vertex-credential' import { resolveProxiedModelCost } from '@/providers/cost-policy' import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models' import type { ProviderRequest } from '@/providers/types' @@ -116,25 +114,12 @@ export class RouterBlockHandler implements BlockHandler { await validateModelProvider(ctx.userId, ctx.workspaceId, resolved.model, ctx) const providerId = getProviderFromModel(resolved.model) - let finalApiKey: string | undefined = routerConfig.apiKey - if (providerId === 'vertex' && routerConfig.vertexCredential) { - finalApiKey = await resolveVertexCredential({ - credentialId: routerConfig.vertexCredential, - actingUserId: ctx.userId, - workspaceId: ctx.workspaceId, - workflowId: ctx.workflowId, - executionId: ctx.executionId, - callerLabel: 'vertex-router', - }) - } - const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, routerConfig) - const providerRequest: ProviderRequest = { model: resolved.model, systemPrompt: resolved.systemPrompt, context: JSON.stringify(messages), temperature: ROUTER.INFERENCE_TEMPERATURE, - apiKey: finalApiKey, + apiKey: routerConfig.apiKey, azureEndpoint: inputs.azureEndpoint, azureApiVersion: inputs.azureApiVersion, vertexProject: routerConfig.vertexProject, @@ -142,7 +127,6 @@ export class RouterBlockHandler implements BlockHandler { bedrockAccessKeyId: routerConfig.bedrockAccessKeyId, bedrockSecretKey: routerConfig.bedrockSecretKey, bedrockRegion: routerConfig.bedrockRegion, - bedrockSession, workflowId: ctx.workflowId, workspaceId: ctx.workspaceId, } @@ -157,6 +141,7 @@ export class RouterBlockHandler implements BlockHandler { providerId, request: providerRequest, resolvedSecretTraceRegistry: modelInputProjection.registry, + credentials: { inputs: routerConfig, callerLabel: 'vertex-router' }, }) const chosenBlockId = result.content.trim().toLowerCase() @@ -290,25 +275,12 @@ export class RouterBlockHandler implements BlockHandler { await validateModelProvider(ctx.userId, ctx.workspaceId, resolved.model, ctx) const providerId = getProviderFromModel(resolved.model) - let finalApiKey: string | undefined = routerConfig.apiKey - if (providerId === 'vertex' && routerConfig.vertexCredential) { - finalApiKey = await resolveVertexCredential({ - credentialId: routerConfig.vertexCredential, - actingUserId: ctx.userId, - workspaceId: ctx.workspaceId, - workflowId: ctx.workflowId, - executionId: ctx.executionId, - callerLabel: 'vertex-router', - }) - } - const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, routerConfig) - const providerRequest: ProviderRequest = { model: resolved.model, systemPrompt: resolved.systemPrompt, context: JSON.stringify(messages), temperature: ROUTER.INFERENCE_TEMPERATURE, - apiKey: finalApiKey, + apiKey: routerConfig.apiKey, azureEndpoint: inputs.azureEndpoint, azureApiVersion: inputs.azureApiVersion, vertexProject: routerConfig.vertexProject, @@ -316,7 +288,6 @@ export class RouterBlockHandler implements BlockHandler { bedrockAccessKeyId: routerConfig.bedrockAccessKeyId, bedrockSecretKey: routerConfig.bedrockSecretKey, bedrockRegion: routerConfig.bedrockRegion, - bedrockSession, workflowId: ctx.workflowId, workspaceId: ctx.workspaceId, responseFormat: { @@ -350,6 +321,7 @@ export class RouterBlockHandler implements BlockHandler { providerId, request: providerRequest, resolvedSecretTraceRegistry: modelInputProjection.registry, + credentials: { inputs: routerConfig, callerLabel: 'vertex-router' }, }) let chosenRouteId: string diff --git a/apps/sim/executor/utils/bedrock-credential.ts b/apps/sim/executor/utils/bedrock-credential.ts deleted file mode 100644 index c6225bebea5..00000000000 --- a/apps/sim/executor/utils/bedrock-credential.ts +++ /dev/null @@ -1,29 +0,0 @@ -import { resolveBedrockCredentialSession } from '@/lib/internal/llm/bedrock-credential' -import { getSelectedBedrockCredential } from '@/blocks/utils' -import type { ExecutionContext } from '@/executor/types' -import type { BedrockCredentialSession } from '@/providers/types' - -/** - * The saved AWS credential session a block's Bedrock call signs with, or - * undefined when the call is not to Bedrock or the block uses access keys. - * Shared by the agent, router, and evaluator handlers. Resolution runs for - * the run's acting user within the workflow's workspace, and a failure fails - * the call rather than letting it fall back to other credentials. - */ -export async function resolveBlockBedrockSession( - ctx: ExecutionContext, - providerId: string, - inputs: Parameters[0] -): Promise { - if (providerId !== 'bedrock') return undefined - const credentialId = getSelectedBedrockCredential(inputs) - if (!credentialId) return undefined - return resolveBedrockCredentialSession({ - credentialId, - actingUserId: ctx.userId, - workflowId: ctx.workflowId, - workspaceId: ctx.workspaceId, - executionId: ctx.executionId, - signal: ctx.abortSignal, - }) -} diff --git a/apps/sim/executor/utils/model-credentials.ts b/apps/sim/executor/utils/model-credentials.ts new file mode 100644 index 00000000000..943ca4bd258 --- /dev/null +++ b/apps/sim/executor/utils/model-credentials.ts @@ -0,0 +1,64 @@ +import { resolveBedrockCredentialSession } from '@/lib/internal/llm/bedrock-credential' +import { getSelectedBedrockCredential } from '@/blocks/utils' +import type { ExecutionContext } from '@/executor/types' +import { resolveVertexCredential } from '@/executor/utils/vertex-credential' +import type { BedrockCredentialSession } from '@/providers/types' + +/** + * The block fields that select a saved model credential. A fallback on another + * provider drops them, so one provider's credential never reaches another. + */ +export const BLOCK_MODEL_CREDENTIAL_FIELDS = [ + 'vertexCredential', + 'bedrockCredential', + 'bedrockAuthMode', +] as const + +/** A block's saved model credential selection. */ +export type BlockModelCredentialInputs = Parameters[0] & { + vertexCredential?: string +} + +/** What a saved credential adds to a provider request; a field is present only when it resolved. */ +export interface BlockModelCredentials { + apiKey?: string + bedrockSession?: BedrockCredentialSession +} + +/** + * Resolves the saved credential a block's model call signs with: a Vertex AI + * OAuth or service-account token (returned as `apiKey`), or a Bedrock AWS + * credential session. Resolution runs for the run's acting user within the + * workflow's workspace, and a failure fails that attempt rather than letting it + * fall back to other credentials. + */ +export async function resolveBlockModelCredentials( + ctx: ExecutionContext, + providerId: string, + inputs: BlockModelCredentialInputs, + callerLabel: string +): Promise { + if (providerId === 'vertex' && inputs.vertexCredential) { + const apiKey = await resolveVertexCredential({ + credentialId: inputs.vertexCredential, + actingUserId: ctx.userId, + workspaceId: ctx.workspaceId, + workflowId: ctx.workflowId, + executionId: ctx.executionId, + callerLabel, + }) + return { apiKey } + } + if (providerId !== 'bedrock') return {} + const credentialId = getSelectedBedrockCredential(inputs) + if (!credentialId) return {} + const bedrockSession = await resolveBedrockCredentialSession({ + credentialId, + actingUserId: ctx.userId, + workflowId: ctx.workflowId, + workspaceId: ctx.workspaceId, + executionId: ctx.executionId, + signal: ctx.abortSignal, + }) + return { bedrockSession } +} diff --git a/apps/sim/executor/utils/model-fallback-request.test.ts b/apps/sim/executor/utils/model-fallback-request.test.ts index ab25ce03736..045c046db46 100644 --- a/apps/sim/executor/utils/model-fallback-request.test.ts +++ b/apps/sim/executor/utils/model-fallback-request.test.ts @@ -1,4 +1,8 @@ import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { + llmBedrockCredentialMock, + llmBedrockCredentialMockFns, +} from '@sim/testing/mocks/llm-bedrock-credential.mock' import { permissionCheckMock, permissionCheckMockFns, @@ -11,8 +15,10 @@ import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-tr import type { SerializedBlock } from '@/serializer/types' const validateModel = permissionCheckMockFns.mockValidateModelProvider +const resolveBedrockSession = llmBedrockCredentialMockFns.mockResolveBedrockCredentialSession providersUtilsMockFns.mockGetProviderFromModel.mockImplementation((model: string) => { if (model.startsWith('claude')) return 'anthropic' + if (model.startsWith('bedrock/')) return 'bedrock' if (model.startsWith('vertex/')) return 'vertex' return 'openai' }) @@ -24,6 +30,7 @@ const { request } = vi.hoisted(() => ({ vi.mock('@/executor/utils/provider-request', () => ({ executeBlockProviderRequest: request })) vi.mock('@/ee/access-control/utils/permission-check', () => permissionCheckMock) vi.mock('@/providers/utils', () => providersUtilsMock) +vi.mock('@/lib/internal/llm/bedrock-credential', () => llmBedrockCredentialMock) function context(): ExecutionContext { return { @@ -246,6 +253,62 @@ describe('executeModelRequestWithFallbacks', () => { expect(request.mock.calls[1][0].request.apiKey).toBeUndefined() }) + it('reuses the saved AWS session for a same-provider fallback after the primary model fails', async () => { + const session = { + credentialId: 'aws-cred', + accessKeyId: 'ASIASESSIONKEY000001', + secretAccessKey: 'session-secret', + sessionToken: 'session-token', + } + resolveBedrockSession.mockResolvedValueOnce(session) + request.mockRejectedValueOnce(new Error('overloaded')) + + const output = await executeModelRequestWithFallbacks({ + ...input(), + providerId: 'bedrock', + configuredModel: 'bedrock/claude-sonnet', + request: { model: 'bedrock/claude-sonnet', systemPrompt: 'Score the content' }, + fallbackModels: [{ model: 'bedrock/claude-haiku' }], + credentials: { + inputs: { bedrockAuthMode: 'credential', bedrockCredential: 'aws-cred' }, + callerLabel: 'test', + }, + }) + + expect(output.result.model).toBe('bedrock/claude-haiku') + expect(resolveBedrockSession).toHaveBeenCalledTimes(1) + expect(request.mock.calls.map(([args]) => args.request.bedrockSession)).toEqual([ + session, + session, + ]) + }) + + it('fails every same-provider candidate after one failed credential resolution, then falls back without the session', async () => { + resolveBedrockSession.mockRejectedValueOnce(new Error('AWS STS is unavailable')) + const options = input() + + const output = await executeModelRequestWithFallbacks({ + ...options, + providerId: 'bedrock', + configuredModel: 'bedrock/claude-sonnet', + request: { model: 'bedrock/claude-sonnet', systemPrompt: 'Score the content' }, + fallbackModels: [{ model: 'bedrock/claude-haiku' }, { model: 'claude-sonnet-5' }], + credentials: { + inputs: { bedrockAuthMode: 'credential', bedrockCredential: 'aws-cred' }, + callerLabel: 'test', + }, + }) + + expect(output.result.model).toBe('claude-sonnet-5') + expect(resolveBedrockSession).toHaveBeenCalledTimes(1) + expect(request.mock.calls.map(([args]) => args.providerId)).toEqual(['anthropic']) + expect(request.mock.calls[0][0].request.bedrockSession).toBeUndefined() + expect(options.ctx.blockLogs[0].modelFallbacks).toEqual([ + 'bedrock/claude-sonnet', + 'bedrock/claude-haiku', + ]) + }) + it('strips the Auto identity prompt from fallbacks and keeps the pool model out of the trace', async () => { const options = input() request.mockRejectedValueOnce(new Error('overloaded')) diff --git a/apps/sim/executor/utils/model-fallback-request.ts b/apps/sim/executor/utils/model-fallback-request.ts index db364312b9e..e4687e58c90 100644 --- a/apps/sim/executor/utils/model-fallback-request.ts +++ b/apps/sim/executor/utils/model-fallback-request.ts @@ -4,6 +4,11 @@ import { resolveFallbackTuning } from '@/lib/workflows/blocks/fallback-models' import { validateModelProvider } from '@/ee/access-control/utils/permission-check' import { isRetryableBlockError } from '@/executor/execution/block-retry' import type { BlockRetryAttempt, ExecutionContext } from '@/executor/types' +import { + type BlockModelCredentialInputs, + type BlockModelCredentials, + resolveBlockModelCredentials, +} from '@/executor/utils/model-credentials' import { assertFallbackFamilyCredentials, getModelFallbacks, @@ -32,11 +37,16 @@ interface ModelFallbackRequestInput { fallbackSystemPrompt: string retry?: BlockRetryAttempt resolvedSecretTraceRegistry: ResolvedSecretTraceRegistry | undefined + /** The block's saved Vertex or Bedrock credential selection, resolved per attempt. */ + credentials?: { inputs: BlockModelCredentialInputs; callerLabel: string } } /** * Shared non-streaming model execution for Router and Evaluator. Block retries * exhaust the primary first; the final try walks the ordered fallback models. + * Saved credentials resolve inside the attempt, so one that cannot be resolved fails + * that attempt and the fallbacks still run. Only candidates on the primary's provider + * resolve them, once per call, and they never reach another provider. * Parsing a successful response remains the handler's job, so a routing decision * such as NO_MATCH is never silently replaced by another model's decision. */ @@ -50,6 +60,7 @@ export async function executeModelRequestWithFallbacks({ fallbackSystemPrompt, retry, resolvedSecretTraceRegistry, + credentials, }: ModelFallbackRequestInput): Promise<{ result: ProviderResponse; usedFallback: boolean }> { const fallbacks = getModelFallbacks(ctx, block, fallbackModels, logger) const candidates = [ @@ -62,6 +73,7 @@ export async function executeModelRequestWithFallbacks({ ] const failedModels: string[] = [] let lastError: unknown + let savedCredentials: Promise | undefined for (const [index, candidate] of candidates.entries()) { const isPrimary = index === 0 @@ -113,6 +125,15 @@ export async function executeModelRequestWithFallbacks({ } try { + if (credentials && candidateProviderId === providerId) { + savedCredentials ??= resolveBlockModelCredentials( + ctx, + providerId, + credentials.inputs, + credentials.callerLabel + ) + candidateRequest = { ...candidateRequest, ...(await savedCredentials) } + } const result = await executeBlockProviderRequest({ ctx, providerId: candidateProviderId, diff --git a/apps/sim/lib/api/contracts/v2/knowledge.ts b/apps/sim/lib/api/contracts/v2/knowledge.ts index a3bc5cb8108..e09ab863dca 100644 --- a/apps/sim/lib/api/contracts/v2/knowledge.ts +++ b/apps/sim/lib/api/contracts/v2/knowledge.ts @@ -1778,7 +1778,7 @@ export const v2KnowledgeConnectorSchema = z .string() .nullable() .describe( - 'OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for API-key and unauthenticated sources.' + 'OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for unauthenticated sources and for `apiKey` sources that use an inline API key.' ), sourceConfig: z .record(z.string(), z.unknown().describe('Connector-specific source configuration value.'))