From 3b3857c500971222d6b1c59ba5a0e8fa858da1b4 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 22 Jul 2026 04:33:51 +0000 Subject: [PATCH] Bound remote connector snapshot reads --- .../mcp/capabilities/registry.node.test.ts | 90 +++++++++++++++++++ .../worker/src/mcp/capabilities/registry.ts | 36 ++++++-- .../snapshot-cache.node.test.ts | 80 ++++++++++++++++- .../src/remote-connector/snapshot-cache.ts | 38 +++++++- 4 files changed, 233 insertions(+), 11 deletions(-) diff --git a/packages/worker/src/mcp/capabilities/registry.node.test.ts b/packages/worker/src/mcp/capabilities/registry.node.test.ts index 796fb8a620..cc80b5622b 100644 --- a/packages/worker/src/mcp/capabilities/registry.node.test.ts +++ b/packages/worker/src/mcp/capabilities/registry.node.test.ts @@ -5,6 +5,10 @@ import { getCapabilityRegistryForContext, getStaticRegistry, } from '#mcp/capabilities/registry.ts' +import { + clearRemoteConnectorSnapshotCacheForTests, + remoteConnectorSnapshotTimeoutMs, +} from '#worker/remote-connector/snapshot-cache.ts' test('getCapabilityRegistryForContext bypasses connector caches when no connectors are attached', async () => { clearCapabilityRegistryCacheForTests() @@ -128,3 +132,89 @@ test('the email domain no longer exposes self-service inbox or sender-identity c expect(capabilityMap.email_message_list).toBeTruthy() expect(capabilityMap.email_message_get).toBeTruthy() }) + +test('getCapabilityRegistryForContext keeps healthy capabilities when a connector snapshot stalls', async () => { + vi.useFakeTimers() + clearCapabilityRegistryCacheForTests() + clearRemoteConnectorSnapshotCacheForTests() + const prepare = vi.fn(() => ({ + bind() { + return this + }, + async all() { + return { results: [], meta: { changes: 0 } } + }, + async first() { + return null + }, + async run() { + return { meta: { changes: 0 } } + }, + })) + const env = { + REMOTE_CONNECTOR_SESSION: { + idFromName(name: string) { + return name + }, + get(name: string) { + const [, instanceId] = JSON.parse(name) as [string, string] + return { + getSnapshot() { + if (instanceId === 'stalled') { + return new Promise(() => undefined) + } + return Promise.resolve({ + connectorId: instanceId, + connectedAt: '2026-03-25T00:00:00.000Z', + lastSeenAt: '2026-03-25T00:00:01.000Z', + tools: [ + { + name: 'roku_press_key', + description: 'Send a Roku ECP keypress.', + inputSchema: { + type: 'object', + properties: { + key: { type: 'string' }, + }, + required: ['key'], + }, + }, + ], + }) + }, + } + }, + }, + APP_DB: { + prepare, + }, + } as unknown as Env + const callerContext = createMcpCallerContext({ + baseUrl: 'https://heykody.dev', + user: { + userId: 'user-1', + email: 'user-1@example.com', + displayName: 'user-1', + }, + remoteConnectors: [{ instanceId: 'healthy' }, { instanceId: 'stalled' }], + }) + + try { + const registryPromise = getCapabilityRegistryForContext({ + env, + callerContext, + }) + await vi.advanceTimersByTimeAsync(remoteConnectorSnapshotTimeoutMs) + const registry = await registryPromise + + expect(registry.capabilityMap.email_send).toBeTruthy() + expect(registry.capabilityMap['remote:healthy:roku_press_key']).toBeTruthy() + expect( + registry.capabilityMap['remote:stalled:roku_press_key'], + ).toBeUndefined() + } finally { + clearCapabilityRegistryCacheForTests() + clearRemoteConnectorSnapshotCacheForTests() + vi.useRealTimers() + } +}) diff --git a/packages/worker/src/mcp/capabilities/registry.ts b/packages/worker/src/mcp/capabilities/registry.ts index d0237148b5..1c62e4dfba 100644 --- a/packages/worker/src/mcp/capabilities/registry.ts +++ b/packages/worker/src/mcp/capabilities/registry.ts @@ -201,6 +201,28 @@ async function loadMcpServerSnapshots(input: { } } +async function loadRemoteConnectorSnapshots(input: { + env: Env + userId: string + refs: ReadonlyArray +}): Promise> { + return await Promise.all( + input.refs.map(async (ref) => { + try { + return await getCachedRemoteConnectorSnapshot({ + env: input.env, + userId: input.userId, + instanceId: ref.instanceId, + }) + } catch { + // A stalled or unavailable connector must not prevent builtin or + // healthy dynamic capabilities from being discovered. + return null + } + }), + ) +} + async function loadOpenApiBindingsForRegistry(input: { env: Env userId: string @@ -262,15 +284,11 @@ export async function getCapabilityRegistryForContext(input: { }) } const [snapshots, mcpServerSnapshots] = await Promise.all([ - Promise.all( - refs.map((ref) => - getCachedRemoteConnectorSnapshot({ - env: input.env, - userId, - instanceId: ref.instanceId, - }), - ), - ), + loadRemoteConnectorSnapshots({ + env: input.env, + userId, + refs, + }), loadMcpServerSnapshots({ env: input.env, userId, diff --git a/packages/worker/src/remote-connector/snapshot-cache.node.test.ts b/packages/worker/src/remote-connector/snapshot-cache.node.test.ts index e59632360f..df95311944 100644 --- a/packages/worker/src/remote-connector/snapshot-cache.node.test.ts +++ b/packages/worker/src/remote-connector/snapshot-cache.node.test.ts @@ -1,8 +1,11 @@ -import { expect, test } from 'vitest' +import { expect, test, vi } from 'vitest' import { createRemoteConnectorMcpClient } from '#worker/remote-connector/client.ts' +import { getRemoteConnectorStatus } from '#worker/remote-connector/status.ts' import { clearRemoteConnectorSnapshotCacheForTests, getCachedRemoteConnectorSnapshot, + RemoteConnectorSnapshotTimeoutError, + remoteConnectorSnapshotTimeoutMs, } from '#worker/remote-connector/snapshot-cache.ts' const runtimeTools = [ @@ -160,3 +163,78 @@ test('getCachedRemoteConnectorSnapshot does not share entries across users', asy expect(getSnapshotCalls()).toBe(2) }) + +test('a stalled snapshot times out, rejects concurrent callers, and is retried', async () => { + vi.useFakeTimers() + clearRemoteConnectorSnapshotCacheForTests() + let attempt = 0 + const { env, getSnapshotCalls } = buildEnv(() => { + attempt += 1 + if (attempt === 1) { + return new Promise(() => undefined) + } + return Promise.resolve({ + connectorKind: 'roku', + connectorId: 'home', + connectedAt: '2026-03-25T00:00:00.000Z', + lastSeenAt: '2026-03-25T00:00:01.000Z', + tools: runtimeTools, + }) + }) + const request = { + env, + userId: 'user-1', + instanceId: 'home', + } + + try { + const first = getCachedRemoteConnectorSnapshot(request) + const second = getCachedRemoteConnectorSnapshot(request) + const firstRejection = expect(first).rejects.toBeInstanceOf( + RemoteConnectorSnapshotTimeoutError, + ) + const secondRejection = expect(second).rejects.toBeInstanceOf( + RemoteConnectorSnapshotTimeoutError, + ) + + expect(getSnapshotCalls()).toBe(1) + await vi.advanceTimersByTimeAsync(remoteConnectorSnapshotTimeoutMs) + await Promise.all([firstRejection, secondRejection]) + + await expect( + getCachedRemoteConnectorSnapshot(request), + ).resolves.toMatchObject({ + connectorId: 'home', + }) + expect(getSnapshotCalls()).toBe(2) + } finally { + clearRemoteConnectorSnapshotCacheForTests() + vi.useRealTimers() + } +}) + +test('remote connector status reports a stalled snapshot as an error', async () => { + vi.useFakeTimers() + clearRemoteConnectorSnapshotCacheForTests() + const { env } = buildEnv(() => new Promise(() => undefined)) + + try { + const statusPromise = getRemoteConnectorStatus({ + env, + userId: 'user-1', + ref: { instanceId: 'home' }, + }) + await vi.advanceTimersByTimeAsync(remoteConnectorSnapshotTimeoutMs) + + await expect(statusPromise).resolves.toMatchObject({ + state: 'error', + connectorId: 'home', + connected: false, + toolCount: 0, + error: expect.stringContaining('snapshot timed out'), + }) + } finally { + clearRemoteConnectorSnapshotCacheForTests() + vi.useRealTimers() + } +}) diff --git a/packages/worker/src/remote-connector/snapshot-cache.ts b/packages/worker/src/remote-connector/snapshot-cache.ts index 8f41415be3..e93be70160 100644 --- a/packages/worker/src/remote-connector/snapshot-cache.ts +++ b/packages/worker/src/remote-connector/snapshot-cache.ts @@ -4,6 +4,39 @@ import { type RemoteConnectorSnapshot } from '#worker/remote-connector/types.ts' export const remoteConnectorSnapshotCacheTtlMs = 30_000 export const remoteConnectorSnapshotCacheLimit = 100 +export const remoteConnectorSnapshotTimeoutMs = 5_000 + +export class RemoteConnectorSnapshotTimeoutError extends Error { + constructor(instanceId: string) { + super( + `Remote connector "${instanceId}" snapshot timed out after ${remoteConnectorSnapshotTimeoutMs}ms.`, + ) + this.name = 'RemoteConnectorSnapshotTimeoutError' + } +} + +async function getSnapshotWithTimeout(input: { + instanceId: string + getSnapshot: () => Promise +}) { + let timeoutId: ReturnType | undefined + try { + return await Promise.race([ + input.getSnapshot(), + new Promise((_resolve, reject) => { + timeoutId = setTimeout( + () => + reject(new RemoteConnectorSnapshotTimeoutError(input.instanceId)), + remoteConnectorSnapshotTimeoutMs, + ) + }), + ]) + } finally { + if (timeoutId !== undefined) { + clearTimeout(timeoutId) + } + } +} function createRemoteConnectorSnapshotCache() { return new PromiseLruCache({ @@ -33,7 +66,10 @@ export function getCachedRemoteConnectorSnapshot(input: { const stub = input.env.REMOTE_CONNECTOR_SESSION.get( input.env.REMOTE_CONNECTOR_SESSION.idFromName(cacheKey), ) - const snapshot = await stub.getSnapshot() + const snapshot = await getSnapshotWithTimeout({ + instanceId: input.instanceId, + getSnapshot: () => stub.getSnapshot(), + }) if (snapshot == null) { // Do not retain disconnected results: a connector that comes online // should be visible on the next lookup instead of after the TTL.