fix(agent): always fetch latest custom tool from DB when customToolId is present (#3208)

* fix(agent): always fetch latest custom tool from DB when customToolId is present

* test(agent): use generic test data for customToolId resolution tests

* fix(agent): mock buildAuthHeaders in tests for CI compatibility

* remove inline mocks in favor of sim/testing ones

apps/sim/app/api/auth/oauth/utils.test.ts

@@ -4,20 +4,10 @@
  * @vitest-environment node
  */
 
-import { loggerMock } from '@sim/testing'
+import { databaseMock, loggerMock } from '@sim/testing'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockReturnValue([]),
-    update: vi.fn().mockReturnThis(),
-    set: vi.fn().mockReturnThis(),
-    orderBy: vi.fn().mockReturnThis(),
-  },
-}))
+vi.mock('@sim/db', () => databaseMock)
 
 vi.mock('@/lib/oauth/oauth', () => ({
   refreshOAuthToken: vi.fn(),
@@ -34,13 +24,36 @@ import {
   refreshTokenIfNeeded,
 } from '@/app/api/auth/oauth/utils'
 
-const mockDbTyped = db as any
+const mockDb = db as any
 const mockRefreshOAuthToken = refreshOAuthToken as any
 
+/**
+ * Creates a chainable mock for db.select() calls.
+ * Returns a nested chain: select() -> from() -> where() -> limit() / orderBy()
+ */
+function mockSelectChain(limitResult: unknown[]) {
+  const mockLimit = vi.fn().mockReturnValue(limitResult)
+  const mockOrderBy = vi.fn().mockReturnValue(limitResult)
+  const mockWhere = vi.fn().mockReturnValue({ limit: mockLimit, orderBy: mockOrderBy })
+  const mockFrom = vi.fn().mockReturnValue({ where: mockWhere })
+  mockDb.select.mockReturnValueOnce({ from: mockFrom })
+  return { mockFrom, mockWhere, mockLimit }
+}
+
+/**
+ * Creates a chainable mock for db.update() calls.
+ * Returns a nested chain: update() -> set() -> where()
+ */
+function mockUpdateChain() {
+  const mockWhere = vi.fn().mockResolvedValue({})
+  const mockSet = vi.fn().mockReturnValue({ where: mockWhere })
+  mockDb.update.mockReturnValueOnce({ set: mockSet })
+  return { mockSet, mockWhere }
+}
+
 describe('OAuth Utils', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockDbTyped.limit.mockReturnValue([])
   })
 
   afterEach(() => {
@@ -50,20 +63,20 @@ describe('OAuth Utils', () => {
   describe('getCredential', () => {
     it('should return credential when found', async () => {
       const mockCredential = { id: 'credential-id', userId: 'test-user-id' }
-      mockDbTyped.limit.mockReturnValueOnce([mockCredential])
+      const { mockFrom, mockWhere, mockLimit } = mockSelectChain([mockCredential])
 
       const credential = await getCredential('request-id', 'credential-id', 'test-user-id')
 
-      expect(mockDbTyped.select).toHaveBeenCalled()
-      expect(mockDbTyped.from).toHaveBeenCalled()
-      expect(mockDbTyped.where).toHaveBeenCalled()
-      expect(mockDbTyped.limit).toHaveBeenCalledWith(1)
+      expect(mockDb.select).toHaveBeenCalled()
+      expect(mockFrom).toHaveBeenCalled()
+      expect(mockWhere).toHaveBeenCalled()
+      expect(mockLimit).toHaveBeenCalledWith(1)
 
       expect(credential).toEqual(mockCredential)
     })
 
     it('should return undefined when credential is not found', async () => {
-      mockDbTyped.limit.mockReturnValueOnce([])
+      mockSelectChain([])
 
       const credential = await getCredential('request-id', 'nonexistent-id', 'test-user-id')
 
@@ -102,11 +115,12 @@ describe('OAuth Utils', () => {
         refreshToken: 'new-refresh-token',
       })
 
+      mockUpdateChain()
+
       const result = await refreshTokenIfNeeded('request-id', mockCredential, 'credential-id')
 
       expect(mockRefreshOAuthToken).toHaveBeenCalledWith('google', 'refresh-token')
-      expect(mockDbTyped.update).toHaveBeenCalled()
-      expect(mockDbTyped.set).toHaveBeenCalled()
+      expect(mockDb.update).toHaveBeenCalled()
       expect(result).toEqual({ accessToken: 'new-token', refreshed: true })
     })
 
@@ -152,7 +166,7 @@ describe('OAuth Utils', () => {
         providerId: 'google',
         userId: 'test-user-id',
       }
-      mockDbTyped.limit.mockReturnValueOnce([mockCredential])
+      mockSelectChain([mockCredential])
 
       const token = await refreshAccessTokenIfNeeded('credential-id', 'test-user-id', 'request-id')
 
@@ -169,7 +183,8 @@ describe('OAuth Utils', () => {
         providerId: 'google',
         userId: 'test-user-id',
       }
-      mockDbTyped.limit.mockReturnValueOnce([mockCredential])
+      mockSelectChain([mockCredential])
+      mockUpdateChain()
 
       mockRefreshOAuthToken.mockResolvedValueOnce({
         accessToken: 'new-token',
@@ -180,13 +195,12 @@ describe('OAuth Utils', () => {
       const token = await refreshAccessTokenIfNeeded('credential-id', 'test-user-id', 'request-id')
 
       expect(mockRefreshOAuthToken).toHaveBeenCalledWith('google', 'refresh-token')
-      expect(mockDbTyped.update).toHaveBeenCalled()
-      expect(mockDbTyped.set).toHaveBeenCalled()
+      expect(mockDb.update).toHaveBeenCalled()
       expect(token).toBe('new-token')
     })
 
     it('should return null if credential not found', async () => {
-      mockDbTyped.limit.mockReturnValueOnce([])
+      mockSelectChain([])
 
       const token = await refreshAccessTokenIfNeeded('nonexistent-id', 'test-user-id', 'request-id')
 
@@ -202,7 +216,7 @@ describe('OAuth Utils', () => {
         providerId: 'google',
         userId: 'test-user-id',
       }
-      mockDbTyped.limit.mockReturnValueOnce([mockCredential])
+      mockSelectChain([mockCredential])
 
       mockRefreshOAuthToken.mockResolvedValueOnce(null)
 

apps/sim/app/api/knowledge/search/utils.test.ts

@@ -4,16 +4,12 @@
  *
  * @vitest-environment node
  */
-import { createEnvMock, createMockLogger } from '@sim/testing'
+import { createEnvMock, databaseMock, loggerMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const loggerMock = vi.hoisted(() => ({
-  createLogger: () => createMockLogger(),
-}))
-
 vi.mock('drizzle-orm')
 vi.mock('@sim/logger', () => loggerMock)
-vi.mock('@sim/db')
+vi.mock('@sim/db', () => databaseMock)
 vi.mock('@/lib/knowledge/documents/utils', () => ({
   retryWithExponentialBackoff: (fn: any) => fn(),
 }))

apps/sim/app/api/schedules/[id]/route.test.ts

@@ -3,17 +3,14 @@
  *
  * @vitest-environment node
  */
-import { loggerMock } from '@sim/testing'
+import { databaseMock, loggerMock } from '@sim/testing'
 import { NextRequest } from 'next/server'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockGetSession, mockAuthorizeWorkflowByWorkspacePermission, mockDbSelect, mockDbUpdate } =
-  vi.hoisted(() => ({
-    mockGetSession: vi.fn(),
-    mockAuthorizeWorkflowByWorkspacePermission: vi.fn(),
-    mockDbSelect: vi.fn(),
-    mockDbUpdate: vi.fn(),
-  }))
+const { mockGetSession, mockAuthorizeWorkflowByWorkspacePermission } = vi.hoisted(() => ({
+  mockGetSession: vi.fn(),
+  mockAuthorizeWorkflowByWorkspacePermission: vi.fn(),
+}))
 
 vi.mock('@/lib/auth', () => ({
   getSession: mockGetSession,
@@ -23,12 +20,7 @@ vi.mock('@/lib/workflows/utils', () => ({
   authorizeWorkflowByWorkspacePermission: mockAuthorizeWorkflowByWorkspacePermission,
 }))
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: mockDbSelect,
-    update: mockDbUpdate,
-  },
-}))
+vi.mock('@sim/db', () => databaseMock)
 
 vi.mock('@sim/db/schema', () => ({
   workflow: { id: 'id', userId: 'userId', workspaceId: 'workspaceId' },
@@ -59,6 +51,9 @@ function createParams(id: string): { params: Promise<{ id: string }> } {
   return { params: Promise.resolve({ id }) }
 }
 
+const mockDbSelect = databaseMock.db.select as ReturnType<typeof vi.fn>
+const mockDbUpdate = databaseMock.db.update as ReturnType<typeof vi.fn>
+
 function mockDbChain(selectResults: unknown[][]) {
   let selectCallIndex = 0
   mockDbSelect.mockImplementation(() => ({

apps/sim/app/api/schedules/route.test.ts

@@ -3,17 +3,14 @@
  *
  * @vitest-environment node
  */
-import { loggerMock } from '@sim/testing'
+import { databaseMock, loggerMock } from '@sim/testing'
 import { NextRequest } from 'next/server'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockGetSession, mockAuthorizeWorkflowByWorkspacePermission, mockDbSelect } = vi.hoisted(
-  () => ({
-    mockGetSession: vi.fn(),
-    mockAuthorizeWorkflowByWorkspacePermission: vi.fn(),
-    mockDbSelect: vi.fn(),
-  })
-)
+const { mockGetSession, mockAuthorizeWorkflowByWorkspacePermission } = vi.hoisted(() => ({
+  mockGetSession: vi.fn(),
+  mockAuthorizeWorkflowByWorkspacePermission: vi.fn(),
+}))
 
 vi.mock('@/lib/auth', () => ({
   getSession: mockGetSession,
@@ -23,11 +20,7 @@ vi.mock('@/lib/workflows/utils', () => ({
   authorizeWorkflowByWorkspacePermission: mockAuthorizeWorkflowByWorkspacePermission,
 }))
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: mockDbSelect,
-  },
-}))
+vi.mock('@sim/db', () => databaseMock)
 
 vi.mock('@sim/db/schema', () => ({
   workflow: { id: 'id', userId: 'userId', workspaceId: 'workspaceId' },
@@ -62,6 +55,8 @@ function createRequest(url: string): NextRequest {
   return new NextRequest(new URL(url), { method: 'GET' })
 }
 
+const mockDbSelect = databaseMock.db.select as ReturnType<typeof vi.fn>
+
 function mockDbChain(results: any[]) {
   let callIndex = 0
   mockDbSelect.mockImplementation(() => ({

apps/sim/app/api/workflows/[id]/route.test.ts

@@ -5,7 +5,7 @@
  * @vitest-environment node
  */
 
-import { loggerMock } from '@sim/testing'
+import { loggerMock, setupGlobalFetchMock } from '@sim/testing'
 import { NextRequest } from 'next/server'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -284,9 +284,7 @@ describe('Workflow By ID API Route', () => {
         where: vi.fn().mockResolvedValue([{ id: 'workflow-123' }]),
       })
 
-      global.fetch = vi.fn().mockResolvedValue({
-        ok: true,
-      })
+      setupGlobalFetchMock({ ok: true })
 
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'DELETE',
@@ -331,9 +329,7 @@ describe('Workflow By ID API Route', () => {
         where: vi.fn().mockResolvedValue([{ id: 'workflow-123' }]),
       })
 
-      global.fetch = vi.fn().mockResolvedValue({
-        ok: true,
-      })
+      setupGlobalFetchMock({ ok: true })
 
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'DELETE',

apps/sim/executor/handlers/agent/agent-handler.test.ts

@@ -1,3 +1,4 @@
+import { setupGlobalFetchMock } from '@sim/testing'
 import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
 import { getAllBlocks } from '@/blocks'
 import { BlockType, isMcpTool } from '@/executor/constants'
@@ -61,6 +62,30 @@ vi.mock('@/providers', () => ({
   }),
 }))
 
+vi.mock('@/executor/utils/http', () => ({
+  buildAuthHeaders: vi.fn().mockResolvedValue({ 'Content-Type': 'application/json' }),
+  buildAPIUrl: vi.fn((path: string, params?: Record<string, string>) => {
+    const url = new URL(path, 'http://localhost:3000')
+    if (params) {
+      for (const [key, value] of Object.entries(params)) {
+        if (value !== undefined && value !== null) {
+          url.searchParams.set(key, value)
+        }
+      }
+    }
+    return url
+  }),
+  extractAPIErrorMessage: vi.fn(async (response: Response) => {
+    const defaultMessage = `API request failed with status ${response.status}`
+    try {
+      const errorData = await response.json()
+      return errorData.error || defaultMessage
+    } catch {
+      return defaultMessage
+    }
+  }),
+}))
+
 vi.mock('@sim/db', () => ({
   db: {
     select: vi.fn().mockReturnValue({
@@ -84,7 +109,7 @@ vi.mock('@sim/db/schema', () => ({
   },
 }))
 
-global.fetch = Object.assign(vi.fn(), { preconnect: vi.fn() }) as typeof fetch
+setupGlobalFetchMock()
 
 const mockGetAllBlocks = getAllBlocks as Mock
 const mockExecuteTool = executeTool as Mock
@@ -1901,5 +1926,301 @@ describe('AgentBlockHandler', () => {
 
       expect(discoveryCalls[0].url).toContain('serverId=mcp-legacy-server')
     })
+
+    describe('customToolId resolution - DB as source of truth', () => {
+      const staleInlineSchema = {
+        function: {
+          name: 'formatReport',
+          description: 'Formats a report',
+          parameters: {
+            type: 'object',
+            properties: {
+              title: { type: 'string', description: 'Report title' },
+              content: { type: 'string', description: 'Report content' },
+            },
+            required: ['title', 'content'],
+          },
+        },
+      }
+
+      const dbSchema = {
+        function: {
+          name: 'formatReport',
+          description: 'Formats a report',
+          parameters: {
+            type: 'object',
+            properties: {
+              title: { type: 'string', description: 'Report title' },
+              content: { type: 'string', description: 'Report content' },
+              format: { type: 'string', description: 'Output format' },
+            },
+            required: ['title', 'content', 'format'],
+          },
+        },
+      }
+
+      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')) {
+            return Promise.resolve({
+              ok: true,
+              headers: { get: () => null },
+              json: () =>
+                Promise.resolve({
+                  data: [
+                    {
+                      id: toolId,
+                      title: 'formatReport',
+                      schema: dbSchema,
+                      code: dbCode,
+                    },
+                  ],
+                }),
+            })
+          }
+          return Promise.resolve({
+            ok: true,
+            headers: { get: () => null },
+            json: () => Promise.resolve({}),
+          })
+        })
+      }
+
+      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({}),
+          })
+        })
+      }
+
+      beforeEach(() => {
+        Object.defineProperty(global, 'window', {
+          value: undefined,
+          writable: true,
+          configurable: true,
+        })
+      })
+
+      it('should always fetch latest schema from DB when customToolId is present', async () => {
+        const toolId = 'custom-tool-123'
+        mockFetchForCustomTool(toolId)
+
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Format a report',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              customToolId: toolId,
+              title: 'formatReport',
+              schema: staleInlineSchema,
+              code: staleInlineCode,
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        await handler.execute(mockContext, mockBlock, inputs)
+
+        expect(mockExecuteProviderRequest).toHaveBeenCalled()
+        const providerCall = mockExecuteProviderRequest.mock.calls[0]
+        const tools = providerCall[1].tools
+
+        expect(tools.length).toBe(1)
+        // DB schema wins over stale inline — includes format param
+        expect(tools[0].parameters.required).toContain('format')
+        expect(tools[0].parameters.properties).toHaveProperty('format')
+      })
+
+      it('should fetch from DB when customToolId has no inline schema', async () => {
+        const toolId = 'custom-tool-123'
+        mockFetchForCustomTool(toolId)
+
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Format a report',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              customToolId: toolId,
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        await handler.execute(mockContext, mockBlock, inputs)
+
+        expect(mockExecuteProviderRequest).toHaveBeenCalled()
+        const providerCall = mockExecuteProviderRequest.mock.calls[0]
+        const tools = providerCall[1].tools
+
+        expect(tools.length).toBe(1)
+        expect(tools[0].name).toBe('formatReport')
+        expect(tools[0].parameters.required).toContain('format')
+      })
+
+      it('should fall back to inline schema when DB fetch fails and inline exists', async () => {
+        mockFetchFailure()
+
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Format a report',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              customToolId: 'custom-tool-123',
+              title: 'formatReport',
+              schema: staleInlineSchema,
+              code: staleInlineCode,
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        await handler.execute(mockContext, mockBlock, inputs)
+
+        expect(mockExecuteProviderRequest).toHaveBeenCalled()
+        const providerCall = mockExecuteProviderRequest.mock.calls[0]
+        const tools = providerCall[1].tools
+
+        expect(tools.length).toBe(1)
+        expect(tools[0].name).toBe('formatReport')
+        expect(tools[0].parameters.required).not.toContain('format')
+      })
+
+      it('should return null when DB fetch fails and no inline schema exists', async () => {
+        mockFetchFailure()
+
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Format a report',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              customToolId: 'custom-tool-123',
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        await handler.execute(mockContext, mockBlock, inputs)
+
+        expect(mockExecuteProviderRequest).toHaveBeenCalled()
+        const providerCall = mockExecuteProviderRequest.mock.calls[0]
+        const tools = providerCall[1].tools
+
+        expect(tools.length).toBe(0)
+      })
+
+      it('should use DB code for executeFunction when customToolId resolves', async () => {
+        const toolId = 'custom-tool-123'
+        mockFetchForCustomTool(toolId)
+
+        let capturedTools: any[] = []
+        Promise.all = vi.fn().mockImplementation((promises: Promise<any>[]) => {
+          const result = originalPromiseAll.call(Promise, promises)
+          result.then((tools: any[]) => {
+            if (tools?.length) {
+              capturedTools = tools.filter((t) => t !== null)
+            }
+          })
+          return result
+        })
+
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Format a report',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              customToolId: toolId,
+              title: 'formatReport',
+              schema: staleInlineSchema,
+              code: staleInlineCode,
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        await handler.execute(mockContext, mockBlock, inputs)
+
+        expect(capturedTools.length).toBe(1)
+        expect(typeof capturedTools[0].executeFunction).toBe('function')
+
+        await capturedTools[0].executeFunction({ title: 'Q1', format: 'pdf' })
+
+        expect(mockExecuteTool).toHaveBeenCalledWith(
+          'function_execute',
+          expect.objectContaining({
+            code: dbCode,
+          }),
+          false,
+          expect.any(Object)
+        )
+      })
+
+      it('should not fetch from DB when no customToolId is present', async () => {
+        const inputs = {
+          model: 'gpt-4o',
+          userPrompt: 'Use the tool',
+          apiKey: 'test-api-key',
+          tools: [
+            {
+              type: 'custom-tool',
+              title: 'formatReport',
+              schema: staleInlineSchema,
+              code: staleInlineCode,
+              usageControl: 'auto' as const,
+            },
+          ],
+        }
+
+        mockGetProviderFromModel.mockReturnValue('openai')
+
+        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(mockExecuteProviderRequest).toHaveBeenCalled()
+        const providerCall = mockExecuteProviderRequest.mock.calls[0]
+        const tools = providerCall[1].tools
+
+        expect(tools.length).toBe(1)
+        expect(tools[0].name).toBe('formatReport')
+        expect(tools[0].parameters.required).not.toContain('format')
+      })
+    })
   })
 })

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

@@ -272,15 +272,16 @@ export class AgentBlockHandler implements BlockHandler {
     let code = tool.code
     let title = tool.title
 
-    if (tool.customToolId && !schema) {
+    if (tool.customToolId) {
       const resolved = await this.fetchCustomToolById(ctx, tool.customToolId)
-      if (!resolved) {
+      if (resolved) {
+        schema = resolved.schema
+        code = resolved.code
+        title = resolved.title
+      } else if (!schema) {
         logger.error(`Custom tool not found: ${tool.customToolId}`)
         return null
       }
-      schema = resolved.schema
-      code = resolved.code
-      title = resolved.title
     }
 
     if (!schema?.function) {

apps/sim/executor/handlers/workflow/workflow-handler.test.ts

@@ -1,3 +1,4 @@
+import { setupGlobalFetchMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
 import { BlockType } from '@/executor/constants'
 import { WorkflowBlockHandler } from '@/executor/handlers/workflow/workflow-handler'
@@ -9,7 +10,7 @@ vi.mock('@/lib/auth/internal', () => ({
 }))
 
 // Mock fetch globally
-global.fetch = vi.fn()
+setupGlobalFetchMock()
 
 describe('WorkflowBlockHandler', () => {
   let handler: WorkflowBlockHandler

apps/sim/lib/messaging/email/mailer.test.ts

@@ -1,4 +1,4 @@
-import { createEnvMock, createMockLogger } from '@sim/testing'
+import { createEnvMock, loggerMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
 
 /**
@@ -10,10 +10,6 @@ import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
  * mock functions can intercept.
  */
 
-const loggerMock = vi.hoisted(() => ({
-  createLogger: () => createMockLogger(),
-}))
-
 const mockSend = vi.fn()
 const mockBatchSend = vi.fn()
 const mockAzureBeginSend = vi.fn()

apps/sim/lib/messaging/email/unsubscribe.test.ts

@@ -1,20 +1,8 @@
-import { createEnvMock, createMockLogger } from '@sim/testing'
+import { createEnvMock, databaseMock, loggerMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import type { EmailType } from '@/lib/messaging/email/mailer'
 
-const loggerMock = vi.hoisted(() => ({
-  createLogger: () => createMockLogger(),
-}))
-
-const mockDb = vi.hoisted(() => ({
-  select: vi.fn(),
-  insert: vi.fn(),
-  update: vi.fn(),
-}))
-
-vi.mock('@sim/db', () => ({
-  db: mockDb,
-}))
+vi.mock('@sim/db', () => databaseMock)
 
 vi.mock('@sim/db/schema', () => ({
   user: { id: 'id', email: 'email' },
@@ -30,6 +18,8 @@ vi.mock('drizzle-orm', () => ({
   eq: vi.fn((a, b) => ({ type: 'eq', left: a, right: b })),
 }))
 
+const mockDb = databaseMock.db as Record<string, ReturnType<typeof vi.fn>>
+
 vi.mock('@/lib/core/config/env', () => createEnvMock({ BETTER_AUTH_SECRET: 'test-secret-key' }))
 
 vi.mock('@sim/logger', () => loggerMock)

apps/sim/lib/workflows/diff/diff-engine.test.ts

@@ -1,18 +1,11 @@
 /**
  * @vitest-environment node
  */
+import { loggerMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types'
 
-// Mock all external dependencies before imports
-vi.mock('@sim/logger', () => ({
-  createLogger: () => ({
-    info: vi.fn(),
-    warn: vi.fn(),
-    error: vi.fn(),
-    debug: vi.fn(),
-  }),
-}))
+vi.mock('@sim/logger', () => loggerMock)
 
 vi.mock('@/stores/workflows/workflow/store', () => ({
   useWorkflowStore: {

apps/sim/lib/workflows/utils.test.ts

@@ -14,22 +14,15 @@ import {
   databaseMock,
   expectWorkflowAccessDenied,
   expectWorkflowAccessGranted,
+  mockAuth,
 } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => databaseMock)
-
-// Mock the auth module
-vi.mock('@/lib/auth', () => ({
-  getSession: vi.fn(),
-}))
-
-import { db } from '@sim/db'
-import { getSession } from '@/lib/auth'
-// Import after mocks are set up
-import { validateWorkflowPermissions } from '@/lib/workflows/utils'
+const mockDb = databaseMock.db
 
 describe('validateWorkflowPermissions', () => {
+  const auth = mockAuth()
+
   const mockSession = createSession({ userId: 'user-1', email: 'user1@test.com' })
   const mockWorkflow = createWorkflowRecord({
     id: 'wf-1',
@@ -42,13 +35,17 @@ describe('validateWorkflowPermissions', () => {
   })
 
   beforeEach(() => {
+    vi.resetModules()
     vi.clearAllMocks()
+
+    vi.doMock('@sim/db', () => databaseMock)
   })
 
   describe('authentication', () => {
     it('should return 401 when no session exists', async () => {
-      vi.mocked(getSession).mockResolvedValue(null)
+      auth.setUnauthenticated()
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 401)
@@ -56,8 +53,9 @@ describe('validateWorkflowPermissions', () => {
     })
 
     it('should return 401 when session has no user id', async () => {
-      vi.mocked(getSession).mockResolvedValue({ user: {} } as any)
+      auth.mockGetSession.mockResolvedValue({ user: {} } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 401)
@@ -66,14 +64,14 @@ describe('validateWorkflowPermissions', () => {
 
   describe('workflow not found', () => {
     it('should return 404 when workflow does not exist', async () => {
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
 
-      // Mock workflow query to return empty
       const mockLimit = vi.fn().mockResolvedValue([])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('non-existent', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 404)
@@ -83,43 +81,42 @@ describe('validateWorkflowPermissions', () => {
 
   describe('owner access', () => {
     it('should deny access to workflow owner without workspace permissions for read action', async () => {
-      const ownerSession = createSession({ userId: 'owner-1' })
-      vi.mocked(getSession).mockResolvedValue(ownerSession as any)
+      auth.setAuthenticated({ id: 'owner-1', email: 'owner-1@test.com' })
 
-      // Mock workflow query
       const mockLimit = vi.fn().mockResolvedValue([mockWorkflow])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 403)
     })
 
     it('should deny access to workflow owner without workspace permissions for write action', async () => {
-      const ownerSession = createSession({ userId: 'owner-1' })
-      vi.mocked(getSession).mockResolvedValue(ownerSession as any)
+      auth.setAuthenticated({ id: 'owner-1', email: 'owner-1@test.com' })
 
       const mockLimit = vi.fn().mockResolvedValue([mockWorkflow])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'write')
 
       expectWorkflowAccessDenied(result, 403)
     })
 
     it('should deny access to workflow owner without workspace permissions for admin action', async () => {
-      const ownerSession = createSession({ userId: 'owner-1' })
-      vi.mocked(getSession).mockResolvedValue(ownerSession as any)
+      auth.setAuthenticated({ id: 'owner-1', email: 'owner-1@test.com' })
 
       const mockLimit = vi.fn().mockResolvedValue([mockWorkflow])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'admin')
 
       expectWorkflowAccessDenied(result, 403)
@@ -128,11 +125,10 @@ describe('validateWorkflowPermissions', () => {
 
   describe('workspace member access with permissions', () => {
     beforeEach(() => {
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
     })
 
     it('should grant read access to user with read permission', async () => {
-      // First call: workflow query, second call: workspace owner, third call: permission
       let callCount = 0
       const mockLimit = vi.fn().mockImplementation(() => {
         callCount++
@@ -141,8 +137,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'read')
 
       expectWorkflowAccessGranted(result)
@@ -157,8 +154,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'write')
 
       expectWorkflowAccessDenied(result, 403)
@@ -174,8 +172,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'write')
 
       expectWorkflowAccessGranted(result)
@@ -190,8 +189,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'write')
 
       expectWorkflowAccessGranted(result)
@@ -206,8 +206,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'admin')
 
       expectWorkflowAccessDenied(result, 403)
@@ -223,8 +224,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'admin')
 
       expectWorkflowAccessGranted(result)
@@ -233,18 +235,19 @@ describe('validateWorkflowPermissions', () => {
 
   describe('no workspace permission', () => {
     it('should deny access to user without any workspace permission', async () => {
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
 
       let callCount = 0
       const mockLimit = vi.fn().mockImplementation(() => {
         callCount++
         if (callCount === 1) return Promise.resolve([mockWorkflow])
-        return Promise.resolve([]) // No permission record
+        return Promise.resolve([])
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 403)
@@ -259,13 +262,14 @@ describe('validateWorkflowPermissions', () => {
         workspaceId: null,
       })
 
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
 
       const mockLimit = vi.fn().mockResolvedValue([workflowWithoutWorkspace])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-2', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 403)
@@ -278,13 +282,14 @@ describe('validateWorkflowPermissions', () => {
         workspaceId: null,
       })
 
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
 
       const mockLimit = vi.fn().mockResolvedValue([workflowWithoutWorkspace])
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-2', 'req-1', 'read')
 
       expectWorkflowAccessDenied(result, 403)
@@ -293,7 +298,7 @@ describe('validateWorkflowPermissions', () => {
 
   describe('default action', () => {
     it('should default to read action when not specified', async () => {
-      vi.mocked(getSession).mockResolvedValue(mockSession as any)
+      auth.mockGetSession.mockResolvedValue(mockSession as any)
 
       let callCount = 0
       const mockLimit = vi.fn().mockImplementation(() => {
@@ -303,8 +308,9 @@ describe('validateWorkflowPermissions', () => {
       })
       const mockWhere = vi.fn(() => ({ limit: mockLimit }))
       const mockFrom = vi.fn(() => ({ where: mockWhere }))
-      vi.mocked(db.select).mockReturnValue({ from: mockFrom } as any)
+      vi.mocked(mockDb.select).mockReturnValue({ from: mockFrom } as any)
 
+      const { validateWorkflowPermissions } = await import('@/lib/workflows/utils')
       const result = await validateWorkflowPermissions('wf-1', 'req-1')
 
       expectWorkflowAccessGranted(result)

apps/sim/lib/workspaces/permissions/utils.test.ts

@@ -1,17 +1,7 @@
-import { drizzleOrmMock } from '@sim/testing/mocks'
+import { databaseMock, drizzleOrmMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn(),
-    from: vi.fn(),
-    where: vi.fn(),
-    limit: vi.fn(),
-    innerJoin: vi.fn(),
-    leftJoin: vi.fn(),
-    orderBy: vi.fn(),
-  },
-}))
+vi.mock('@sim/db', () => databaseMock)
 
 vi.mock('@sim/db/schema', () => ({
   permissions: {