Skip to content

Commit a472092

Browse files
waleedlatif1waleedlatif1
committed
fix(executor): resolve saved model credentials inside each fallback attempt in Router and Evaluator
1 parent 005cb97 commit a472092

11 files changed

Lines changed: 263 additions & 101 deletions

File tree

‎apps/docs/openapi-v2-knowledge.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6933,7 +6933,7 @@
69336933
"type": "null"
69346934
}
69356935
],
6936-
"description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for API-key and unauthenticated sources."
6936+
"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."
69376937
},
69386938
"sourceConfig": {
69396939
"type": "object",
@@ -7331,7 +7331,7 @@
73317331
"type": "null"
73327332
}
73337333
],
7334-
"description": "OAuth credential identifier, or the saved credential an `apiKey` source syncs with (see `auth.savedCredential`); null for API-key and unauthenticated sources."
7334+
"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."
73357335
},
73367336
"sourceConfig": {
73377337
"type": "object",

‎apps/sim/executor/handlers/agent/handler.ts‎

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,12 @@ import type {
9191
StreamingExecution,
9292
UserFile,
9393
} from '@/executor/types'
94-
import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential'
9594
import { collectBlockData } from '@/executor/utils/block-data'
9695
import { stringifyJSON } from '@/executor/utils/json'
96+
import {
97+
BLOCK_MODEL_CREDENTIAL_FIELDS,
98+
resolveBlockModelCredentials,
99+
} from '@/executor/utils/model-credentials'
97100
import {
98101
assertFallbackFamilyCredentials,
99102
getModelFallbacks,
@@ -109,7 +112,6 @@ import type {
109112
ResolvedSecretTraceRegistry,
110113
} from '@/executor/utils/resolved-secret-trace-registry'
111114
import { annotateDuplicateToolBindings } from '@/executor/utils/tool-binding-labels'
112-
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
113115
import { executeProviderRequest } from '@/providers'
114116
import {
115117
formatAttachmentSizes,
@@ -2710,9 +2712,7 @@ export class AgentBlockHandler implements BlockHandler {
27102712
? config.modelInputs
27112713
: omit(config.modelInputs, [
27122714
...PROVIDER_FAMILY_CREDENTIAL_FIELDS,
2713-
'vertexCredential',
2714-
'bedrockCredential',
2715-
'bedrockAuthMode',
2715+
...BLOCK_MODEL_CREDENTIAL_FIELDS,
27162716
])),
27172717
apiKey: resolveFallbackApiKey({
27182718
candidate,
@@ -2989,19 +2989,12 @@ export class AgentBlockHandler implements BlockHandler {
29892989
const model = providerRequest.model
29902990

29912991
try {
2992-
let finalApiKey: string | undefined = providerRequest.apiKey
2993-
2994-
if (providerId === 'vertex' && providerRequest.vertexCredential) {
2995-
finalApiKey = await resolveVertexCredential({
2996-
credentialId: providerRequest.vertexCredential,
2997-
actingUserId: ctx.userId,
2998-
workspaceId: ctx.workspaceId,
2999-
workflowId: ctx.workflowId,
3000-
executionId: ctx.executionId,
3001-
callerLabel: 'vertex-agent',
3002-
})
3003-
}
3004-
const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, providerRequest)
2992+
const credentials = await resolveBlockModelCredentials(
2993+
ctx,
2994+
providerId,
2995+
providerRequest,
2996+
'vertex-agent'
2997+
)
30052998

30062999
const { blockData, blockNameMapping } = collectBlockData(ctx)
30073000
const agentMemoryRetrieval = agentConversation?.memoryId
@@ -3024,15 +3017,15 @@ export class AgentBlockHandler implements BlockHandler {
30243017
: providerRequest.tools,
30253018
temperature: providerRequest.temperature,
30263019
maxTokens: providerRequest.maxTokens,
3027-
apiKey: finalApiKey,
3020+
apiKey: credentials.apiKey ?? providerRequest.apiKey,
30283021
azureEndpoint: providerRequest.azureEndpoint,
30293022
azureApiVersion: providerRequest.azureApiVersion,
30303023
vertexProject: providerRequest.vertexProject,
30313024
vertexLocation: providerRequest.vertexLocation,
30323025
bedrockAccessKeyId: providerRequest.bedrockAccessKeyId,
30333026
bedrockSecretKey: providerRequest.bedrockSecretKey,
30343027
bedrockRegion: providerRequest.bedrockRegion,
3035-
bedrockSession,
3028+
bedrockSession: credentials.bedrockSession,
30363029
responseFormat: providerRequest.responseFormat,
30373030
workflowId: providerRequest.workflowId,
30383031
workspaceId: ctx.workspaceId,

‎apps/sim/executor/handlers/evaluator/handler.test.ts‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
import {
66
llmBedrockCredentialMock,
77
llmBedrockCredentialMockFns,
8+
MockBedrockCredentialError,
89
} from '@sim/testing/mocks/llm-bedrock-credential.mock'
910
import { permissionsMock, permissionsMockFns } from '@sim/testing/mocks/permissions.mock'
1011
import '@sim/testing/mocks/executor'
@@ -252,6 +253,50 @@ describe('EvaluatorBlockHandler', () => {
252253
expect(JSON.stringify(fallbackRequest)).not.toContain('saved-secret')
253254
})
254255

256+
it('scores with a fallback model when the saved AWS credential is unavailable on the final try', async () => {
257+
mockGetProviderFromModel.mockImplementation((model: string) =>
258+
model.startsWith('claude') ? 'anthropic' : 'bedrock'
259+
)
260+
mockResolveBedrockCredentialSession.mockRejectedValueOnce(
261+
new MockBedrockCredentialError('AWS STS is unavailable', { reason: 'unavailable' })
262+
)
263+
mockExecuteProviderRequest.mockResolvedValueOnce({
264+
content: '{"score1":7}',
265+
model: 'claude-sonnet-5',
266+
tokens: { input: 12, output: 3, total: 15 },
267+
cost: 0,
268+
})
269+
mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry()
270+
mockContext.blockLogs = [
271+
{
272+
blockId: mockBlock.id,
273+
startedAt: '',
274+
endedAt: '',
275+
durationMs: 0,
276+
success: false,
277+
executionOrder: 1,
278+
},
279+
]
280+
281+
const output = await handler.execute(
282+
mockContext,
283+
mockBlock,
284+
{
285+
...admissionInputs,
286+
model: 'bedrock/claude',
287+
bedrockAuthMode: 'credential',
288+
bedrockCredential: 'aws-cred',
289+
fallbackModels: [{ model: 'claude-sonnet-5' }],
290+
},
291+
{ nodeId: mockBlock.id, retry: { attempt: 2, maxTries: 2, isFinalTry: true } }
292+
)
293+
294+
expect(output).toMatchObject({ model: 'claude-sonnet-5', score1: 7 })
295+
expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1)
296+
expect(mockExecuteProviderRequest.mock.calls[0][0]).toBe('anthropic')
297+
expect(mockContext.blockLogs[0].modelFallbacks).toEqual(['bedrock/claude'])
298+
})
299+
255300
it('fails a Bedrock evaluation whose AWS credential cannot be used instead of sending it', async () => {
256301
mockGetProviderFromModel.mockReturnValue('bedrock')
257302
mockResolveBedrockCredentialSession.mockRejectedValueOnce(

‎apps/sim/executor/handlers/evaluator/handler.ts‎

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import type { BlockOutput } from '@/blocks/types'
1010
import { validateModelProvider } from '@/ee/access-control/utils/permission-check'
1111
import { BlockType, DEFAULTS, EVALUATOR } from '@/executor/constants'
1212
import type { BlockHandler, BlockNodeMetadata, ExecutionContext } from '@/executor/types'
13-
import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential'
1413
import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json'
1514
import { executeModelRequestWithFallbacks } from '@/executor/utils/model-fallback-request'
1615
import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection'
@@ -19,7 +18,6 @@ import type {
1918
ResolvedSecretInputPath,
2019
ResolvedSecretTraceRegistry,
2120
} from '@/executor/utils/resolved-secret-trace-registry'
22-
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
2321
import { resolveProxiedModelCost } from '@/providers/cost-policy'
2422
import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models'
2523
import type { ProviderRequest } from '@/providers/types'
@@ -181,19 +179,6 @@ export class EvaluatorBlockHandler implements BlockHandler {
181179
await validateModelProvider(ctx.userId, ctx.workspaceId, model, ctx)
182180
const providerId = getProviderFromModel(model)
183181

184-
let finalApiKey: string | undefined = evaluatorConfig.apiKey
185-
if (providerId === 'vertex' && evaluatorConfig.vertexCredential) {
186-
finalApiKey = await resolveVertexCredential({
187-
credentialId: evaluatorConfig.vertexCredential,
188-
actingUserId: ctx.userId,
189-
workspaceId: ctx.workspaceId,
190-
workflowId: ctx.workflowId,
191-
executionId: ctx.executionId,
192-
callerLabel: 'vertex-evaluator',
193-
})
194-
}
195-
const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, evaluatorConfig)
196-
197182
try {
198183
const providerRequest: ProviderRequest = {
199184
model,
@@ -208,15 +193,14 @@ export class EvaluatorBlockHandler implements BlockHandler {
208193
]),
209194

210195
temperature: EVALUATOR.DEFAULT_TEMPERATURE,
211-
apiKey: finalApiKey,
196+
apiKey: evaluatorConfig.apiKey,
212197
azureEndpoint: inputs.azureEndpoint,
213198
azureApiVersion: inputs.azureApiVersion,
214199
vertexProject: evaluatorConfig.vertexProject,
215200
vertexLocation: evaluatorConfig.vertexLocation,
216201
bedrockAccessKeyId: evaluatorConfig.bedrockAccessKeyId,
217202
bedrockSecretKey: evaluatorConfig.bedrockSecretKey,
218203
bedrockRegion: evaluatorConfig.bedrockRegion,
219-
bedrockSession,
220204
workflowId: ctx.workflowId,
221205
workspaceId: ctx.workspaceId,
222206
}
@@ -231,6 +215,7 @@ export class EvaluatorBlockHandler implements BlockHandler {
231215
providerId,
232216
request: providerRequest,
233217
resolvedSecretTraceRegistry: modelInputProjection.registry,
218+
credentials: { inputs: evaluatorConfig, callerLabel: 'vertex-evaluator' },
234219
})
235220

236221
const parsedContent = this.extractJSONFromResponse(

‎apps/sim/executor/handlers/router/handler.test.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
import {
66
llmBedrockCredentialMock,
77
llmBedrockCredentialMockFns,
8+
MockBedrockCredentialError,
89
} from '@sim/testing/mocks/llm-bedrock-credential.mock'
910
import { permissionsMock, permissionsMockFns } from '@sim/testing/mocks/permissions.mock'
1011
import '@sim/testing/mocks/executor'
@@ -223,6 +224,53 @@ describe('RouterBlockHandler', () => {
223224
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
224225
})
225226

227+
it('routes with a fallback model when the saved AWS credential is unavailable on the final try', async () => {
228+
mockGetProviderFromModel.mockImplementation((model: string) =>
229+
model.startsWith('claude') ? 'anthropic' : 'bedrock'
230+
)
231+
mockResolveBedrockCredentialSession.mockRejectedValueOnce(
232+
new MockBedrockCredentialError('AWS STS is unavailable', { reason: 'unavailable' })
233+
)
234+
mockExecuteProviderRequest.mockResolvedValueOnce({
235+
content: 'target-block-1',
236+
model: 'claude-sonnet-5',
237+
tokens: { input: 10, output: 2, total: 12 },
238+
cost: 0,
239+
})
240+
mockContext.resolvedSecretTraceRegistry = new ResolvedSecretTraceRegistry()
241+
mockContext.blockLogs = [
242+
{
243+
blockId: mockBlock.id,
244+
startedAt: '',
245+
endedAt: '',
246+
durationMs: 0,
247+
success: false,
248+
executionOrder: 1,
249+
},
250+
]
251+
252+
const output = await handler.execute(
253+
mockContext,
254+
mockBlock,
255+
{
256+
prompt: 'Pick a destination',
257+
model: 'bedrock/claude',
258+
bedrockAuthMode: 'credential',
259+
bedrockCredential: 'aws-cred',
260+
fallbackModels: [{ model: 'claude-sonnet-5' }],
261+
},
262+
{ nodeId: mockBlock.id, retry: { attempt: 2, maxTries: 2, isFinalTry: true } }
263+
)
264+
265+
expect(output).toMatchObject({
266+
model: 'claude-sonnet-5',
267+
selectedPath: { blockId: 'target-block-1' },
268+
})
269+
expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1)
270+
expect(mockExecuteProviderRequest.mock.calls[0][0]).toBe('anthropic')
271+
expect(mockContext.blockLogs[0].modelFallbacks).toEqual(['bedrock/claude'])
272+
})
273+
226274
it('selects the same legacy destination when a fallback provider answers', async () => {
227275
mockExecuteProviderRequest
228276
.mockRejectedValueOnce(new Error('overloaded'))

‎apps/sim/executor/handlers/router/handler.ts‎

Lines changed: 4 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,9 @@ import {
1717
ROUTER,
1818
} from '@/executor/constants'
1919
import type { BlockHandler, BlockNodeMetadata, ExecutionContext } from '@/executor/types'
20-
import { resolveBlockBedrockSession } from '@/executor/utils/bedrock-credential'
2120
import { executeModelRequestWithFallbacks } from '@/executor/utils/model-fallback-request'
2221
import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal'
2322
import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry'
24-
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
2523
import { resolveProxiedModelCost } from '@/providers/cost-policy'
2624
import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models'
2725
import type { ProviderRequest } from '@/providers/types'
@@ -116,33 +114,19 @@ export class RouterBlockHandler implements BlockHandler {
116114
await validateModelProvider(ctx.userId, ctx.workspaceId, resolved.model, ctx)
117115
const providerId = getProviderFromModel(resolved.model)
118116

119-
let finalApiKey: string | undefined = routerConfig.apiKey
120-
if (providerId === 'vertex' && routerConfig.vertexCredential) {
121-
finalApiKey = await resolveVertexCredential({
122-
credentialId: routerConfig.vertexCredential,
123-
actingUserId: ctx.userId,
124-
workspaceId: ctx.workspaceId,
125-
workflowId: ctx.workflowId,
126-
executionId: ctx.executionId,
127-
callerLabel: 'vertex-router',
128-
})
129-
}
130-
const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, routerConfig)
131-
132117
const providerRequest: ProviderRequest = {
133118
model: resolved.model,
134119
systemPrompt: resolved.systemPrompt,
135120
context: JSON.stringify(messages),
136121
temperature: ROUTER.INFERENCE_TEMPERATURE,
137-
apiKey: finalApiKey,
122+
apiKey: routerConfig.apiKey,
138123
azureEndpoint: inputs.azureEndpoint,
139124
azureApiVersion: inputs.azureApiVersion,
140125
vertexProject: routerConfig.vertexProject,
141126
vertexLocation: routerConfig.vertexLocation,
142127
bedrockAccessKeyId: routerConfig.bedrockAccessKeyId,
143128
bedrockSecretKey: routerConfig.bedrockSecretKey,
144129
bedrockRegion: routerConfig.bedrockRegion,
145-
bedrockSession,
146130
workflowId: ctx.workflowId,
147131
workspaceId: ctx.workspaceId,
148132
}
@@ -157,6 +141,7 @@ export class RouterBlockHandler implements BlockHandler {
157141
providerId,
158142
request: providerRequest,
159143
resolvedSecretTraceRegistry: modelInputProjection.registry,
144+
credentials: { inputs: routerConfig, callerLabel: 'vertex-router' },
160145
})
161146

162147
const chosenBlockId = result.content.trim().toLowerCase()
@@ -290,33 +275,19 @@ export class RouterBlockHandler implements BlockHandler {
290275
await validateModelProvider(ctx.userId, ctx.workspaceId, resolved.model, ctx)
291276
const providerId = getProviderFromModel(resolved.model)
292277

293-
let finalApiKey: string | undefined = routerConfig.apiKey
294-
if (providerId === 'vertex' && routerConfig.vertexCredential) {
295-
finalApiKey = await resolveVertexCredential({
296-
credentialId: routerConfig.vertexCredential,
297-
actingUserId: ctx.userId,
298-
workspaceId: ctx.workspaceId,
299-
workflowId: ctx.workflowId,
300-
executionId: ctx.executionId,
301-
callerLabel: 'vertex-router',
302-
})
303-
}
304-
const bedrockSession = await resolveBlockBedrockSession(ctx, providerId, routerConfig)
305-
306278
const providerRequest: ProviderRequest = {
307279
model: resolved.model,
308280
systemPrompt: resolved.systemPrompt,
309281
context: JSON.stringify(messages),
310282
temperature: ROUTER.INFERENCE_TEMPERATURE,
311-
apiKey: finalApiKey,
283+
apiKey: routerConfig.apiKey,
312284
azureEndpoint: inputs.azureEndpoint,
313285
azureApiVersion: inputs.azureApiVersion,
314286
vertexProject: routerConfig.vertexProject,
315287
vertexLocation: routerConfig.vertexLocation,
316288
bedrockAccessKeyId: routerConfig.bedrockAccessKeyId,
317289
bedrockSecretKey: routerConfig.bedrockSecretKey,
318290
bedrockRegion: routerConfig.bedrockRegion,
319-
bedrockSession,
320291
workflowId: ctx.workflowId,
321292
workspaceId: ctx.workspaceId,
322293
responseFormat: {
@@ -350,6 +321,7 @@ export class RouterBlockHandler implements BlockHandler {
350321
providerId,
351322
request: providerRequest,
352323
resolvedSecretTraceRegistry: modelInputProjection.registry,
324+
credentials: { inputs: routerConfig, callerLabel: 'vertex-router' },
353325
})
354326

355327
let chosenRouteId: string

0 commit comments

Comments
 (0)