Directly query db for custom tool id

Author: Theodore Li

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]

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<string, string> = {}
-
-      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

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,60 +364,29 @@ 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
 ): Promise<ToolConfig | undefined> {
   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<string, string> = {}
-    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
   }
 }