Skip to content

Commit 4966fe7

Browse files
committed
fix(mothership): bound best-effort mcp discovery and stop retrying a down server every query
Broad chat catalog discovery started best-effort discovery for every usable workspace server at once. The 3s wait only ends the caller's wait, so each attempt kept its sockets open for up to ~2 min, and a workspace with many servers could burst past the connection pool and evict other users' warm connections. Best-effort failures write no shared cooldown, so a down server was reattempted (with its DB and Redis reads) on every query. - Run best-effort discovery through mapSettledWithConcurrency with 8 at once; each server's 3s wait counts from its own start, and servers still queued 3s into the query are skipped unstarted, bounding query latency at ~6s. - Remember, per process and per requester (workspace:server:user, 60s, capped at 1000 entries), a best-effort server that failed or outlived its wait, and skip it until the memo expires or the late discovery succeeds. Nothing is written to shared state. - Named servers are unchanged: listed at once, waited on to completion.
1 parent fd4323a commit 4966fe7

2 files changed

Lines changed: 234 additions & 15 deletions

File tree

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

Lines changed: 158 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1+
import { flushMicrotasks } from '@sim/testing'
12
import { mcpUseCasesMock, mcpUseCasesMockFns } from '@sim/testing/mocks/mcp-use-cases.mock'
2-
import { describe, expect, it, vi } from 'vitest'
3+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
34

45
const { discoverServerTools } = vi.hoisted(() => ({
56
discoverServerTools: vi.fn(),
@@ -10,6 +11,7 @@ vi.mock('@/lib/internal/mcp/discover-tools', () => ({
1011
}))
1112
vi.mock('@/lib/mcp/application/use-cases', () => mcpUseCasesMock)
1213

14+
import { DeadlineExceededError } from '@/lib/core/utils/deadline'
1315
import { buildMcpToolSchemas } from '@/lib/mothership/mcp-tools'
1416

1517
mcpUseCasesMockFns.mockDiscoverMcpServerToolsUseCase.mockImplementation(async (args: unknown) => ({
@@ -81,3 +83,158 @@ describe('mothership MCP tool schemas', () => {
8183
expect(discovery.failures).toEqual([])
8284
})
8385
})
86+
87+
/**
88+
* Broad chat discovery reaches every usable workspace server. Its failure modes are fan-out
89+
* (every server's connect started at once), latency that grows with the server count, a slow
90+
* server cut short by a wait that began before it started, and a down server reattempted on
91+
* every query. None of it may touch the servers the request names.
92+
*/
93+
describe('mothership MCP best-effort discovery bounds', () => {
94+
const hanging = () => new Promise<never>(() => undefined)
95+
const toolsOf = (serverId: string) => [{ serverId, name: 'search', inputSchema: {} }]
96+
const serverIdOf = (args: unknown) => (args as { input: { serverId: string } }).input.serverId
97+
const startedServers = () => discoverServerTools.mock.calls.map(([args]) => serverIdOf(args))
98+
const ids = (prefix: string, count: number) =>
99+
Array.from({ length: count }, (_, index) => `${prefix}-${index + 1}`)
100+
101+
beforeEach(() => {
102+
vi.useFakeTimers()
103+
discoverServerTools.mockReset()
104+
})
105+
afterEach(() => {
106+
vi.useRealTimers()
107+
})
108+
109+
it('starts at most eight best-effort discoveries at once and skips the rest once the start window closes', async () => {
110+
discoverServerTools.mockImplementation(hanging)
111+
const servers = ids('mcp-hang', 20)
112+
113+
const pending = buildMcpToolSchemas('user-cap', 'ws-cap', { complete: [], bestEffort: servers })
114+
await flushMicrotasks(10)
115+
expect(startedServers()).toEqual(servers.slice(0, 8))
116+
117+
await vi.advanceTimersByTimeAsync(3_000)
118+
const discovery = await pending
119+
expect(startedServers()).toEqual(servers.slice(0, 8))
120+
expect(discovery.failures.map(({ serverId }) => serverId)).toEqual(servers)
121+
expect(discovery.failures.every(({ error }) => error instanceof DeadlineExceededError)).toBe(
122+
true
123+
)
124+
})
125+
126+
it('keeps discovering fast servers while slow ones hold their slots', async () => {
127+
const fast = ids('mcp-fast', 10)
128+
discoverServerTools.mockImplementation(async (args: unknown) => {
129+
const serverId = serverIdOf(args)
130+
return serverId === 'mcp-slow' ? hanging() : toolsOf(serverId)
131+
})
132+
133+
const pending = buildMcpToolSchemas('user-fast', 'ws-fast', {
134+
complete: [],
135+
bestEffort: ['mcp-slow', ...fast],
136+
})
137+
await flushMicrotasks(50)
138+
expect(startedServers()).toEqual(['mcp-slow', ...fast])
139+
140+
await vi.advanceTimersByTimeAsync(3_000)
141+
const discovery = await pending
142+
expect(discovery.tools.map((tool) => tool.service)).toEqual(
143+
fast.map((serverId) => `mcp:${serverId}`)
144+
)
145+
expect(discovery.failures.map(({ serverId }) => serverId)).toEqual(['mcp-slow'])
146+
})
147+
148+
it("measures each server's wait from its own start, not the query's", async () => {
149+
const blockers = ids('mcp-block', 8)
150+
discoverServerTools.mockImplementation((args: unknown) => {
151+
const serverId = serverIdOf(args)
152+
if (serverId === 'mcp-block-1')
153+
return new Promise((resolve) => setTimeout(() => resolve(toolsOf(serverId)), 1_000))
154+
if (serverId === 'mcp-late') {
155+
const resolveAt = 3_500 - (Date.now() - startedAt)
156+
return new Promise((resolve) => setTimeout(() => resolve(toolsOf(serverId)), resolveAt))
157+
}
158+
return hanging()
159+
})
160+
const startedAt = Date.now()
161+
162+
const pending = buildMcpToolSchemas('user-own', 'ws-own', {
163+
complete: [],
164+
bestEffort: [...blockers, 'mcp-late'],
165+
})
166+
await vi.advanceTimersByTimeAsync(4_000)
167+
const discovery = await pending
168+
169+
expect(discovery.tools.map((tool) => tool.service)).toEqual(['mcp:mcp-block-1', 'mcp:mcp-late'])
170+
expect(discovery.failures.map(({ serverId }) => serverId)).toEqual(blockers.slice(1))
171+
})
172+
173+
it('lists every named server at once and waits for each however long it takes', async () => {
174+
const named = ids('mcp-named', 12)
175+
discoverServerTools.mockImplementation((args: unknown) => {
176+
const serverId = serverIdOf(args)
177+
return serverId.startsWith('mcp-named')
178+
? new Promise((resolve) => setTimeout(() => resolve(toolsOf(serverId)), 30_000))
179+
: hanging()
180+
})
181+
182+
const pending = buildMcpToolSchemas('user-named', 'ws-named', {
183+
complete: named,
184+
bestEffort: ids('mcp-optional', 10),
185+
})
186+
await flushMicrotasks(10)
187+
expect(startedServers()).toEqual([...named, ...ids('mcp-optional', 8)])
188+
189+
await vi.advanceTimersByTimeAsync(30_000)
190+
const discovery = await pending
191+
expect(discovery.tools.map((tool) => tool.service)).toEqual(
192+
named.map((serverId) => `mcp:${serverId}`)
193+
)
194+
expect(discovery.failures.map(({ serverId }) => serverId)).toEqual(ids('mcp-optional', 10))
195+
})
196+
197+
it('does not reattempt a failed best-effort server for the same requester until the memo expires', async () => {
198+
discoverServerTools.mockRejectedValue(new Error('connect ECONNREFUSED'))
199+
const query = (userId: string) =>
200+
buildMcpToolSchemas(userId, 'ws-memo', { complete: [], bestEffort: ['mcp-down'] })
201+
202+
expect((await query('user-memo-a')).failures.map(({ error }) => error.message)).toEqual([
203+
'connect ECONNREFUSED',
204+
])
205+
expect((await query('user-memo-a')).failures.map(({ error }) => error.message)).toEqual([
206+
'connect ECONNREFUSED',
207+
])
208+
expect(discoverServerTools).toHaveBeenCalledTimes(1)
209+
210+
await query('user-memo-b')
211+
expect(discoverServerTools).toHaveBeenCalledTimes(2)
212+
213+
vi.advanceTimersByTime(60_000)
214+
await query('user-memo-a')
215+
expect(discoverServerTools).toHaveBeenCalledTimes(3)
216+
})
217+
218+
it('skips a server whose discovery outlived the wait until that discovery settles, then serves it', async () => {
219+
discoverServerTools.mockImplementation(
220+
(args: unknown) =>
221+
new Promise((resolve) => setTimeout(() => resolve(toolsOf(serverIdOf(args))), 5_000))
222+
)
223+
const query = () =>
224+
buildMcpToolSchemas('user-slow', 'ws-slow', { complete: [], bestEffort: ['mcp-slow'] })
225+
226+
const first = query()
227+
await vi.advanceTimersByTimeAsync(3_000)
228+
expect((await first).failures.map(({ serverId }) => serverId)).toEqual(['mcp-slow'])
229+
230+
const second = query()
231+
await flushMicrotasks(10)
232+
expect(discoverServerTools).toHaveBeenCalledTimes(1)
233+
expect((await second).failures.map(({ serverId }) => serverId)).toEqual(['mcp-slow'])
234+
235+
await vi.advanceTimersByTimeAsync(2_000)
236+
discoverServerTools.mockResolvedValueOnce(toolsOf('mcp-slow'))
237+
expect((await query()).tools.map((tool) => tool.service)).toEqual(['mcp:mcp-slow'])
238+
expect(discoverServerTools).toHaveBeenCalledTimes(2)
239+
})
240+
})

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

Lines changed: 76 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { toError } from '@sim/utils/errors'
2-
import { withinDeadline } from '@/lib/core/utils/deadline'
2+
import { mapSettledWithConcurrency } from '@/lib/core/utils/concurrency'
3+
import { DeadlineExceededError, withinDeadline } from '@/lib/core/utils/deadline'
34
import { discoverMcpServerToolsAsExecutor } from '@/lib/internal/mcp/discover-tools'
45
import type { InternalToolOperationContext } from '@/lib/internal/tool-operations/types'
56
import { MCP_SERVER_DELEGATION_AUDIENCE } from '@/lib/mcp/application/authorization'
@@ -10,11 +11,51 @@ import { createCopilotChatPrincipal } from '@/lib/mothership/auth/application-de
1011
import type { ToolSchema } from '@/lib/mothership/chat/payload'
1112

1213
/**
13-
* How long one best-effort server may hold up a catalog query. A warm server is
14-
* a cache read; a cold or slow one is skipped for this query while its
15-
* discovery keeps running and fills the cache for the next one.
14+
* How long a catalog query waits on one best-effort server, counted from that
15+
* server's own start. A warm server is a cache read; a cold or slow one is
16+
* skipped for this query while its discovery keeps running and fills the cache
17+
* for the next one.
1618
*/
1719
const BEST_EFFORT_DISCOVERY_WAIT_MS = 3_000
20+
/**
21+
* Best-effort discoveries a query runs at once. The wait does not cancel a
22+
* discovery, so this also caps the connections one query can leave running.
23+
*/
24+
const BEST_EFFORT_DISCOVERY_CONCURRENCY = 8
25+
/**
26+
* How long after a query begins it may still start a best-effort discovery.
27+
* Servers still queued then are skipped for this query, so with the wait it
28+
* bounds the query's latency at about six seconds however many servers the
29+
* workspace has.
30+
*/
31+
const BEST_EFFORT_DISCOVERY_START_WINDOW_MS = 3_000
32+
/** How long a requester skips a best-effort server whose discovery failed or outlived the wait. */
33+
const BEST_EFFORT_RETRY_AFTER_MS = 60_000
34+
const MAX_RECENT_BEST_EFFORT_FAILURES = 1_000
35+
36+
/**
37+
* Best-effort servers that recently failed for a requester, keyed by
38+
* `workspaceId:serverId:userId`, oldest first. It lives in this process because
39+
* best-effort discovery must leave the server's shared status and cooldown to
40+
* the discoveries someone asked for.
41+
*/
42+
const recentBestEffortFailures = new Map<string, { error: Error; retryAfter: number }>()
43+
44+
function recentBestEffortFailure(key: string): Error | undefined {
45+
const failure = recentBestEffortFailures.get(key)
46+
if (!failure) return undefined
47+
if (failure.retryAfter > Date.now()) return failure.error
48+
recentBestEffortFailures.delete(key)
49+
return undefined
50+
}
51+
52+
function rememberBestEffortFailure(key: string, error: Error): void {
53+
recentBestEffortFailures.delete(key)
54+
recentBestEffortFailures.set(key, { error, retryAfter: Date.now() + BEST_EFFORT_RETRY_AFTER_MS })
55+
if (recentBestEffortFailures.size <= MAX_RECENT_BEST_EFFORT_FAILURES) return
56+
const oldest = recentBestEffortFailures.keys().next()
57+
if (!oldest.done) recentBestEffortFailures.delete(oldest.value)
58+
}
1859

1960
interface McpServerDiscoveryFailure {
2061
serverId: string
@@ -31,8 +72,13 @@ interface McpDiscoveryServers {
3172
/** Servers the request names: listed live and completely, however long that takes. */
3273
complete: string[]
3374
/**
34-
* Chat-only: servers available without being named. Served from the discovery
35-
* cache when fresh, otherwise skipped after {@link BEST_EFFORT_DISCOVERY_WAIT_MS}.
75+
* Chat-only: servers available without being named. At most
76+
* {@link BEST_EFFORT_DISCOVERY_CONCURRENCY} are discovered at once, each served
77+
* from the discovery cache when fresh and otherwise skipped after
78+
* {@link BEST_EFFORT_DISCOVERY_WAIT_MS}; one still queued after
79+
* {@link BEST_EFFORT_DISCOVERY_START_WINDOW_MS} is skipped unstarted. One that
80+
* failed or outlived its wait is skipped for the same requester for
81+
* {@link BEST_EFFORT_RETRY_AFTER_MS}, or until that late discovery succeeds.
3682
* One with no permitted operations contributes nothing and is not a failure.
3783
* Their discovery never touches a server's shared status or cooldown, and skips a
3884
* server the requester cannot authenticate to.
@@ -92,23 +138,39 @@ export async function buildMcpToolSchemas(
92138
if (!tools.length) throw new Error(`No permitted MCP operations are available for ${serverId}`)
93139
return tools
94140
}
141+
const startWindowEndsAt = Date.now() + BEST_EFFORT_DISCOVERY_START_WINDOW_MS
95142
/** Ignores the deadline's signal so a slow server's discovery still finishes and warms the cache. */
96143
const discoverBestEffort = async (serverId: string) => {
144+
const failureKey = `${workspaceId}:${serverId}:${userId}`
145+
const recentFailure = recentBestEffortFailure(failureKey)
146+
if (recentFailure) throw recentFailure
147+
if (Date.now() >= startWindowEndsAt) throw new DeadlineExceededError()
97148
const discovery = discoverMcpServerToolsUseCase.execute({
98149
principal: principal(),
99150
input: { workspaceId, serverId, bestEffort: true },
100151
})
101-
const { tools } = await withinDeadline(
102-
() => discovery,
103-
Date.now() + BEST_EFFORT_DISCOVERY_WAIT_MS
104-
)
105-
return tools
152+
try {
153+
const { tools } = await withinDeadline(
154+
() => discovery,
155+
Date.now() + BEST_EFFORT_DISCOVERY_WAIT_MS
156+
)
157+
return tools
158+
} catch (error) {
159+
rememberBestEffortFailure(failureKey, toError(error))
160+
if (error instanceof DeadlineExceededError)
161+
discovery.then(
162+
() => recentBestEffortFailures.delete(failureKey),
163+
(lateError: unknown) => rememberBestEffortFailure(failureKey, toError(lateError))
164+
)
165+
throw error
166+
}
106167
}
107168

108-
const outcomes = await Promise.allSettled([
109-
...complete.map(discoverComplete),
110-
...bestEffort.map(discoverBestEffort),
169+
const [completeOutcomes, bestEffortOutcomes] = await Promise.all([
170+
Promise.allSettled(complete.map(discoverComplete)),
171+
mapSettledWithConcurrency(bestEffort, BEST_EFFORT_DISCOVERY_CONCURRENCY, discoverBestEffort),
111172
])
173+
const outcomes = [...completeOutcomes, ...bestEffortOutcomes]
112174

113175
const seen = new Set<string>()
114176
const discovery: McpDiscovery = { tools: [], failures: [] }

0 commit comments

Comments
 (0)