fix(custom-tools): count skipped titles and skip title conflicts atomically

Author: waleedlatif1

apps/sim/lib/workflows/custom-tools/operations.ts

@@ -17,23 +17,13 @@ import {
 } from '@/lib/api/list-query'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { assertValidCustomToolDeclaration } from '@/lib/custom-tools/schema'
+import type { DbTransaction } from '@/lib/db/types'
 
 const logger = createLogger('CustomToolsOperations')
 
 export type CustomToolSortBy = 'title' | 'createdAt' | 'updatedAt'
 
-/**
- * What `upsertCustomTools` does with a tool it would insert under a title the
- * workspace already holds: `'error'` throws and rolls back the whole batch,
- * `'skip'` leaves the existing row untouched and persists the rest.
- */
-type CustomToolTitleConflictPolicy = 'error' | 'skip'
-
-/**
- * Internal function to create/update custom tools
- * Can be called from API routes or internal services
- */
-export async function upsertCustomTools(params: {
+interface CustomToolWriteParams {
   tools: Array<{
     id?: string
     title: string
@@ -43,112 +33,162 @@ export async function upsertCustomTools(params: {
   workspaceId: string
   userId: string
   requestId?: string
-  onTitleConflict?: CustomToolTitleConflictPolicy
-}) {
-  const {
-    tools,
-    workspaceId,
-    userId,
-    requestId = generateRequestId(),
-    onTitleConflict = 'error',
-  } = params
-
-  /**
-   * Ahead of the transaction so a batch is rejected whole rather than landing
-   * the tools that preceded the unstorable one.
-   */
+}
+
+/**
+ * What a batch write does with a tool it would insert under a title the
+ * workspace already holds: `'error'` throws and rolls back the whole batch,
+ * `'skip'` leaves the existing row untouched and writes the rest.
+ */
+type CustomToolTitleConflictPolicy = 'error' | 'skip'
+
+/**
+ * Ahead of the transaction so a batch is rejected whole rather than landing
+ * the tools that preceded the unstorable one.
+ */
+function assertStorableBatch(tools: CustomToolWriteParams['tools']) {
   for (const tool of tools) {
     assertValidCustomToolDeclaration(tool.schema)
   }
+}
 
-  return await db.transaction(async (tx) => {
-    for (const tool of tools) {
-      const nowTime = new Date()
+/** Writes the batch inside `tx`, returning the titles skipped under the `'skip'` policy. */
+async function writeCustomTools(
+  tx: DbTransaction,
+  params: CustomToolWriteParams,
+  onTitleConflict: CustomToolTitleConflictPolicy
+): Promise<string[]> {
+  const { tools, workspaceId, userId, requestId = generateRequestId() } = params
+  const skippedTitles: string[] = []
+
+  for (const tool of tools) {
+    const nowTime = new Date()
 
-      if (tool.id) {
-        const existingWorkspaceTool = await tx
-          .select()
-          .from(customTools)
+    if (tool.id) {
+      const existingWorkspaceTool = await tx
+        .select()
+        .from(customTools)
+        .where(and(eq(customTools.id, tool.id), eq(customTools.workspaceId, workspaceId)))
+        .limit(1)
+
+      if (existingWorkspaceTool.length > 0) {
+        await tx
+          .update(customTools)
+          .set({
+            title: tool.title,
+            schema: tool.schema,
+            code: tool.code,
+            updatedAt: nowTime,
+          })
           .where(and(eq(customTools.id, tool.id), eq(customTools.workspaceId, workspaceId)))
-          .limit(1)
-
-        if (existingWorkspaceTool.length > 0) {
-          await tx
-            .update(customTools)
-            .set({
-              title: tool.title,
-              schema: tool.schema,
-              code: tool.code,
-              updatedAt: nowTime,
-            })
-            .where(and(eq(customTools.id, tool.id), eq(customTools.workspaceId, workspaceId)))
-          continue
-        }
-
-        const existingLegacyTool = await tx
-          .select()
-          .from(customTools)
-          .where(
-            and(
-              eq(customTools.id, tool.id),
-              isNull(customTools.workspaceId),
-              eq(customTools.userId, userId)
-            )
-          )
-          .limit(1)
-
-        if (existingLegacyTool.length > 0) {
-          await tx
-            .update(customTools)
-            .set({
-              title: tool.title,
-              schema: tool.schema,
-              code: tool.code,
-              updatedAt: nowTime,
-            })
-            .where(eq(customTools.id, tool.id))
-
-          logger.info(`[${requestId}] Updated legacy tool ${tool.id}`)
-          continue
-        }
+        continue
       }
 
-      const duplicateTitle = await tx
+      const existingLegacyTool = await tx
         .select()
         .from(customTools)
-        .where(and(eq(customTools.workspaceId, workspaceId), eq(customTools.title, tool.title)))
+        .where(
+          and(
+            eq(customTools.id, tool.id),
+            isNull(customTools.workspaceId),
+            eq(customTools.userId, userId)
+          )
+        )
         .limit(1)
 
-      if (duplicateTitle.length > 0) {
-        if (onTitleConflict === 'skip') {
-          logger.info(`[${requestId}] Skipped custom tool "${tool.title}": title already exists`)
-          continue
-        }
-        throw new Error(`A tool with the title "${tool.title}" already exists in this workspace`)
+      if (existingLegacyTool.length > 0) {
+        await tx
+          .update(customTools)
+          .set({
+            title: tool.title,
+            schema: tool.schema,
+            code: tool.code,
+            updatedAt: nowTime,
+          })
+          .where(eq(customTools.id, tool.id))
+
+        logger.info(`[${requestId}] Updated legacy tool ${tool.id}`)
+        continue
       }
+    }
+
+    const row = {
+      id: generateShortId(),
+      workspaceId,
+      userId,
+      title: tool.title,
+      schema: tool.schema,
+      code: tool.code,
+      createdAt: nowTime,
+      updatedAt: nowTime,
+    }
 
-      await tx.insert(customTools).values({
-        id: generateShortId(),
-        workspaceId,
-        userId,
-        title: tool.title,
-        schema: tool.schema,
-        code: tool.code,
-        createdAt: nowTime,
-        updatedAt: nowTime,
-      })
+    if (onTitleConflict === 'skip') {
+      /** A conflict-target insert, so a title a concurrent save just committed is skipped too. */
+      const inserted = await tx
+        .insert(customTools)
+        .values(row)
+        .onConflictDoNothing({ target: [customTools.workspaceId, customTools.title] })
+        .returning({ id: customTools.id })
+      if (inserted.length === 0) {
+        logger.info(`[${requestId}] Skipped custom tool "${tool.title}": title already exists`)
+        skippedTitles.push(tool.title)
+      }
+      continue
     }
 
+    const duplicateTitle = await tx
+      .select()
+      .from(customTools)
+      .where(and(eq(customTools.workspaceId, workspaceId), eq(customTools.title, tool.title)))
+      .limit(1)
+
+    if (duplicateTitle.length > 0) {
+      throw new Error(`A tool with the title "${tool.title}" already exists in this workspace`)
+    }
+
+    await tx.insert(customTools).values(row)
+  }
+
+  return skippedTitles
+}
+
+/**
+ * Internal function to create/update custom tools
+ * Can be called from API routes or internal services
+ */
+export async function upsertCustomTools(params: CustomToolWriteParams) {
+  assertStorableBatch(params.tools)
+
+  return await db.transaction(async (tx) => {
+    await writeCustomTools(tx, params, 'error')
+
     const resultTools = await tx
       .select()
       .from(customTools)
-      .where(eq(customTools.workspaceId, workspaceId))
+      .where(eq(customTools.workspaceId, params.workspaceId))
       .orderBy(desc(customTools.createdAt))
 
     return resultTools
   })
 }
 
+/**
+ * `upsertCustomTools` for writes that re-offer tools they may already have
+ * stored (workflow saves and imports): a tool whose title the workspace already
+ * holds is skipped, leaving that row untouched, instead of failing the batch.
+ *
+ * @returns The titles that were skipped.
+ */
+export async function upsertCustomToolsSkippingTakenTitles(
+  params: CustomToolWriteParams
+): Promise<{ skippedTitles: string[] }> {
+  assertStorableBatch(params.tools)
+
+  const skippedTitles = await db.transaction((tx) => writeCustomTools(tx, params, 'skip'))
+  return { skippedTitles }
+}
+
 export async function listCustomTools(params: { userId: string; workspaceId?: string }) {
   const { userId, workspaceId } = params
   return workspaceId

apps/sim/lib/workflows/persistence/custom-tools-persistence.integration.ts

@@ -94,7 +94,7 @@ describe('extractAndPersistCustomTools', () => {
       owner
     )
 
-    expect(second.errors).toEqual([])
+    expect(second).toEqual({ saved: 1, errors: [] })
     expect(await workspaceTools()).toEqual([
       { title: 'Lookup order', code: 'return 1' },
       { title: 'Refund order', code: 'return 2' },
@@ -114,9 +114,22 @@ describe('extractAndPersistCustomTools', () => {
       owner
     )
 
-    expect(result.errors).toEqual([])
+    expect(result).toEqual({ saved: 0, errors: [] })
     expect(await workspaceTools()).toEqual([{ title: 'Lookup order', code: 'return "library"' }])
   })
+
+  it('lets concurrent saves of the same new inline tool both succeed with one row', async () => {
+    const state = workflowStateWithInlineTools([{ title: 'Lookup order', code: 'return 1' }])
+
+    const results = await Promise.all([
+      extractAndPersistCustomTools(state, workspaceId, owner),
+      extractAndPersistCustomTools(state, workspaceId, owner),
+    ])
+
+    expect(results.map((result) => result.errors)).toEqual([[], []])
+    expect(results.map((result) => result.saved).sort()).toEqual([0, 1])
+    expect(await workspaceTools()).toEqual([{ title: 'Lookup order', code: 'return 1' }])
+  })
 })
 
 describe('upsertCustomTools', () => {

apps/sim/lib/workflows/persistence/custom-tools-persistence.test.ts

@@ -5,7 +5,7 @@ const { mockUpsertCustomTools } = vi.hoisted(() => ({
 }))
 
 vi.mock('@/lib/workflows/custom-tools/operations', () => ({
-  upsertCustomTools: mockUpsertCustomTools,
+  upsertCustomToolsSkippingTakenTitles: mockUpsertCustomTools,
 }))
 
 import { persistCustomToolsToDatabase } from '@/lib/workflows/persistence/custom-tools-persistence'
@@ -34,7 +34,7 @@ function customTool(overrides: Record<string, unknown> = {}) {
 
 describe('persistCustomToolsToDatabase', () => {
   beforeEach(() => {
-    mockUpsertCustomTools.mockResolvedValue([])
+    mockUpsertCustomTools.mockResolvedValue({ skippedTitles: [] })
   })
 
   it('skips a declaration missing the function discriminator the public API republishes', async () => {

apps/sim/lib/workflows/persistence/custom-tools-persistence.ts

@@ -1,7 +1,7 @@
 import { createLogger } from '@sim/logger'
 import { toError } from '@sim/utils/errors'
 import { isValidCustomToolDeclaration } from '@/lib/custom-tools/schema'
-import { upsertCustomTools } from '@/lib/workflows/custom-tools/operations'
+import { upsertCustomToolsSkippingTakenTitles } from '@/lib/workflows/custom-tools/operations'
 
 const logger = createLogger('CustomToolsPersistence')
 
@@ -164,7 +164,7 @@ export async function persistCustomToolsToDatabase(
   }
 
   try {
-    await upsertCustomTools({
+    const { skippedTitles } = await upsertCustomToolsSkippingTakenTitles({
       tools: validTools.map((tool) => ({
         id: tool.toolId,
         title: tool.title,
@@ -173,11 +173,10 @@ export async function persistCustomToolsToDatabase(
       })),
       workspaceId,
       userId,
-      onTitleConflict: 'skip',
     })
 
-    saved = validTools.length
-    logger.info(`Persisted ${saved} custom tool(s)`, { workspaceId })
+    saved = validTools.length - skippedTitles.length
+    logger.info(`Persisted ${saved} custom tool(s)`, { workspaceId, skippedTitles })
   } catch (error) {
     const errorMsg = `Failed to persist custom tools: ${toError(error).message}`
     logger.error(errorMsg, { error })