Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -73,3 +73,7 @@ start-collector.sh
## Helm Chart Tests
helm/sim/test
i18n.cache

## Claude Code
.claude/launch.json
.claude/worktrees/
97 changes: 35 additions & 62 deletions apps/sim/executor/handlers/agent/agent-handler.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
import { db } from '@sim/db'
import { account, mcpServers } from '@sim/db/schema'
import { mcpServers } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { and, eq, inArray, isNull } from 'drizzle-orm'
import { createMcpToolId } from '@/lib/mcp/utils'
import { refreshTokenIfNeeded, resolveOAuthAccountId } from '@/app/api/auth/oauth/utils'
import { getAllBlocks } from '@/blocks'
import type { BlockOutput } from '@/blocks/types'
import {
Expand All @@ -30,6 +29,7 @@ import type { BlockHandler, ExecutionContext, StreamingExecution } from '@/execu
import { collectBlockData } from '@/executor/utils/block-data'
import { buildAPIUrl, buildAuthHeaders } from '@/executor/utils/http'
import { stringifyJSON } from '@/executor/utils/json'
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
import { executeProviderRequest } from '@/providers'
import { getProviderFromModel, transformBlockTool } from '@/providers/utils'
import type { SerializedBlock } from '@/serializer/types'
Expand Down Expand Up @@ -439,24 +439,15 @@ export class AgentBlockHandler implements BlockHandler {
tool: ToolInput
): Promise<any> {
const { serverId, toolName, serverName, ...userProvidedParams } = tool.params || {}

const { filterSchemaForLLM } = await import('@/tools/params')
const filteredSchema = filterSchemaForLLM(
tool.schema || { type: 'object', properties: {} },
userProvidedParams
)

const toolId = createMcpToolId(serverId, toolName)

return {
id: toolId,
name: toolName,
return this.buildMcpTool({
serverId,
toolName,
description:
tool.schema?.description || `MCP tool ${toolName} from ${serverName || serverId}`,
parameters: filteredSchema,
params: userProvidedParams,
usageControl: tool.usageControl || 'auto',
}
schema: tool.schema || { type: 'object', properties: {} },
userProvidedParams,
usageControl: tool.usageControl,
})
}

/**
Expand Down Expand Up @@ -585,22 +576,35 @@ export class AgentBlockHandler implements BlockHandler {
serverId: string
): Promise<any> {
const { toolName, ...userProvidedParams } = tool.params || {}
return this.buildMcpTool({
serverId,
toolName,
description: mcpTool.description || `MCP tool ${toolName} from ${mcpTool.serverName}`,
schema: mcpTool.inputSchema || { type: 'object', properties: {} },
userProvidedParams,
usageControl: tool.usageControl,
})
}

private async buildMcpTool(config: {
serverId: string
toolName: string
description: string
schema: any
userProvidedParams: Record<string, any>
usageControl?: string
}): Promise<any> {
const { filterSchemaForLLM } = await import('@/tools/params')
const filteredSchema = filterSchemaForLLM(
mcpTool.inputSchema || { type: 'object', properties: {} },
userProvidedParams
)

const toolId = createMcpToolId(serverId, toolName)
const filteredSchema = filterSchemaForLLM(config.schema, config.userProvidedParams)
const toolId = createMcpToolId(config.serverId, config.toolName)

return {
id: toolId,
name: toolName,
description: mcpTool.description || `MCP tool ${toolName} from ${mcpTool.serverName}`,
name: config.toolName,
description: config.description,
parameters: filteredSchema,
params: userProvidedParams,
usageControl: tool.usageControl || 'auto',
params: config.userProvidedParams,
usageControl: config.usageControl || 'auto',
}
}

Expand Down Expand Up @@ -924,9 +928,9 @@ export class AgentBlockHandler implements BlockHandler {
let finalApiKey: string | undefined = providerRequest.apiKey

if (providerId === 'vertex' && providerRequest.vertexCredential) {
finalApiKey = await this.resolveVertexCredential(
finalApiKey = await resolveVertexCredential(
providerRequest.vertexCredential,
ctx.workflowId
'vertex-agent'
)
}

Expand Down Expand Up @@ -973,37 +977,6 @@ export class AgentBlockHandler implements BlockHandler {
}
}

/**
* Resolves a Vertex AI OAuth credential to an access token
*/
private async resolveVertexCredential(credentialId: string, workflowId: string): Promise<string> {
const requestId = `vertex-${Date.now()}`

logger.info(`[${requestId}] Resolving Vertex AI credential: ${credentialId}`)

const resolved = await resolveOAuthAccountId(credentialId)
if (!resolved) {
throw new Error(`Vertex AI credential is not a valid OAuth credential: ${credentialId}`)
}

const credential = await db.query.account.findFirst({
where: eq(account.id, resolved.accountId),
})

if (!credential) {
throw new Error(`Vertex AI credential not found: ${credentialId}`)
}

const { accessToken } = await refreshTokenIfNeeded(requestId, credential, resolved.accountId)

if (!accessToken) {
throw new Error('Failed to get Vertex AI access token')
}

logger.info(`[${requestId}] Successfully resolved Vertex AI credential`)
return accessToken
}

private handleExecutionError(
error: any,
startTime: number,
Expand Down Expand Up @@ -1187,7 +1160,7 @@ export class AgentBlockHandler implements BlockHandler {
},
toolCalls: {
list: result.toolCalls?.map(this.formatToolCall.bind(this)) || [],
count: result.toolCalls?.length || DEFAULTS.EXECUTION_TIME,
count: result.toolCalls?.length ?? 0,
},
providerTiming: result.timing,
cost: result.cost,
Expand Down
45 changes: 7 additions & 38 deletions apps/sim/executor/handlers/evaluator/evaluator-handler.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,11 @@
import { db } from '@sim/db'
import { account } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { eq } from 'drizzle-orm'
import { refreshTokenIfNeeded, resolveOAuthAccountId } from '@/app/api/auth/oauth/utils'
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, ExecutionContext } from '@/executor/types'
import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http'
import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json'
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
import { calculateCost, getProviderFromModel } from '@/providers/utils'
import type { SerializedBlock } from '@/serializer/types'

Expand Down Expand Up @@ -44,7 +41,10 @@ export class EvaluatorBlockHandler implements BlockHandler {

let finalApiKey: string | undefined = evaluatorConfig.apiKey
if (providerId === 'vertex' && evaluatorConfig.vertexCredential) {
finalApiKey = await this.resolveVertexCredential(evaluatorConfig.vertexCredential)
finalApiKey = await resolveVertexCredential(
evaluatorConfig.vertexCredential,
'vertex-evaluator'
)
}

const processedContent = this.processContent(inputs.content)
Expand Down Expand Up @@ -234,7 +234,7 @@ export class EvaluatorBlockHandler implements BlockHandler {
if (Object.keys(parsedContent).length === 0) {
validMetrics.forEach((metric: any) => {
if (metric?.name) {
metricScores[metric.name.toLowerCase()] = DEFAULTS.EXECUTION_TIME
metricScores[metric.name.toLowerCase()] = 0
}
})
return metricScores
Expand Down Expand Up @@ -273,37 +273,6 @@ export class EvaluatorBlockHandler implements BlockHandler {
}

logger.warn(`Metric "${metricName}" not found in LLM response`)
return DEFAULTS.EXECUTION_TIME
}

/**
* Resolves a Vertex AI OAuth credential to an access token
*/
private async resolveVertexCredential(credentialId: string): Promise<string> {
const requestId = `vertex-evaluator-${Date.now()}`

logger.info(`[${requestId}] Resolving Vertex AI credential: ${credentialId}`)

const resolved = await resolveOAuthAccountId(credentialId)
if (!resolved) {
throw new Error(`Vertex AI credential is not a valid OAuth credential: ${credentialId}`)
}

const credential = await db.query.account.findFirst({
where: eq(account.id, resolved.accountId),
})

if (!credential) {
throw new Error(`Vertex AI credential not found: ${credentialId}`)
}

const { accessToken } = await refreshTokenIfNeeded(requestId, credential, resolved.accountId)

if (!accessToken) {
throw new Error('Failed to get Vertex AI access token')
}

logger.info(`[${requestId}] Successfully resolved Vertex AI credential`)
return accessToken
return 0
}
}
Loading