address comments

Author: icecrasher321

apps/sim/app/api/mcp/serve/[serverId]/route.ts

@@ -560,9 +560,7 @@ async function handleToolsList(
       .orderBy(asc(workflowMcpTool.id))
       .limit(MAX_MCP_TOOLS_LIST_COUNT + 1)
 
-    const hasNextPage = toolSizes.length > MAX_MCP_TOOLS_LIST_COUNT
-    const pageSizes = hasNextPage ? toolSizes.slice(0, MAX_MCP_TOOLS_LIST_COUNT) : toolSizes
-    const nextCursor = hasNextPage ? pageSizes.at(-1)?.id : undefined
+    const pageSizes = toolSizes.slice(0, MAX_MCP_TOOLS_LIST_COUNT)
 
     let estimatedSchemaBytes = 0
     let estimatedMetadataBytes = 0
@@ -602,7 +600,9 @@ async function handleToolsList(
       .orderBy(asc(workflowMcpTool.id))
       .limit(MAX_MCP_TOOLS_LIST_COUNT + 1)
 
+    const hasNextPage = tools.length > MAX_MCP_TOOLS_LIST_COUNT
     const pageTools = tools.slice(0, MAX_MCP_TOOLS_LIST_COUNT)
+    const nextCursor = hasNextPage ? pageTools.at(-1)?.id : undefined
     let schemaBytes = 0
     const result: ListToolsResult = {
       tools: pageTools.map((tool) => {

apps/sim/lib/mcp/server-locks.ts

@@ -6,7 +6,9 @@ const MCP_SERVER_LOCK_TIMEOUT_MS = 3_000
 const LOCK_NOT_AVAILABLE_SQLSTATE = '55P03'
 
 export async function setWorkflowMcpTransactionLockTimeout(tx: DbOrTx): Promise<void> {
-  await tx.execute(sql.raw(`SET LOCAL lock_timeout = '${MCP_SERVER_LOCK_TIMEOUT_MS}ms'`))
+  await tx.execute(
+    sql`select set_config('lock_timeout', ${`${MCP_SERVER_LOCK_TIMEOUT_MS}ms`}, true)`
+  )
 }
 
 export async function acquireWorkflowMcpServerLock(tx: DbOrTx, serverId: string): Promise<void> {

apps/sim/lib/mcp/workflow-mcp-sync.ts

@@ -50,7 +50,8 @@ interface ServerMetadataUsageState {
 async function listWorkflowMcpToolSyncPage(
   tx: DbOrTx,
   workflowId: string,
-  afterToolId?: string
+  afterToolId?: string,
+  serverIds?: string[]
 ): Promise<WorkflowMcpToolSyncRow[]> {
   return tx
     .select({
@@ -64,6 +65,9 @@ async function listWorkflowMcpToolSyncPage(
       and(
         eq(workflowMcpTool.workflowId, workflowId),
         isNull(workflowMcpTool.archivedAt),
+        serverIds && serverIds.length > 0
+          ? inArray(workflowMcpTool.serverId, serverIds)
+          : undefined,
         afterToolId ? gt(workflowMcpTool.id, afterToolId) : undefined
       )
     )
@@ -188,10 +192,13 @@ export async function syncMcpToolsForWorkflow(
 
     const affectedServerIds = new Set<string>()
     const lockedServers = await collectWorkflowMcpToolServerIds(tx, workflowId)
+    if (lockedServers.length === 0) return []
+
     for (const { serverId } of lockedServers) {
       await acquireWorkflowMcpServerLock(tx, serverId)
       affectedServerIds.add(serverId)
     }
+    const lockedServerIds = [...affectedServerIds]
 
     const usageStateByServer = new Map<string, ServerMetadataUsageState>()
     for (const { serverId } of lockedServers) {
@@ -206,7 +213,7 @@ export async function syncMcpToolsForWorkflow(
     let afterToolId: string | undefined
 
     while (true) {
-      const page = await listWorkflowMcpToolSyncPage(tx, workflowId, afterToolId)
+      const page = await listWorkflowMcpToolSyncPage(tx, workflowId, afterToolId, lockedServerIds)
       if (page.length === 0) break
 
       const pageTools = page.slice(0, MCP_SYNC_TOOLS_PAGE_SIZE)
@@ -222,7 +229,9 @@ export async function syncMcpToolsForWorkflow(
         left.localeCompare(right)
       )) {
         const usageState = usageStateByServer.get(serverId)
-        if (!usageState) continue
+        if (!usageState) {
+          throw new Error(`Missing locked MCP server usage state for server ${serverId}`)
+        }
         const schemaToolIds: string[] = []
         const emptySchemaToolIds: string[] = []