fix(mothership): keep the best-effort memo in lru-cache and let only its own attempt settle it

- Replace the hand-rolled TTL map with LRUCache (max 1000, ttl 60s), per the
  repo's caching rule.
- A timed-out attempt's late outcome now settles the memo only while the entry
  is still that attempt's: an older attempt's late success no longer erases a
  newer attempt's failure, and its late failure no longer recreates a memo
  after a newer attempt succeeded.

Author: waleedlatif1

apps/sim/lib/mothership/mcp-tools.test.ts

@@ -98,11 +98,26 @@ describe('mothership MCP best-effort discovery bounds', () => {
   const ids = (prefix: string, count: number) =>
     Array.from({ length: count }, (_, index) => `${prefix}-${index + 1}`)
 
+  let clockStart = 0
+  let clockOffset = 0
+  /**
+   * `lru-cache` captures `performance` at load, so fake timers cannot move the memo's clock: spy it
+   * forward with the faked `Date` and timers.
+   */
+  async function advanceClock(ms: number): Promise<void> {
+    clockOffset += ms
+    vi.spyOn(performance, 'now').mockImplementation(() => clockStart + clockOffset)
+    await vi.advanceTimersByTimeAsync(ms)
+  }
+
   beforeEach(() => {
-    vi.useFakeTimers()
+    vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout', 'Date'] })
+    clockStart = performance.now()
+    clockOffset = 0
     discoverServerTools.mockReset()
   })
   afterEach(() => {
+    vi.runOnlyPendingTimers()
     vi.useRealTimers()
   })
 
@@ -210,7 +225,7 @@ describe('mothership MCP best-effort discovery bounds', () => {
     await query('user-memo-b')
     expect(discoverServerTools).toHaveBeenCalledTimes(2)
 
-    vi.advanceTimersByTime(60_000)
+    await advanceClock(61_000)
     await query('user-memo-a')
     expect(discoverServerTools).toHaveBeenCalledTimes(3)
   })
@@ -237,4 +252,54 @@ describe('mothership MCP best-effort discovery bounds', () => {
     expect((await query()).tools.map((tool) => tool.service)).toEqual(['mcp:mcp-slow'])
     expect(discoverServerTools).toHaveBeenCalledTimes(2)
   })
+
+  it("keeps a newer attempt's failure when an older attempt succeeds late", async () => {
+    discoverServerTools
+      .mockImplementationOnce(
+        () => new Promise((resolve) => setTimeout(() => resolve(toolsOf('mcp-flaky')), 70_000))
+      )
+      .mockRejectedValueOnce(new Error('connect ECONNREFUSED'))
+    const query = () =>
+      buildMcpToolSchemas('user-race-a', 'ws-race', { complete: [], bestEffort: ['mcp-flaky'] })
+
+    const first = query()
+    await advanceClock(3_000)
+    expect((await first).failures.map(({ serverId }) => serverId)).toEqual(['mcp-flaky'])
+    await advanceClock(61_000)
+    expect((await query()).failures.map(({ error }) => error.message)).toEqual([
+      'connect ECONNREFUSED',
+    ])
+
+    await advanceClock(6_000)
+    const third = query()
+    await flushMicrotasks(10)
+    expect(discoverServerTools).toHaveBeenCalledTimes(2)
+    expect((await third).failures.map(({ error }) => error.message)).toEqual([
+      'connect ECONNREFUSED',
+    ])
+  })
+
+  it('does not remember a server whose older attempt fails late after a newer one succeeded', async () => {
+    discoverServerTools
+      .mockImplementationOnce(
+        () =>
+          new Promise((_resolve, reject) =>
+            setTimeout(() => reject(new Error('connect ETIMEDOUT')), 70_000)
+          )
+      )
+      .mockResolvedValueOnce(toolsOf('mcp-flaky'))
+      .mockResolvedValueOnce(toolsOf('mcp-flaky'))
+    const query = () =>
+      buildMcpToolSchemas('user-race-b', 'ws-race', { complete: [], bestEffort: ['mcp-flaky'] })
+
+    const first = query()
+    await advanceClock(3_000)
+    expect((await first).failures.map(({ serverId }) => serverId)).toEqual(['mcp-flaky'])
+    await advanceClock(61_000)
+    expect((await query()).tools.map((tool) => tool.service)).toEqual(['mcp:mcp-flaky'])
+
+    await advanceClock(6_000)
+    expect((await query()).tools.map((tool) => tool.service)).toEqual(['mcp:mcp-flaky'])
+    expect(discoverServerTools).toHaveBeenCalledTimes(3)
+  })
 })

apps/sim/lib/mothership/mcp-tools.ts

@@ -1,4 +1,5 @@
 import { toError } from '@sim/utils/errors'
+import { LRUCache } from 'lru-cache'
 import { mapSettledWithConcurrency } from '@/lib/core/utils/concurrency'
 import { DeadlineExceededError, withinDeadline } from '@/lib/core/utils/deadline'
 import { discoverMcpServerToolsAsExecutor } from '@/lib/internal/mcp/discover-tools'
@@ -33,29 +34,21 @@ const BEST_EFFORT_DISCOVERY_START_WINDOW_MS = 3_000
 const BEST_EFFORT_RETRY_AFTER_MS = 60_000
 const MAX_RECENT_BEST_EFFORT_FAILURES = 1_000
 
+/** One attempt's failure; the entry's identity is that attempt's, so a later attempt's entry is never its to settle. */
+interface BestEffortFailure {
+  error: Error
+}
+
 /**
  * Best-effort servers that recently failed for a requester, keyed by
- * `workspaceId:serverId:userId`, oldest first. It lives in this process because
- * best-effort discovery must leave the server's shared status and cooldown to
- * the discoveries someone asked for.
+ * `workspaceId:serverId:userId`. It lives in this process because best-effort
+ * discovery must leave the server's shared status and cooldown to the
+ * discoveries someone asked for.
  */
-const recentBestEffortFailures = new Map<string, { error: Error; retryAfter: number }>()
-
-function recentBestEffortFailure(key: string): Error | undefined {
-  const failure = recentBestEffortFailures.get(key)
-  if (!failure) return undefined
-  if (failure.retryAfter > Date.now()) return failure.error
-  recentBestEffortFailures.delete(key)
-  return undefined
-}
-
-function rememberBestEffortFailure(key: string, error: Error): void {
-  recentBestEffortFailures.delete(key)
-  recentBestEffortFailures.set(key, { error, retryAfter: Date.now() + BEST_EFFORT_RETRY_AFTER_MS })
-  if (recentBestEffortFailures.size <= MAX_RECENT_BEST_EFFORT_FAILURES) return
-  const oldest = recentBestEffortFailures.keys().next()
-  if (!oldest.done) recentBestEffortFailures.delete(oldest.value)
-}
+const recentBestEffortFailures = new LRUCache<string, BestEffortFailure>({
+  max: MAX_RECENT_BEST_EFFORT_FAILURES,
+  ttl: BEST_EFFORT_RETRY_AFTER_MS,
+})
 
 interface McpServerDiscoveryFailure {
   serverId: string
@@ -142,8 +135,8 @@ export async function buildMcpToolSchemas(
   /** Ignores the deadline's signal so a slow server's discovery still finishes and warms the cache. */
   const discoverBestEffort = async (serverId: string) => {
     const failureKey = `${workspaceId}:${serverId}:${userId}`
-    const recentFailure = recentBestEffortFailure(failureKey)
-    if (recentFailure) throw recentFailure
+    const recentFailure = recentBestEffortFailures.get(failureKey)
+    if (recentFailure) throw recentFailure.error
     if (Date.now() >= startWindowEndsAt) throw new DeadlineExceededError()
     const discovery = discoverMcpServerToolsUseCase.execute({
       principal: principal(),
@@ -156,12 +149,20 @@ export async function buildMcpToolSchemas(
       )
       return tools
     } catch (error) {
-      rememberBestEffortFailure(failureKey, toError(error))
-      if (error instanceof DeadlineExceededError)
+      const failure: BestEffortFailure = { error: toError(error) }
+      recentBestEffortFailures.set(failureKey, failure)
+      if (error instanceof DeadlineExceededError) {
+        const isLatest = () =>
+          recentBestEffortFailures.peek(failureKey, { allowStale: true }) === failure
         discovery.then(
-          () => recentBestEffortFailures.delete(failureKey),
-          (lateError: unknown) => rememberBestEffortFailure(failureKey, toError(lateError))
+          () => {
+            if (isLatest()) recentBestEffortFailures.delete(failureKey)
+          },
+          (lateError: unknown) => {
+            if (isLatest()) recentBestEffortFailures.set(failureKey, { error: toError(lateError) })
+          }
         )
+      }
       throw error
     }
   }