From b842827310f1f792a6b69f69b869fdc59a52bbbd Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Tue, 31 Mar 2026 17:22:05 -0700 Subject: [PATCH 1/4] Directly query db for custom tool id --- .../handlers/agent/agent-handler.test.ts | 69 +++++++------------ .../executor/handlers/agent/agent-handler.ts | 35 ++-------- apps/sim/tools/utils.ts | 57 ++++----------- 3 files changed, 41 insertions(+), 120 deletions(-) diff --git a/apps/sim/executor/handlers/agent/agent-handler.test.ts b/apps/sim/executor/handlers/agent/agent-handler.test.ts index 87fba03a417..d64983129a4 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.test.ts +++ b/apps/sim/executor/handlers/agent/agent-handler.test.ts @@ -110,6 +110,12 @@ vi.mock('@sim/db/schema', () => ({ }, })) +const mockGetCustomToolById = vi.fn() + +vi.mock('@/lib/workflows/custom-tools/operations', () => ({ + getCustomToolById: (...args: unknown[]) => mockGetCustomToolById(...args), +})) + setupGlobalFetchMock() const mockGetAllBlocks = getAllBlocks as Mock @@ -1957,49 +1963,22 @@ describe('AgentBlockHandler', () => { const staleInlineCode = 'return { title, content };' const dbCode = 'return { title, content, format };' - function mockFetchForCustomTool(toolId: string) { - mockFetch.mockImplementation((url: string) => { - if (typeof url === 'string' && url.includes('/api/tools/custom')) { + function mockDBForCustomTool(toolId: string) { + mockGetCustomToolById.mockImplementation(({ toolId: id }: { toolId: string }) => { + if (id === toolId) { return Promise.resolve({ - ok: true, - headers: { get: () => null }, - json: () => - Promise.resolve({ - data: [ - { - id: toolId, - title: 'formatReport', - schema: dbSchema, - code: dbCode, - }, - ], - }), + id: toolId, + title: 'formatReport', + schema: dbSchema, + code: dbCode, }) } - return Promise.resolve({ - ok: true, - headers: { get: () => null }, - json: () => Promise.resolve({}), - }) + return Promise.resolve(null) }) } - function mockFetchFailure() { - mockFetch.mockImplementation((url: string) => { - if (typeof url === 'string' && url.includes('/api/tools/custom')) { - return Promise.resolve({ - ok: false, - status: 500, - headers: { get: () => null }, - json: () => Promise.resolve({}), - }) - } - return Promise.resolve({ - ok: true, - headers: { get: () => null }, - json: () => Promise.resolve({}), - }) - }) + function mockDBFailure() { + mockGetCustomToolById.mockRejectedValue(new Error('DB connection failed')) } beforeEach(() => { @@ -2008,11 +1987,12 @@ describe('AgentBlockHandler', () => { writable: true, configurable: true, }) + mockGetCustomToolById.mockReset() }) it('should always fetch latest schema from DB when customToolId is present', async () => { const toolId = 'custom-tool-123' - mockFetchForCustomTool(toolId) + mockDBForCustomTool(toolId) const inputs = { model: 'gpt-4o', @@ -2046,7 +2026,7 @@ describe('AgentBlockHandler', () => { it('should fetch from DB when customToolId has no inline schema', async () => { const toolId = 'custom-tool-123' - mockFetchForCustomTool(toolId) + mockDBForCustomTool(toolId) const inputs = { model: 'gpt-4o', @@ -2075,7 +2055,7 @@ describe('AgentBlockHandler', () => { }) it('should fall back to inline schema when DB fetch fails and inline exists', async () => { - mockFetchFailure() + mockDBFailure() const inputs = { model: 'gpt-4o', @@ -2107,7 +2087,7 @@ describe('AgentBlockHandler', () => { }) it('should return null when DB fetch fails and no inline schema exists', async () => { - mockFetchFailure() + mockDBFailure() const inputs = { model: 'gpt-4o', @@ -2135,7 +2115,7 @@ describe('AgentBlockHandler', () => { it('should use DB schema when customToolId resolves', async () => { const toolId = 'custom-tool-123' - mockFetchForCustomTool(toolId) + mockDBForCustomTool(toolId) const inputs = { model: 'gpt-4o', @@ -2185,10 +2165,7 @@ describe('AgentBlockHandler', () => { await handler.execute(mockContext, mockBlock, inputs) - const customToolFetches = mockFetch.mock.calls.filter( - (call: any[]) => typeof call[0] === 'string' && call[0].includes('/api/tools/custom') - ) - expect(customToolFetches.length).toBe(0) + expect(mockGetCustomToolById).not.toHaveBeenCalled() expect(mockExecuteProviderRequest).toHaveBeenCalled() const providerCall = mockExecuteProviderRequest.mock.calls[0] diff --git a/apps/sim/executor/handlers/agent/agent-handler.ts b/apps/sim/executor/handlers/agent/agent-handler.ts index 492d07c9b00..1e30b60ce9b 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.ts +++ b/apps/sim/executor/handlers/agent/agent-handler.ts @@ -3,6 +3,7 @@ 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 { getCustomToolById } from '@/lib/workflows/custom-tools/operations' import { getAllBlocks } from '@/blocks' import type { BlockOutput } from '@/blocks/types' import { @@ -289,38 +290,12 @@ export class AgentBlockHandler implements BlockHandler { } try { - const headers = await buildAuthHeaders(ctx.userId) - const params: Record = {} - - if (ctx.workspaceId) { - params.workspaceId = ctx.workspaceId - } - if (ctx.workflowId) { - params.workflowId = ctx.workflowId - } - if (ctx.userId) { - params.userId = ctx.userId - } - - const url = buildAPIUrl('/api/tools/custom', params) - const response = await fetch(url.toString(), { - method: 'GET', - headers, + const tool = await getCustomToolById({ + toolId: customToolId, + userId: ctx.userId, + workspaceId: ctx.workspaceId, }) - if (!response.ok) { - await response.text().catch(() => {}) - logger.error(`Failed to fetch custom tools: ${response.status}`) - return null - } - - const data = await response.json() - if (!data.data || !Array.isArray(data.data)) { - logger.error('Invalid custom tools API response') - return null - } - - const tool = data.data.find((t: any) => t.id === customToolId) if (!tool) { logger.warn(`Custom tool not found by ID: ${customToolId}`) return null diff --git a/apps/sim/tools/utils.ts b/apps/sim/tools/utils.ts index 581d4a6ac58..762ad2f2124 100644 --- a/apps/sim/tools/utils.ts +++ b/apps/sim/tools/utils.ts @@ -1,6 +1,6 @@ import { createLogger } from '@sim/logger' import { getMaxExecutionTimeout } from '@/lib/core/execution-limits' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' +import { getCustomToolById, listCustomTools } from '@/lib/workflows/custom-tools/operations' import { AGENT, isCustomTool } from '@/executor/constants' import { getCustomTool } from '@/hooks/queries/custom-tools' import { useEnvironmentStore } from '@/stores/settings/environment' @@ -320,7 +320,7 @@ export async function getToolAsync( // Check if it's a custom tool if (isCustomTool(toolId)) { - return fetchCustomToolFromAPI(toolId, workflowId, userId) + return fetchCustomToolFromDB(toolId, workflowId, userId) } return undefined @@ -364,8 +364,8 @@ function createToolConfig(customTool: any, customToolId: string): ToolConfig { } } -// Create a tool config from a custom tool definition by fetching from API -async function fetchCustomToolFromAPI( +// Create a tool config from a custom tool definition by querying the database directly +async function fetchCustomToolFromDB( customToolId: string, workflowId?: string, userId?: string @@ -373,51 +373,20 @@ async function fetchCustomToolFromAPI( const identifier = customToolId.replace('custom_', '') try { - const baseUrl = getInternalApiBaseUrl() - const url = new URL('/api/tools/custom', baseUrl) - - if (workflowId) { - url.searchParams.append('workflowId', workflowId) - } - if (userId) { - url.searchParams.append('userId', userId) - } - - // For server-side calls (during workflow execution), use internal JWT token - const headers: Record = {} - if (typeof window === 'undefined') { - try { - const { generateInternalToken } = await import('@/lib/auth/internal') - const internalToken = await generateInternalToken(userId) - headers.Authorization = `Bearer ${internalToken}` - } catch (error) { - logger.warn('Failed to generate internal token for custom tools fetch', { error }) - // Continue without token - will fail auth and be reported upstream - } - } - - const response = await fetch(url.toString(), { - headers, - }) - - if (!response.ok) { - await response.text().catch(() => {}) - logger.error(`Failed to fetch custom tools: ${response.statusText}`) + if (!userId) { + logger.error(`Cannot fetch custom tool without userId: ${identifier}`) return undefined } - const result = await response.json() + // Try to find by ID first + let customTool = await getCustomToolById({ toolId: identifier, userId }) - if (!result.data || !Array.isArray(result.data)) { - logger.error(`Invalid response when fetching custom tools: ${JSON.stringify(result)}`) - return undefined + // Fall back to searching by title + if (!customTool) { + const allTools = await listCustomTools({ userId }) + customTool = allTools.find((t) => t.title === identifier) ?? null } - // Try to find the tool by ID or title - const customTool = result.data.find( - (tool: any) => tool.id === identifier || tool.title === identifier - ) - if (!customTool) { logger.error(`Custom tool not found: ${identifier}`) return undefined @@ -458,7 +427,7 @@ async function fetchCustomToolFromAPI( }, } } catch (error) { - logger.error(`Error fetching custom tool ${identifier} from API:`, error) + logger.error(`Error fetching custom tool ${identifier} from DB:`, error) return undefined } } From 3c92b344206ecd009edd5aa424c407d3206af6f3 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Wed, 1 Apr 2026 16:15:25 -0700 Subject: [PATCH 2/4] Switch back to inline imports --- .../executor/handlers/agent/agent-handler.ts | 2 +- apps/sim/tools/utils.server.ts | 57 ++++--------- apps/sim/tools/utils.ts | 84 ------------------- 3 files changed, 16 insertions(+), 127 deletions(-) diff --git a/apps/sim/executor/handlers/agent/agent-handler.ts b/apps/sim/executor/handlers/agent/agent-handler.ts index 9926e4d5b14..204b2cdbf9a 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.ts +++ b/apps/sim/executor/handlers/agent/agent-handler.ts @@ -3,7 +3,6 @@ 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 { getCustomToolById } from '@/lib/workflows/custom-tools/operations' import { getAllBlocks } from '@/blocks' import type { BlockOutput } from '@/blocks/types' import { @@ -284,6 +283,7 @@ export class AgentBlockHandler implements BlockHandler { } try { + const { getCustomToolById } = await import('@/lib/workflows/custom-tools/operations') const tool = await getCustomToolById({ toolId: customToolId, userId: ctx.userId, diff --git a/apps/sim/tools/utils.server.ts b/apps/sim/tools/utils.server.ts index 7abce99e125..7123590850e 100644 --- a/apps/sim/tools/utils.server.ts +++ b/apps/sim/tools/utils.server.ts @@ -1,10 +1,8 @@ import { createLogger } from '@sim/logger' -import { generateInternalToken } from '@/lib/auth/internal' import { secureFetchWithPinnedIP, validateUrlWithDNS, } from '@/lib/core/security/input-validation.server' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' import { isCustomTool } from '@/executor/constants' import type { CustomToolDefinition } from '@/hooks/queries/custom-tools' import { extractErrorMessage } from '@/tools/error-extractors' @@ -97,13 +95,13 @@ export async function getToolAsync( if (builtInTool) return builtInTool if (isCustomTool(toolId)) { - return fetchCustomToolFromAPI(toolId, context) + return fetchCustomToolFromDB(toolId, context) } return undefined } -async function fetchCustomToolFromAPI( +async function fetchCustomToolFromDB( customToolId: string, context: GetToolAsyncContext ): Promise { @@ -111,53 +109,28 @@ async function fetchCustomToolFromAPI( const identifier = customToolId.replace('custom_', '') try { - const baseUrl = getInternalApiBaseUrl() - const url = new URL('/api/tools/custom', baseUrl) - - if (workflowId) { - url.searchParams.append('workflowId', workflowId) - } - if (userId) { - url.searchParams.append('userId', userId) - } - if (workspaceId) { - url.searchParams.append('workspaceId', workspaceId) - } - - const headers: Record = {} - - try { - const internalToken = await generateInternalToken(userId) - headers.Authorization = `Bearer ${internalToken}` - } catch (error) { - logger.warn('Failed to generate internal token for custom tools fetch', { error }) - } - - const response = await fetch(url.toString(), { headers }) - - if (!response.ok) { - await response.text().catch(() => {}) - logger.error(`Failed to fetch custom tools: ${response.statusText}`) + if (!userId) { + logger.error(`Cannot fetch custom tool without userId: ${identifier}`) return undefined } - const result = await response.json() - - if (!result.data || !Array.isArray(result.data)) { - logger.error(`Invalid response when fetching custom tools: ${JSON.stringify(result)}`) - return undefined - } + const { getCustomToolById, listCustomTools } = await import( + '@/lib/workflows/custom-tools/operations' + ) - const customTool = result.data.find( - (tool: CustomToolDefinition) => tool.id === identifier || tool.title === identifier - ) as CustomToolDefinition | undefined + // Try to find by ID first, fall back to searching by title + const customTool = + (await getCustomToolById({ toolId: identifier, userId, workspaceId })) ?? + (await listCustomTools({ userId, workspaceId })).find( + (t: { title: string }) => t.title === identifier + ) if (!customTool) { logger.error(`Custom tool not found: ${identifier}`) return undefined } - const toolConfig = createToolConfig(customTool, customToolId) + const toolConfig = createToolConfig(customTool as unknown as CustomToolDefinition, customToolId) return { ...toolConfig, @@ -168,7 +141,7 @@ async function fetchCustomToolFromAPI( }, } } catch (error) { - logger.error(`Error fetching custom tool ${identifier} from API:`, error) + logger.error(`Error fetching custom tool ${identifier} from DB:`, error) return undefined } } diff --git a/apps/sim/tools/utils.ts b/apps/sim/tools/utils.ts index 0b827abb1cb..d8a7044e62c 100644 --- a/apps/sim/tools/utils.ts +++ b/apps/sim/tools/utils.ts @@ -1,6 +1,5 @@ import { createLogger } from '@sim/logger' import { getMaxExecutionTimeout } from '@/lib/core/execution-limits' -import { getCustomToolById, listCustomTools } from '@/lib/workflows/custom-tools/operations' import { AGENT, isCustomTool } from '@/executor/constants' import type { CustomToolDefinition } from '@/hooks/queries/custom-tools' import { useEnvironmentStore } from '@/stores/settings/environment' @@ -295,24 +294,6 @@ export function getTool(toolId: string, _workspaceId?: string): ToolConfig | und return undefined } -// Get a tool by its ID asynchronously (supports server-side) -export async function getToolAsync( - toolId: string, - workflowId?: string, - userId?: string -): Promise { - // Check for built-in tools - const builtInTool = tools[toolId] - if (builtInTool) return builtInTool - - // Check if it's a custom tool - if (isCustomTool(toolId)) { - return fetchCustomToolFromDB(toolId, workflowId, userId) - } - - return undefined -} - // Helper function to create a tool config from a custom tool export function createToolConfig( customTool: CustomToolDefinition, @@ -354,69 +335,4 @@ export function createToolConfig( } } -// Create a tool config from a custom tool definition by querying the database directly -async function fetchCustomToolFromDB( - customToolId: string, - workflowId?: string, - userId?: string -): Promise { - const identifier = customToolId.replace('custom_', '') - - try { - if (!userId) { - logger.error(`Cannot fetch custom tool without userId: ${identifier}`) - return undefined - } - - // Try to find by ID first, fall back to searching by title - const customTool = - (await getCustomToolById({ toolId: identifier, userId })) ?? - (await listCustomTools({ userId })).find((t) => t.title === identifier) - - if (!customTool) { - logger.error(`Custom tool not found: ${identifier}`) - return undefined - } - - const schema = customTool.schema as Record - - // Create a parameter schema - const params = createParamSchema(customTool) - - // Create a tool config for the custom tool - return { - id: customToolId, - name: customTool.title, - description: schema.function?.description || '', - version: '1.0.0', - params, - - // Request configuration - for custom tools we'll use the execute endpoint - request: { - url: '/api/function/execute', - method: 'POST', - headers: () => ({ 'Content-Type': 'application/json' }), - body: createCustomToolRequestBody(customTool, false, workflowId), - }, - - // Same response handling as client-side - transformResponse: async (response: Response) => { - const data = await response.json() - - if (!data.success) { - throw new Error(data.error || 'Custom tool execution failed') - } - - return { - success: true, - output: data.output.result || data.output, - error: undefined, - } - }, - } - } catch (error) { - logger.error(`Error fetching custom tool ${identifier} from DB:`, error) - return undefined - } -} From 5b858b3f1598c839796cc2c4e7cf3f0c1b418246 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Wed, 1 Apr 2026 16:15:44 -0700 Subject: [PATCH 3/4] Fix lint --- apps/sim/tools/utils.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/apps/sim/tools/utils.ts b/apps/sim/tools/utils.ts index d8a7044e62c..2f944c18bd4 100644 --- a/apps/sim/tools/utils.ts +++ b/apps/sim/tools/utils.ts @@ -1,6 +1,5 @@ import { createLogger } from '@sim/logger' import { getMaxExecutionTimeout } from '@/lib/core/execution-limits' -import { AGENT, isCustomTool } from '@/executor/constants' import type { CustomToolDefinition } from '@/hooks/queries/custom-tools' import { useEnvironmentStore } from '@/stores/settings/environment' import { tools } from '@/tools/registry' @@ -334,5 +333,3 @@ export function createToolConfig( }, } } - - From ed468d78282e753cb3df9fb022f51c9d9284fd14 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Thu, 2 Apr 2026 00:16:33 -0700 Subject: [PATCH 4/4] Fix test --- apps/sim/tools/index.test.ts | 74 +++++++++++++++++++----------------- 1 file changed, 40 insertions(+), 34 deletions(-) diff --git a/apps/sim/tools/index.test.ts b/apps/sim/tools/index.test.ts index a55aefa1ee2..4b3ec00d8cb 100644 --- a/apps/sim/tools/index.test.ts +++ b/apps/sim/tools/index.test.ts @@ -16,19 +16,27 @@ import { import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' // Hoisted mock state - these are available to vi.mock factories -const { mockIsHosted, mockEnv, mockGetBYOKKey, mockGetToolAsync, mockRateLimiterFns } = vi.hoisted( - () => ({ - mockIsHosted: { value: false }, - mockEnv: { NEXT_PUBLIC_APP_URL: 'http://localhost:3000' } as Record, - mockGetBYOKKey: vi.fn(), - mockGetToolAsync: vi.fn(), - mockRateLimiterFns: { - acquireKey: vi.fn(), - preConsumeCapacity: vi.fn(), - consumeCapacity: vi.fn(), - }, - }) -) +const { + mockIsHosted, + mockEnv, + mockGetBYOKKey, + mockGetToolAsync, + mockRateLimiterFns, + mockGetCustomToolById, + mockListCustomTools, +} = vi.hoisted(() => ({ + mockIsHosted: { value: false }, + mockEnv: { NEXT_PUBLIC_APP_URL: 'http://localhost:3000' } as Record, + mockGetBYOKKey: vi.fn(), + mockGetToolAsync: vi.fn(), + mockRateLimiterFns: { + acquireKey: vi.fn(), + preConsumeCapacity: vi.fn(), + consumeCapacity: vi.fn(), + }, + mockGetCustomToolById: vi.fn(), + mockListCustomTools: vi.fn(), +})) // Mock feature flags vi.mock('@/lib/core/config/feature-flags', () => ({ @@ -214,6 +222,11 @@ vi.mock('@/hooks/queries/utils/custom-tool-cache', () => { } }) +vi.mock('@/lib/workflows/custom-tools/operations', () => ({ + getCustomToolById: mockGetCustomToolById, + listCustomTools: mockListCustomTools, +})) + vi.mock('@/tools/utils.server', async (importOriginal) => { const actual = await importOriginal() mockGetToolAsync.mockImplementation(actual.getToolAsync) @@ -307,30 +320,23 @@ describe('Custom Tools', () => { }) it('resolves custom tools through the async helper', async () => { - setupFetchMock({ - json: { - data: [ - { - id: 'remote-tool-123', - title: 'Custom Weather Tool', - schema: { - function: { - name: 'weather_tool', - description: 'Get weather information', - parameters: { - type: 'object', - properties: { - location: { type: 'string', description: 'City name' }, - }, - required: ['location'], - }, - }, + mockGetCustomToolById.mockResolvedValue({ + id: 'remote-tool-123', + title: 'Custom Weather Tool', + schema: { + function: { + name: 'weather_tool', + description: 'Get weather information', + parameters: { + type: 'object', + properties: { + location: { type: 'string', description: 'City name' }, }, + required: ['location'], }, - ], + }, }, - status: 200, - headers: { 'content-type': 'application/json' }, + code: '', }) const customTool = await getToolAsync('custom_remote-tool-123', {