fix(mcp): stop restore spending its budget on servers it cannot restore

A server can hold both a live registration and archived ones for the same
workflow: the partial unique index constrains only the live row. Such a server
was counted twice — once shrinking the restore budget, once consuming one of
its slots — before the liveness check under the lock skipped it. While the
bound was the full server limit that waste was invisible; once the bound became
the remaining headroom, every slot spent that way cost a registration that
could have been restored.

The candidate query now excludes servers the workflow is already live on, so
the budget and the candidate set agree. The exclusion sits on the inner stage,
before deduplication and the bound, and is a pre-lock optimisation only: the
check under the server lock stays authoritative, because the query can go stale
between reading and unarchiving.

Candidates rejected for a tool-name collision, the per-server cap, or the
metadata budget are still not replaced. That case is only knowable under the
lock, so replacing it would mean fetching past the bound and locking servers
outside the candidate set, widening an ordering every writer here relies on.

Author: waleedlatif1

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

@@ -94,6 +94,7 @@ interface FakeSelect {
   order: { type: string; column: unknown }[]
   limit: number | null
   source: FakeSelect | null
+  where: unknown
 }
 
 interface FakeOrderNode {
@@ -112,15 +113,24 @@ interface FakeOrderNode {
  * statement first. A test can therefore observe which rows a bounded, possibly
  * two-stage candidate query actually returns rather than being handed a canned
  * answer.
+ *
+ * `liveServerIds` models the one correlated anti-join the restore issues: a
+ * select whose WHERE carries a `notExists` drops the rows whose `serverId` the
+ * workflow already holds a live registration on, exactly as the subquery would.
+ * A statement that does not ask for the anti-join is filtered by nothing, so a
+ * fixture can tell the two apart.
  */
-function createFakeTx(results: QueuedRead[]) {
+function createFakeTx(results: QueuedRead[], liveServerIds: string[] = []) {
   const writes: RecordedWrite[] = []
   const queue = [...results]
+  const liveServers = new Set(liveServerIds)
   let pending: RecordedWrite | null = null
-  let current: FakeSelect | null = null
+  const selects: FakeSelect[] = []
+
+  const top = (): FakeSelect | null => selects.at(-1) ?? null
 
   const startSelect = (projection: Record<string, unknown>, distinctOn: unknown[]): void => {
-    current = { projection, distinctOn, order: [], limit: null, source: null }
+    selects.push({ projection, distinctOn, order: [], limit: null, source: null, where: undefined })
   }
 
   /** Resolves an order/distinct column token to the row key it was selected as. */
@@ -144,6 +154,12 @@ function createFakeTx(results: QueuedRead[]) {
       rows = typeof next === 'function' ? next() : next
     }
 
+    if (hasMockCondition(select.where, (node) => node.type === 'notExists')) {
+      rows = rows.filter(
+        (row) => !liveServers.has((row as { serverId?: string }).serverId as string)
+      )
+    }
+
     if (select.order.length > 0) {
       rows = [...rows].sort((left, right) => {
         for (const node of select.order) {
@@ -186,20 +202,22 @@ function createFakeTx(results: QueuedRead[]) {
     },
     from(source: unknown) {
       const handle = (source as Record<symbol, FakeSelect> | null)?.[FAKE_SUBQUERY]
-      if (current && handle) current.source = handle
+      const select = top()
+      if (select && handle) select.source = handle
       return builder
     },
     orderBy(...nodes: FakeOrderNode[]) {
-      if (current) current.order = nodes
+      const select = top()
+      if (select) select.order = nodes
       return builder
     },
     limit(value: number) {
-      if (current) current.limit = value
+      const select = top()
+      if (select) select.limit = value
       return builder
     },
     as() {
-      const select = current
-      current = null
+      const select = selects.pop() ?? null
       return new Proxy(
         {},
         {
@@ -226,14 +244,24 @@ function createFakeTx(results: QueuedRead[]) {
       }
       return builder
     },
+    /**
+     * A correlated subquery is built — and finished by its own `where` — while
+     * its enclosing statement is still open, so a completed select is popped
+     * here to let the outer statement's `where` land on the outer select.
+     */
     where(condition: unknown) {
-      if (pending) pending.where = condition
-      pending = null
+      if (pending) {
+        pending.where = condition
+        pending = null
+        return builder
+      }
+      while (selects.length > 1 && top()?.where !== undefined) selects.pop()
+      const select = top()
+      if (select && select.where === undefined) select.where = condition
       return builder
     },
     then(onFulfilled: (value: unknown) => unknown) {
-      const select = current
-      current = null
+      const select = selects.pop() ?? null
       const rows = select ? evaluate(select) : (queue.shift() ?? [])
       return Promise.resolve(rows).then(onFulfilled)
     },
@@ -528,6 +556,55 @@ describe('workflow MCP tool withdrawal and restore', () => {
     )
   })
 
+  /**
+   * A `(server, workflow)` pair can hold both a live row and archived
+   * generations: `workflow_mcp_tool_server_workflow_unique` is scoped to
+   * unarchived rows, so it constrains only the live one. Such a server was
+   * counted twice — once by the live-server count that shrinks the budget, once
+   * as a candidate — and the slot it consumed was spent on a candidate the
+   * post-lock liveness check could only ever skip, leaving a genuinely
+   * restorable server beyond the bound unfetched. The candidate query must
+   * exclude it so the budget and the candidate set agree.
+   */
+  it('never spends a restore slot on a server the workflow is already live on', async () => {
+    const headroom = 2
+    const sharedServerId = 'srv-shared'
+    const liveServers = [
+      ...liveServerRows(MAX_MCP_SERVERS_PER_WORKFLOW - headroom - 1),
+      { serverId: sharedServerId },
+    ]
+    const candidates = [
+      archivedRow('t-arch-a', 'srv-arch-a', 'orders_a', new Date(2020, 0, 1)),
+      archivedRow('t-arch-b', 'srv-arch-b', 'orders_b', new Date(2020, 0, 2)),
+      archivedRow('t-arch-shared', sharedServerId, 'orders_shared', new Date(2020, 0, 3)),
+    ]
+
+    const { tx, writes } = createFakeTx(
+      [liveServers, candidates, liveServers, [], [], []],
+      [sharedServerId]
+    )
+
+    await syncMcpToolsForWorkflow({
+      workflowId: WORKFLOW_ID,
+      requestId: REQUEST_ID,
+      state: { blocks: {} },
+      tx,
+      notify: false,
+      throwOnError: true,
+    })
+
+    expect(mocks.acquireLock.mock.calls.map((call) => call[1])).toEqual([
+      'srv-arch-a',
+      'srv-arch-b',
+    ])
+
+    const restore = writes.find((write) => write.values?.archivedAt === null)
+    const restoredIds = flattenMockConditions(restore?.where).find(
+      (node) => node.type === 'inArray'
+    )?.values as string[]
+    expect(restoredIds).toEqual(['t-arch-b', 't-arch-a'])
+  })
+
   /**
    * The headroom that bounds the candidate query is read before any lock is
    * taken, so a manual `tools create` on another server can consume part of it

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

@@ -1,6 +1,7 @@
 import { db, workflowMcpServer, workflowMcpTool } from '@sim/db'
 import { createLogger } from '@sim/logger'
-import { and, asc, desc, eq, gt, inArray, isNotNull, isNull } from 'drizzle-orm'
+import { and, asc, desc, eq, gt, inArray, isNotNull, isNull, notExists } from 'drizzle-orm'
+import { alias } from 'drizzle-orm/pg-core'
 import type { DbOrTx } from '@/lib/db/types'
 import { MAX_MCP_SERVERS_PER_WORKFLOW, MAX_MCP_TOOLS_PER_SERVER } from '@/lib/mcp/constants'
 import { acquireWorkflowMcpServerLock } from '@/lib/mcp/server-locks'
@@ -230,6 +231,18 @@ async function getRestoreSkipReason(
  * locked, and each accepted candidate spends one unit of it, so the restore can
  * never select more servers than the workflow has room for.
  *
+ * The candidate query also excludes, with a correlated `NOT EXISTS`, every
+ * server the workflow already holds a live row on. Those servers are counted by
+ * {@link countWorkflowLiveMcpServers} and so already shrink the budget, while
+ * `workflow_mcp_tool_server_workflow_unique` — scoped to unarchived rows —
+ * leaves their archived generations unconstrained. Without the exclusion such a
+ * server was counted twice: once against the budget and once as a candidate that
+ * consumed one of the remaining slots only to be skipped under the lock, which
+ * kept a genuinely restorable server beyond the bound from ever being fetched.
+ * The exclusion is a pre-lock optimisation and can go stale, so the
+ * `(server_id, workflow_id)` liveness check in {@link getRestoreSkipReason}
+ * stays authoritative.
+ *
  * Candidates are decided in recency order — the order the bounded query already
  * ranks them in — so when there is headroom for some but not all, the servers
  * the workflow most recently published on are the ones that come back. Lock
@@ -268,6 +281,13 @@ async function restoreArchivedMcpToolsForWorkflow(
     return
   }
 
+  /**
+   * Self-join alias correlating a candidate row with the workflow's live rows
+   * on the same server: `workflow_mcp_tool` appears on both sides of the
+   * candidate query, so the inner reference needs its own name.
+   */
+  const liveWorkflowMcpTool = alias(workflowMcpTool, 'live_workflow_mcp_tool')
+
   const latestPerServer = tx
     .selectDistinctOn([workflowMcpTool.serverId], {
       id: workflowMcpTool.id,
@@ -278,7 +298,24 @@ async function restoreArchivedMcpToolsForWorkflow(
       updatedAt: workflowMcpTool.updatedAt,
     })
     .from(workflowMcpTool)
-    .where(and(eq(workflowMcpTool.workflowId, workflowId), isNotNull(workflowMcpTool.archivedAt)))
+    .where(
+      and(
+        eq(workflowMcpTool.workflowId, workflowId),
+        isNotNull(workflowMcpTool.archivedAt),
+        notExists(
+          tx
+            .select({ id: liveWorkflowMcpTool.id })
+            .from(liveWorkflowMcpTool)
+            .where(
+              and(
+                eq(liveWorkflowMcpTool.workflowId, workflowId),
+                eq(liveWorkflowMcpTool.serverId, workflowMcpTool.serverId),
+                isNull(liveWorkflowMcpTool.archivedAt)
+              )
+            )
+        )
+      )
+    )
     .orderBy(
       asc(workflowMcpTool.serverId),
       desc(workflowMcpTool.updatedAt),