diff --git a/apps/sim/lib/mcp/application/use-cases.ts b/apps/sim/lib/mcp/application/use-cases.ts index 68732ed1ecf..fa85d70fb51 100644 --- a/apps/sim/lib/mcp/application/use-cases.ts +++ b/apps/sim/lib/mcp/application/use-cases.ts @@ -143,6 +143,8 @@ export interface DiscoverMcpServerToolsInput { refresh?: boolean signal?: AbortSignal requireComplete?: boolean + /** Discovery nobody asked for: reads the server's shared health state but never writes it. */ + bestEffort?: boolean } /** @@ -192,14 +194,18 @@ export const discoverMcpServerToolsUseCase = defineAuthorizedWorkspaceUseCase({ }) const refresh = input.refresh ? 'skip-cache' : 'cache-aside' const tools = - input.signal || input.requireComplete + input.signal || input.requireComplete || input.bestEffort ? await mcpService.discoverServerTools( userId, context.server.id, context.workspaceId, refresh, undefined, - { signal: input.signal, requireComplete: input.requireComplete } + { + signal: input.signal, + requireComplete: input.requireComplete, + bestEffort: input.bestEffort, + } ) : await mcpService.discoverServerTools( userId, diff --git a/apps/sim/lib/mcp/oauth/index.ts b/apps/sim/lib/mcp/oauth/index.ts index 7e89f72d7c3..a7b1e5d5b97 100644 --- a/apps/sim/lib/mcp/oauth/index.ts +++ b/apps/sim/lib/mcp/oauth/index.ts @@ -18,6 +18,7 @@ export { clearTokens, clearVerifier, getOrCreateOauthRow, + hasDecryptableOauthTokens, loadOauthRow, loadOauthRowByState, type McpOauthRow, diff --git a/apps/sim/lib/mcp/oauth/provider.test.ts b/apps/sim/lib/mcp/oauth/provider.test.ts new file mode 100644 index 00000000000..a24c3d06268 --- /dev/null +++ b/apps/sim/lib/mcp/oauth/provider.test.ts @@ -0,0 +1,94 @@ +import { auth } from '@modelcontextprotocol/sdk/client/auth.js' +import { dbChainMockFns, encryptionMock, resetDbChainMock } from '@sim/testing' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('@/lib/core/security/encryption', () => encryptionMock) + +import { SimMcpOauthProvider } from '@/lib/mcp/oauth/provider' +import type { McpOauthRow } from '@/lib/mcp/oauth/storage' + +const SERVER_URL = 'https://mcp.example.com/mcp' +const ISSUER = 'https://auth.example.com' + +function grantRow(): McpOauthRow { + return { + id: 'oauth-row-1', + mcpServerId: 'server-1', + userId: 'authorizer-1', + workspaceId: 'workspace-1', + organizationId: null, + clientInformation: { client_id: 'client-1', issuer: ISSUER }, + tokens: { access_token: 'expired', token_type: 'bearer', refresh_token: 'refresh-1' }, + codeVerifier: null, + state: null, + stateCreatedAt: null, + updatedAt: new Date(), + } +} + +/** An authorization server whose token endpoint answers every refresh with `token`. */ +function authorizationServer(token: { status: number; body: Record }) { + return vi.fn(async (input: string | URL) => { + const url = String(input) + const json = (status: number, body: unknown) => + new Response(JSON.stringify(body), { + status, + headers: { 'content-type': 'application/json' }, + }) + if (url.startsWith('https://mcp.example.com/.well-known/oauth-protected-resource')) + return json(200, { resource: SERVER_URL, authorization_servers: [ISSUER] }) + if (url.startsWith(`${ISSUER}/.well-known/`)) + return json(200, { + issuer: ISSUER, + authorization_endpoint: `${ISSUER}/authorize`, + token_endpoint: `${ISSUER}/token`, + response_types_supported: ['code'], + code_challenge_methods_supported: ['S256'], + }) + if (url === `${ISSUER}/token`) return json(token.status, token.body) + return json(404, {}) + }) +} + +/** + * Best-effort discovery connects with the workspace's one shared grant. A refresh that fails + * there must leave the grant as it was: the SDK otherwise clears the tokens and starts a new + * authorization flow, disconnecting the server for every member. + */ +describe('SimMcpOauthProvider preserving the shared grant', () => { + beforeEach(() => { + resetDbChainMock() + }) + + it('writes nothing to the OAuth row when a refresh is rejected', async () => { + const provider = new SimMcpOauthProvider({ row: grantRow(), preserveGrant: true }) + + await expect( + auth(provider, { + serverUrl: SERVER_URL, + fetchFn: authorizationServer({ status: 400, body: { error: 'invalid_grant' } }), + }) + ).rejects.toThrow('MCP OAuth grant left unchanged') + + expect(dbChainMockFns.update).not.toHaveBeenCalled() + expect(dbChainMockFns.insert).not.toHaveBeenCalled() + expect(provider.tokens()?.refresh_token).toBe('refresh-1') + }) + + it('persists a successful refresh, which keeps the grant valid for everyone', async () => { + const provider = new SimMcpOauthProvider({ row: grantRow(), preserveGrant: true }) + + await expect( + auth(provider, { + serverUrl: SERVER_URL, + fetchFn: authorizationServer({ + status: 200, + body: { access_token: 'fresh', token_type: 'bearer', refresh_token: 'refresh-2' }, + }), + }) + ).resolves.toBe('AUTHORIZED') + + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(provider.tokens()?.access_token).toBe('fresh') + }) +}) diff --git a/apps/sim/lib/mcp/oauth/provider.ts b/apps/sim/lib/mcp/oauth/provider.ts index 5a7d93bde9a..ece15d37b29 100644 --- a/apps/sim/lib/mcp/oauth/provider.ts +++ b/apps/sim/lib/mcp/oauth/provider.ts @@ -71,17 +71,37 @@ interface SimMcpOauthProviderInit { * Dynamic Client Registration and uses these for the auth/token exchange. */ preregistered?: PreregisteredClient + /** + * The grant belongs to every member and this caller did not ask to reauthorize it: a + * refreshed token still persists, but clearing the grant, registering a client or starting + * an authorization flow fails with {@link McpOauthGrantPreservedError} before any write. + */ + preserveGrant?: boolean +} + +/** A refresh failed where reauthorizing the shared grant is left to an explicit reconnect. */ +class McpOauthGrantPreservedError extends Error { + constructor(step: string) { + super(`MCP OAuth grant left unchanged: ${step} requires an explicit reconnect`) + this.name = 'McpOauthGrantPreservedError' + } } export class SimMcpOauthProvider implements OAuthClientProvider { private row: McpOauthRow private readonly scope?: string private readonly preregistered?: PreregisteredClient + private readonly preserveGrant: boolean - constructor({ row, scope, preregistered }: SimMcpOauthProviderInit) { + constructor({ row, scope, preregistered, preserveGrant = false }: SimMcpOauthProviderInit) { this.row = row this.scope = preregistered?.scope ?? scope this.preregistered = preregistered + this.preserveGrant = preserveGrant + } + + private assertMayChangeGrant(step: string): void { + if (this.preserveGrant) throw new McpOauthGrantPreservedError(step) } /** Deployment registrations may restrict consent even when discovery advertises more tools. */ @@ -108,6 +128,7 @@ export class SimMcpOauthProvider implements OAuthClientProvider { } async state(): Promise { + this.assertMayChangeGrant('authorization') const state = generateId() await saveState(this.row.id, state, 'provider.state') return state @@ -132,6 +153,7 @@ export class SimMcpOauthProvider implements OAuthClientProvider { async saveClientInformation(info: OAuthClientInformationMixed): Promise { if (this.preregistered) return + this.assertMayChangeGrant('client registration') await saveClientInformationDb(this.row.id, info) this.row.clientInformation = info } @@ -150,6 +172,7 @@ export class SimMcpOauthProvider implements OAuthClientProvider { } async saveCodeVerifier(codeVerifier: string): Promise { + this.assertMayChangeGrant('authorization') await saveCodeVerifierDb(this.row.id, codeVerifier) this.row.codeVerifier = codeVerifier } @@ -164,6 +187,7 @@ export class SimMcpOauthProvider implements OAuthClientProvider { async invalidateCredentials( scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery' ): Promise { + if (scope !== 'discovery') this.assertMayChangeGrant(`invalidating ${scope} credentials`) if (scope === 'all' || scope === 'client') { await clearClient(this.row.id) this.row.clientInformation = null diff --git a/apps/sim/lib/mcp/oauth/storage.test.ts b/apps/sim/lib/mcp/oauth/storage.test.ts index 7202b1d1c33..dc7dcdcf3d1 100644 --- a/apps/sim/lib/mcp/oauth/storage.test.ts +++ b/apps/sim/lib/mcp/oauth/storage.test.ts @@ -14,7 +14,12 @@ afterAll(resetRedisConfigMock) vi.mock('@/lib/core/security/encryption', () => encryptionMock) -import { getOrCreateOauthRow, loadOauthRow, withMcpOauthRefreshLock } from './storage' +import { + getOrCreateOauthRow, + hasDecryptableOauthTokens, + loadOauthRow, + withMcpOauthRefreshLock, +} from './storage' describe('MCP OAuth storage', () => { beforeEach(() => { @@ -77,6 +82,21 @@ describe('MCP OAuth storage', () => { expect(row.userId).toBe('authorizer-1') expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) + + /** The row is the workspace's one grant; a presence check must never be what erases it. */ + it('reports undecryptable tokens as unusable without clearing the shared row', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ tokens: 'ciphertext-from-a-rotated-key' }]) + encryptionMockFns.mockDecryptSecret.mockRejectedValueOnce(new Error('bad decrypt')) + + expect(await hasDecryptableOauthTokens('server-1')).toBe(false) + expect(dbChainMockFns.update).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() + + dbChainMockFns.limit.mockResolvedValueOnce([{ tokens: 'ciphertext' }]) + expect(await hasDecryptableOauthTokens('server-1')).toBe(true) + dbChainMockFns.limit.mockResolvedValueOnce([{ tokens: null }]) + expect(await hasDecryptableOauthTokens('server-1')).toBe(false) + }) }) describe('withMcpOauthRefreshLock', () => { diff --git a/apps/sim/lib/mcp/oauth/storage.ts b/apps/sim/lib/mcp/oauth/storage.ts index 1e1f006b195..2d974a5af45 100644 --- a/apps/sim/lib/mcp/oauth/storage.ts +++ b/apps/sim/lib/mcp/oauth/storage.ts @@ -156,6 +156,26 @@ export async function loadOauthRow(params: { mcpServerId: string }): Promise { + const [row] = await db + .select({ tokens: mcpServerOauth.tokens }) + .from(mcpServerOauth) + .where(eq(mcpServerOauth.mcpServerId, mcpServerId)) + .limit(1) + if (!row?.tokens) return false + try { + await decryptSecret(row.tokens) + return true + } catch { + return false + } +} + export async function setOauthRowUser(rowId: string, userId: string): Promise { await db .update(mcpServerOauth) diff --git a/apps/sim/lib/mcp/resolve-config.test.ts b/apps/sim/lib/mcp/resolve-config.test.ts index 54e2244dcb6..baa6e0a6985 100644 --- a/apps/sim/lib/mcp/resolve-config.test.ts +++ b/apps/sim/lib/mcp/resolve-config.test.ts @@ -1,6 +1,6 @@ import { environmentUtilsMockFns } from '@sim/testing/mocks/environment-utils.mock' import { beforeEach, describe, expect, it, vi } from 'vitest' -import { resolveMcpConfigEnvVars } from '@/lib/mcp/resolve-config' +import { findMissingMcpConfigEnvVars, resolveMcpConfigEnvVars } from '@/lib/mcp/resolve-config' const mockGetEffectiveEnvironmentSnapshot = environmentUtilsMockFns.mockGetEffectiveEnvironmentSnapshot @@ -171,3 +171,63 @@ describe('resolveMcpConfigEnvVars secret provenance', () => { ]) }) }) + +describe('findMissingMcpConfigEnvVars', () => { + it('names the variables this requester has no value for, without reading env for literal configs', async () => { + mockGetEffectiveEnvironmentSnapshot.mockReset() + mockGetEffectiveEnvironmentSnapshot.mockResolvedValue({ + personalEncrypted: {}, + workspaceEncrypted: {}, + personalDecrypted: { MCP_HOST: 'api.example.com' }, + workspaceDecrypted: {}, + conflicts: [], + decryptionFailures: [], + }) + + expect(await findMissingMcpConfigEnvVars(BASE_CONFIG, 'user-b', 'workspace-1')).toEqual([ + 'MCP_TOKEN', + ]) + expect( + await findMissingMcpConfigEnvVars( + { ...BASE_CONFIG, url: 'https://api.example.com/mcp', headers: {} }, + 'user-b', + 'workspace-1' + ) + ).toEqual([]) + expect(mockGetEffectiveEnvironmentSnapshot).toHaveBeenCalledTimes(1) + }) + + it('counts a variable that failed to decrypt as missing', async () => { + mockGetEffectiveEnvironmentSnapshot.mockResolvedValue({ + personalEncrypted: { MCP_HOST: 'host-encrypted', MCP_TOKEN: 'token-encrypted' }, + workspaceEncrypted: {}, + personalDecrypted: { MCP_HOST: 'api.example.com', MCP_TOKEN: '' }, + workspaceDecrypted: {}, + conflicts: [], + decryptionFailures: ['MCP_TOKEN'], + }) + + expect(await findMissingMcpConfigEnvVars(BASE_CONFIG, 'user-b', 'workspace-1')).toEqual([ + 'MCP_TOKEN', + ]) + }) + + it('counts an intentionally empty value as present', async () => { + mockGetEffectiveEnvironmentSnapshot.mockResolvedValue({ + personalEncrypted: { MCP_PATH_PREFIX: 'prefix-encrypted' }, + workspaceEncrypted: {}, + personalDecrypted: { MCP_PATH_PREFIX: '' }, + workspaceDecrypted: {}, + conflicts: [], + decryptionFailures: [], + }) + + expect( + await findMissingMcpConfigEnvVars( + { ...BASE_CONFIG, url: 'https://example.com{{MCP_PATH_PREFIX}}/mcp', headers: {} }, + 'user-b', + 'workspace-1' + ) + ).toEqual([]) + }) +}) diff --git a/apps/sim/lib/mcp/resolve-config.ts b/apps/sim/lib/mcp/resolve-config.ts index a707160056a..59bd8a3bd89 100644 --- a/apps/sim/lib/mcp/resolve-config.ts +++ b/apps/sim/lib/mcp/resolve-config.ts @@ -6,13 +6,14 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' +import { omit } from '@sim/utils/object' import { type EnvironmentResolutionSnapshot, getEffectiveEnvironmentSnapshot, } from '@/lib/environment/utils' import type { McpServerConfig } from '@/lib/mcp/types' import { recordSecretUsage } from '@/lib/secrets/usage/record' -import { resolveEnvVarReferences } from '@/executor/utils/reference-validation' +import { createEnvVarPattern, resolveEnvVarReferences } from '@/executor/utils/reference-validation' import { createIncompleteResolvedSecretTraceRegistry, createResolvedSecretTraceRegistry, @@ -149,3 +150,36 @@ export async function resolveMcpConfigEnvVars( resolvedSecretTraceProvenance, } } + +/** + * The variables `config` references (url, headers) that `userId` has no value for in this + * workspace. Read-only: it records no secret usage, so it can check whether a requester can + * authenticate to a server before anything connects. + */ +export async function findMissingMcpConfigEnvVars( + config: McpServerConfig, + userId: string, + workspaceId?: string +): Promise { + const values = [config.url, ...Object.values(config.headers ?? {})].filter( + (value): value is string => typeof value === 'string' + ) + if (!values.some((value) => createEnvVarPattern().test(value))) return [] + const env = await getEffectiveEnvironmentSnapshot(userId, workspaceId) + const merged = { ...env.personalDecrypted, ...env.workspaceDecrypted } + /** + * A value that failed to decrypt arrives as '' and is named in `decryptionFailures`; an + * intentionally empty one is only ''. The list does not say which slice failed, so a key + * counts as failed only when its effective value is also empty. The one ambiguity left, an + * intentionally empty workspace value over a personal one that failed, counts as missing, + * so best-effort discovery skips rather than connects. + */ + const failed = new Set(env.decryptionFailures) + const envVars = omit( + merged, + Object.keys(merged).filter((key) => failed.has(key) && merged[key] === '') + ) + const missingKeys: string[] = [] + for (const value of values) resolveEnvVarReferences(value, envVars, { missingKeys }) + return [...new Set(missingKeys)] +} diff --git a/apps/sim/lib/mcp/service.test.ts b/apps/sim/lib/mcp/service.test.ts index 476330ed9fb..b5839680ff7 100644 --- a/apps/sim/lib/mcp/service.test.ts +++ b/apps/sim/lib/mcp/service.test.ts @@ -1,6 +1,10 @@ import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' import { dbChainMockFns, loggerMock, resetDbChainMock } from '@sim/testing' -import { McpOauthRedirectRequiredMock, mcpOauthMock } from '@sim/testing/mocks/mcp-oauth.mock' +import { + McpOauthRedirectRequiredMock, + mcpOauthMock, + mcpOauthMockFns, +} from '@sim/testing/mocks/mcp-oauth.mock' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { @@ -10,6 +14,7 @@ const { mockDisconnect, mockGetWorkspaceServersRows, mockResolveEnvVars, + mockFindMissingEnvVars, mockValidateDomain, mockValidateSsrf, mockIsDomainAllowed, @@ -64,6 +69,7 @@ const { mockDisconnect, mockGetWorkspaceServersRows: vi.fn(), mockResolveEnvVars: vi.fn(), + mockFindMissingEnvVars: vi.fn(), mockValidateDomain: vi.fn(), mockValidateSsrf: vi.fn(), mockIsDomainAllowed: vi.fn(() => true), @@ -108,6 +114,7 @@ vi.mock('@/lib/mcp/oauth', () => mcpOauthMock) vi.mock('@/lib/mcp/resolve-config', () => ({ resolveMcpConfigEnvVars: (...args: unknown[]) => mockResolveEnvVars(...args), + findMissingMcpConfigEnvVars: (...args: unknown[]) => mockFindMissingEnvVars(...args), })) vi.mock('@/lib/mcp/storage', () => ({ @@ -482,3 +489,167 @@ describe('McpService.discoverTools per-server caching', () => { expect(mockListTools).toHaveBeenCalledTimes(2) }) }) + +/** + * A server's status, failure count, cooldown and OAuth row are shared by the workspace, but + * credentials resolve per requester. Discovery a member never asked for (an untagged server in + * a broad catalog query) must neither mark the server failing for everyone nor connect with + * credentials the requester lacks. + */ +describe('McpService best-effort discovery', () => { + const MEMBER_B = 'user-b' + + beforeEach(async () => { + resetDbChainMock() + wireSelectsToWorkspaceRows() + dbChainMockFns.returning.mockResolvedValue([{ id: 'mcp-a' }]) + mockListTools.mockReset() + mockConnect.mockReset() + mockConnect.mockResolvedValue(undefined) + mockDisconnect.mockResolvedValue(undefined) + mockIsDomainAllowed.mockReturnValue(true) + mockValidateSsrf.mockResolvedValue('1.2.3.4') + mockValidateDomain.mockImplementation(() => undefined) + mockResolveEnvVars.mockImplementation((config: { url: string }) => + Promise.resolve({ config: { ...config }, missingVars: [] }) + ) + mockFindMissingEnvVars.mockReset() + mockFindMissingEnvVars.mockResolvedValue([]) + mcpOauthMockFns.mockLoadOauthRow.mockReset() + mcpOauthMockFns.mockHasDecryptableOauthTokens.mockReset() + mcpOauthMockFns.mockGetOrCreateOauthRow.mockReset() + await mcpService.clearCache() + mockCacheAdapter.set.mockClear() + mockCacheAdapter.delete.mockClear() + }) + + function expectNoSharedStateWrites() { + expect(dbChainMockFns.set).not.toHaveBeenCalled() + expect(mockCacheAdapter.set).not.toHaveBeenCalledWith( + `workspace:${WORKSPACE_ID}:server:mcp-a:failure`, + expect.anything(), + expect.any(Number) + ) + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + } + + it("leaves status and cooldown alone on failure, so another member's targeted discovery still runs", async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockRejectedValueOnce(new Error('Connection refused')) + + await expect( + mcpService.discoverServerTools(MEMBER_B, 'mcp-a', WORKSPACE_ID, 'cache-aside', undefined, { + bestEffort: true, + }) + ).rejects.toThrow('Connection refused') + expectNoSharedStateWrites() + + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + const tools = await mcpService.discoverServerTools( + USER_ID, + 'mcp-a', + WORKSPACE_ID, + 'cache-aside', + undefined, + { requireComplete: true } + ) + expect(tools.map((t) => t.name)).toEqual(['a1']) + }) + + it('skips a server whose variables the requester lacks without connecting, even when cached', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([ + dbRow('mcp-a', 'A', { headers: { Authorization: 'Bearer {{MY_TOKEN}}' } }), + ]) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + expect(await mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)).toHaveLength(1) + resetDbChainMock() + wireSelectsToWorkspaceRows() + mockCacheAdapter.set.mockClear() + mockCacheAdapter.delete.mockClear() + mockListTools.mockClear() + mockConnect.mockClear() + + mockFindMissingEnvVars.mockImplementation(async (_config: unknown, userId: string) => + userId === MEMBER_B ? ['MY_TOKEN'] : [] + ) + await expect( + mcpService.discoverServerTools(MEMBER_B, 'mcp-a', WORKSPACE_ID, 'cache-aside', undefined, { + bestEffort: true, + }) + ).rejects.toThrow('MY_TOKEN') + expect(mockConnect).not.toHaveBeenCalled() + expect(mockListTools).not.toHaveBeenCalled() + expectNoSharedStateWrites() + + const tools = await mcpService.discoverServerTools( + USER_ID, + 'mcp-a', + WORKSPACE_ID, + 'cache-aside', + undefined, + { bestEffort: true } + ) + expect(tools.map((t) => t.name)).toEqual(['a1']) + }) + + it('skips an OAuth server nobody authorized without creating its OAuth row or marking it pending', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A', { authType: 'oauth' })]) + mcpOauthMockFns.mockHasDecryptableOauthTokens.mockResolvedValue(false) + + await expect( + mcpService.discoverServerTools(MEMBER_B, 'mcp-a', WORKSPACE_ID, 'cache-aside', undefined, { + bestEffort: true, + }) + ).rejects.toBeInstanceOf(McpOauthAuthorizationRequiredError) + expect(mcpOauthMockFns.mockGetOrCreateOauthRow).not.toHaveBeenCalled() + // loadOauthRow clears a column it cannot decrypt; the gate must not risk that. + expect(mcpOauthMockFns.mockLoadOauthRow).not.toHaveBeenCalled() + expect(mockConnect).not.toHaveBeenCalled() + expectNoSharedStateWrites() + }) + + it('connects an authorized OAuth server with a grant-preserving provider and writes nothing when it is rejected', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A', { authType: 'oauth' })]) + mcpOauthMockFns.mockHasDecryptableOauthTokens.mockResolvedValue(true) + mcpOauthMockFns.mockGetOrCreateOauthRow.mockResolvedValue({ + id: 'oauth-row-1', + tokens: { access_token: 'expired', token_type: 'bearer' }, + }) + mockListTools.mockRejectedValueOnce(new UnauthorizedError('refresh rejected')) + MockMcpClient.mockClear() + + await expect( + mcpService.discoverServerTools(MEMBER_B, 'mcp-a', WORKSPACE_ID, 'cache-aside', undefined, { + bestEffort: true, + }) + ).rejects.toThrow('refresh rejected') + + expect(MockMcpClient.mock.calls[0]?.[0]).toMatchObject({ + oauthCredentials: { initialProvider: { preserveGrant: true } }, + }) + expectNoSharedStateWrites() + }) + + it('warms the shared tool cache on success without stamping status', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + await mcpService.discoverServerTools( + MEMBER_B, + 'mcp-a', + WORKSPACE_ID, + 'cache-aside', + undefined, + { + bestEffort: true, + } + ) + + expect(dbChainMockFns.set).not.toHaveBeenCalled() + expect(mockCacheAdapter.set).toHaveBeenCalledWith( + `workspace:${WORKSPACE_ID}:server:mcp-a`, + [tool('a1', 'mcp-a')], + expect.any(Number) + ) + }) +}) diff --git a/apps/sim/lib/mcp/service.ts b/apps/sim/lib/mcp/service.ts index c558301875e..ee2cb01ee68 100644 --- a/apps/sim/lib/mcp/service.ts +++ b/apps/sim/lib/mcp/service.ts @@ -26,18 +26,20 @@ import { } from '@/lib/mcp/domain-check' import { getOrCreateOauthRow, + hasDecryptableOauthTokens, loadPreregisteredClient, McpOauthRedirectRequired, SimMcpOauthProvider, } from '@/lib/mcp/oauth' import type { McpOauthCredentials } from '@/lib/mcp/oauth/coordinated-fetch' -import { resolveMcpConfigEnvVars } from '@/lib/mcp/resolve-config' +import { findMissingMcpConfigEnvVars, resolveMcpConfigEnvVars } from '@/lib/mcp/resolve-config' import { createMcpCacheAdapter, getMcpCacheType, type McpCacheStorageAdapter, } from '@/lib/mcp/storage' import { + McpConnectionError, McpOauthAuthorizationRequiredError, type McpServerConfig, McpServerCooldownError, @@ -72,6 +74,19 @@ type ResolvedSecretTraceProvenanceCallback = (provenance: ResolvedSecretTracePro interface McpRequestOptions { signal?: AbortSignal requireComplete?: boolean + /** + * Discovery the requester did not ask for, such as an untagged server in a broad catalog + * query. The server's status, failure count, cooldown and OAuth row are shared by every + * member while credentials resolve per requester, so one member's missing credential must + * not mark the server failing for the rest. The guarantee is no destructive or health + * write: no status, failure count, cooldown or cache eviction, and the OAuth grant is never + * cleared, re-registered or restarted (a failed refresh skips the server). Two writes are + * allowed: the positive tool cache on success, and a successful token refresh, which keeps + * the grant valid for everyone and must persist where the server rotates refresh tokens. + * A server the requester cannot authenticate to is skipped, even when another member's + * discovery left its tools cached, without connecting. + */ + bestEffort?: boolean } interface McpToolExecutionOptions extends McpRequestOptions { @@ -412,7 +427,8 @@ class McpService { resolvedIP: string | null, userId?: string, resolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceV1, - signal?: AbortSignal + signal?: AbortSignal, + preserveGrant = false ): Promise { const securityPolicy = { requireConsent: true, @@ -447,7 +463,7 @@ class McpService { throw new McpOauthAuthorizationRequiredError(config.id, config.name) } const preregistered = await loadPreregisteredClient(config.id) - return new SimMcpOauthProvider({ row, preregistered }) + return new SimMcpOauthProvider({ row, preregistered, preserveGrant }) } const client = new McpClient({ config, @@ -580,7 +596,8 @@ class McpService { workspaceId: string, extraHeaders?: Record, onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback, - signal?: AbortSignal + signal?: AbortSignal, + preserveGrant = false ): () => Promise { return async () => { const { @@ -601,7 +618,8 @@ class McpService { resolvedIP, userId, resolvedSecretTraceProvenance, - signal + signal, + preserveGrant ) } } @@ -612,6 +630,10 @@ class McpService { * so the retry re-acquires a fresh one that re-resolves the credential. (OAuth * 401s are left to the caller's oauth-pending handling.) `listTools` is * idempotent, so the retry is always safe. + * + * `preserveGrant` connects an OAuth server with a provider that may persist a refreshed + * grant but never clears, re-registers or restarts it. Such a client is pooled apart, so + * neither kind of caller borrows the other's provider. */ private async fetchServerTools( config: McpServerConfig, @@ -619,14 +641,17 @@ class McpService { workspaceId: string, onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback, signal?: AbortSignal, - requireComplete = false + requireComplete = false, + preserveGrant = false ): Promise { + const poolKey = this.poolKey(config.id, workspaceId, userId) for (let attempt = 0; ; attempt++) { signal?.throwIfAborted() try { return await this.withServerClient( { - key: this.poolKey(config.id, workspaceId, userId), + key: + preserveGrant && config.authType === 'oauth' ? `${poolKey}:preserve-grant` : poolKey, serverId: config.id, allowPool: true, }, @@ -636,7 +661,8 @@ class McpService { workspaceId, undefined, onResolvedSecretTraceProvenance, - signal + signal, + preserveGrant ), (client) => { reportRetainedClientProvenance( @@ -1157,12 +1183,13 @@ class McpService { workspaceId, refresh, createInvocationProvenanceReporter(onResolvedSecretTraceProvenance), - options.signal, - options.requireComplete + options ) } - const inflightKey = `${workspaceId}:${serverId}:${userId}:${refresh}:partial-ok` + const bestEffort = options.bestEffort === true + // A best-effort discovery skips the bookkeeping a tracked one owes, so they never share. + const inflightKey = `${workspaceId}:${serverId}:${userId}:${refresh}:${bestEffort ? 'best-effort' : 'partial-ok'}` const existing = this.inflightServerDiscovery.get(inflightKey) if (existing) return existing @@ -1172,8 +1199,9 @@ class McpService { workspaceId, refresh, undefined, - undefined, - false + { + bestEffort, + } ).finally(() => { this.inflightServerDiscovery.delete(inflightKey) }) @@ -1187,14 +1215,18 @@ class McpService { workspaceId: string, refresh: McpDiscoveryRefresh, onResolvedSecretTraceProvenance?: ResolvedSecretTraceProvenanceCallback, - signal?: AbortSignal, - requireComplete = false + { signal, requireComplete = false, bestEffort = false }: McpRequestOptions = {} ): Promise { signal?.throwIfAborted() const requestId = generateRequestId() const discoveryStartedAt = new Date() const maxRetries = 2 + // Before the cache: another member's cached tools never reach a requester who cannot call them. + const requesterConfig = bestEffort + ? await this.loadAuthenticatableServerConfig(serverId, userId, workspaceId) + : undefined + if (refresh === 'cache-aside' && !requireComplete) { try { const cached = await this.cacheAdapter.get(serverCacheKey(workspaceId, serverId)) @@ -1220,7 +1252,7 @@ class McpService { `[${requestId}] Discovering tools from server ${serverId} for user ${userId}${attempt > 0 ? ` (attempt ${attempt + 1})` : ''}` ) - const config = await this.getServerConfig(serverId, workspaceId) + const config = requesterConfig ?? (await this.getServerConfig(serverId, workspaceId)) if (!config) { throw new Error(`Server ${serverId} not found or not accessible`) } @@ -1232,22 +1264,26 @@ class McpService { workspaceId, onResolvedSecretTraceProvenance, signal, - requireComplete + requireComplete, + bestEffort ) logger.info(`[${requestId}] Discovered ${tools.length} tools from server ${config.name}`) - await Promise.allSettled([ - this.cacheAdapter - .set(serverCacheKey(workspaceId, serverId), tools, this.cacheTimeout) - .catch((err) => - logger.warn(`[${requestId}] Cache write failed for ${config.name}:`, err) - ), - this.clearServerFailure(workspaceId, serverId), - this.updateServerStatus(serverId, workspaceId, { - outcome: 'connected', - toolCount: tools.length, - discoveryStartedAt, - }), - ]) + const cacheWrite = this.cacheAdapter + .set(serverCacheKey(workspaceId, serverId), tools, this.cacheTimeout) + .catch((err) => logger.warn(`[${requestId}] Cache write failed for ${config.name}:`, err)) + await Promise.allSettled( + bestEffort + ? [cacheWrite] + : [ + cacheWrite, + this.clearServerFailure(workspaceId, serverId), + this.updateServerStatus(serverId, workspaceId, { + outcome: 'connected', + toolCount: tools.length, + discoveryStartedAt, + }), + ] + ) return tools } catch (error) { signal?.throwIfAborted() @@ -1263,6 +1299,7 @@ class McpService { signal?.throwIfAborted() continue } + if (bestEffort) throw error // Drop positive cache so a follow-up doesn't return stale tools. const statusApplied = isOauthAuthorizationError(error, authType) ? await this.markServerOauthPending(serverId, workspaceId, discoveryStartedAt) @@ -1288,6 +1325,28 @@ class McpService { throw new Error(`Failed to discover tools from server ${serverId} after ${maxRetries} attempts`) } + /** + * The server's config, once it is known `userId` can authenticate to it without trying: + * every variable it references resolves for them, and an OAuth server has been authorized. + */ + private async loadAuthenticatableServerConfig( + serverId: string, + userId: string, + workspaceId: string + ): Promise { + const config = await this.getServerConfig(serverId, workspaceId) + if (!config) throw new Error(`Server ${serverId} not found or not accessible`) + const missing = await findMissingMcpConfigEnvVars(config, userId, workspaceId) + if (missing.length) + throw new McpConnectionError( + `Missing environment variables for this user: ${missing.join(', ')}`, + config.name + ) + if (config.authType === 'oauth' && !(await hasDecryptableOauthTokens(config.id))) + throw new McpOauthAuthorizationRequiredError(config.id, config.name) + return config + } + async getServerSummaries(userId: string, workspaceId: string): Promise { const requestId = generateRequestId() diff --git a/apps/sim/lib/mothership/integrations/application/catalog.test.ts b/apps/sim/lib/mothership/integrations/application/catalog.test.ts index b495ab34592..277d2359097 100644 --- a/apps/sim/lib/mothership/integrations/application/catalog.test.ts +++ b/apps/sim/lib/mothership/integrations/application/catalog.test.ts @@ -379,11 +379,18 @@ describe('integration catalog ranking', () => { description: 'Set a key only if it does not exist', input_schema: { type: 'object', properties: {} }, } + const doe = { + name: 'crm_find_doe', + description: 'Find a Jane Doe record', + input_schema: { type: 'object', properties: {} }, + } expect( - projectIntegrationCatalog([setnx], { ...input, query: 'what does it do', limit: 0 }) - ).toMatchObject({ - total: 0, - }) + projectIntegrationCatalog([setnx, doe], { + ...input, + query: 'what does setnx do', + limit: 0, + }).operations.map((operation) => operation.toolId) + ).toEqual(['redis_setnx']) }) it('folds -es plurals onto their singular', () => { @@ -407,12 +414,22 @@ describe('integration catalog ranking', () => { ).toEqual([query.startsWith('branch') ? 'github_create_branch' : 'snowflake_get_warehouse']) }) - it('matches nothing for a query made only of stop words and one-character terms', () => { - for (const query of ['a the to of', 'to a', 'x y z', 'gets uses']) - expect(projectIntegrationCatalog(emailCapable, { ...input, query, limit: 0 })).toMatchObject({ - total: 0, - operations: [], - }) + /** An agent reading total 0 would conclude the workspace has no integrations at all. */ + it('lists every operation, compactly, for a query with no content terms', () => { + const listing = projectIntegrationCatalog(emailCapable, { ...input, limit: 0 }) + for (const query of [ + 'all', + 'any', + 'what can I use', + 'a the to of', + 'x y z', + 'gets uses', + '发送邮件', + ]) + expect(projectIntegrationCatalog(emailCapable, { ...input, query, limit: 0 })).toEqual( + listing + ) + expect(listing.total).toBe(emailCapable.length) }) it('only counts operations that share a content term with the intent', () => { @@ -892,11 +909,19 @@ describe('untagged workspace MCP servers', () => { expect(toolIds(result)).toEqual(['mcp-crm-lookup', 'mcp-docs-lookup']) expect(mocks.discover).toHaveBeenCalledWith({ principal: expect.objectContaining({ workspaceId: 'workspace-1', subjectUserId: 'actor' }), - input: { workspaceId: 'workspace-1', serverId: 'mcp-docs' }, + input: { workspaceId: 'workspace-1', serverId: 'mcp-docs', bestEffort: true }, }) } ) + it('keeps built-ins and tagged servers when the workspace server list fails', async () => { + mocks.listServers.mockRejectedValue(new Error('connection terminated')) + const result = await catalog({ mcpServerIds: ['mcp-docs'] }) + expect(toolIds(result)).toEqual(['mcp-docs-lookup', 'gmail_send', 'slack_send']) + expect(discoveredServerIds()).toEqual(['mcp-docs']) + await expect(catalog({ service: 'mcp:mcp-crm' })).rejects.toThrow('connection terminated') + }) + it('keeps Search free of MCP even with usable workspace servers', async () => { const result = await catalog({}, 'assistant') expect(toolIds(result)).toEqual(['gmail_send', 'slack_send']) diff --git a/apps/sim/lib/mothership/integrations/application/catalog.ts b/apps/sim/lib/mothership/integrations/application/catalog.ts index b767836b110..8924d8be430 100644 --- a/apps/sim/lib/mothership/integrations/application/catalog.ts +++ b/apps/sim/lib/mothership/integrations/application/catalog.ts @@ -166,7 +166,8 @@ interface CatalogProjectionOptions { * Return schemas only for targeted discovery (query, service or exact ID), while keeping * explicit full listings compact. Free-text queries rank by BM25, with operations of the * requester's connected services boosted; an exact operation ID always ranks first, and - * preferred services break ties in score. + * preferred services break ties in score. A query with no content terms ("what can I + * use", or a script the tokenizer does not split) asks for the listing, as an empty one does. */ export function projectIntegrationCatalog( tools: ToolSchema[], @@ -183,11 +184,12 @@ export function projectIntegrationCatalog( (!input.toolId || tool.name === input.toolId) && (!services || (tool.service && services.has(tool.service.toLowerCase()))) ) - const connected = normalized ? resolveConnectedServices(candidates, connectedProviderIds) : null - const scores = normalized ? rankByBm25(candidates, queryTerms) : [] + const ranked = queryTerms.length > 0 + const connected = ranked ? resolveConnectedServices(candidates, connectedProviderIds) : null + const scores = ranked ? rankByBm25(candidates, queryTerms) : [] const matches = candidates.flatMap((tool, index) => { - if (!normalized) return [{ tool, exact: false, score: 0 }] - const exact = tool.name.toLowerCase() === normalized + const exact = Boolean(normalized) && tool.name.toLowerCase() === normalized + if (!ranked) return [{ tool, exact, score: 0 }] const score = (scores[index] ?? 0) * (tool.service && connected?.has(tool.service.toLowerCase()) ? CONNECTED_SERVICE_BOOST : 1) @@ -201,7 +203,7 @@ export function projectIntegrationCatalog( a.tool.name.localeCompare(b.tool.name) ) const selected = input.limit === 0 ? matches : matches.slice(0, input.limit) - const targeted = Boolean(input.toolId || normalized || services) + const targeted = Boolean(input.toolId || ranked || services) return { total: matches.length, truncated: selected.length < matches.length, @@ -361,14 +363,32 @@ const catalogUseCase = defineAuthorizedChatUseCase({ let complete = named let bestEffort: string[] = [] if (!selectionScoped) { - const usable = await listUsableMcpServerIds(context.userId, workspaceId) - const untagged = [...usable].filter( + /** Broad discovery is optional: without the server list it keeps the tagged servers and built-ins. */ + const usable = await listUsableMcpServerIds(context.userId, workspaceId).catch( + (error: unknown) => { + if (targeted) throw error + logger.warn('Workspace MCP servers unavailable; discovering only tagged servers', { + workspaceId, + error: getErrorMessage(error), + }) + return null + } + ) + const untagged = [...(usable ?? [])].filter( (serverId) => isTarget(serverId) && !taggedIds.has(serverId) ) /** A targeted server is listed completely; a tagged one keeps its own error even once unusable. */ if (targeted) complete = [...named, ...untagged] - else { + else if (usable) { complete = named.filter((serverId) => usable.has(serverId)) + for (const serverId of named) + if (!usable.has(serverId)) + logger.warn( + 'Omitting a tagged MCP server that is no longer usable from the integration catalog', + { + serverId, + } + ) bestEffort = untagged } } diff --git a/apps/sim/lib/mothership/mcp-tools.ts b/apps/sim/lib/mothership/mcp-tools.ts index 8d5d67ae562..d0099273b4e 100644 --- a/apps/sim/lib/mothership/mcp-tools.ts +++ b/apps/sim/lib/mothership/mcp-tools.ts @@ -34,6 +34,8 @@ interface McpDiscoveryServers { * Chat-only: servers available without being named. Served from the discovery * cache when fresh, otherwise skipped after {@link BEST_EFFORT_DISCOVERY_WAIT_MS}. * One with no permitted operations contributes nothing and is not a failure. + * Their discovery never touches a server's shared status or cooldown, and skips a + * server the requester cannot authenticate to. */ bestEffort?: string[] } @@ -94,7 +96,7 @@ export async function buildMcpToolSchemas( const discoverBestEffort = async (serverId: string) => { const discovery = discoverMcpServerToolsUseCase.execute({ principal: principal(), - input: { workspaceId, serverId }, + input: { workspaceId, serverId, bestEffort: true }, }) const { tools } = await withinDeadline( () => discovery, diff --git a/packages/testing/src/mocks/mcp-oauth.mock.ts b/packages/testing/src/mocks/mcp-oauth.mock.ts index 96ab558ef2a..952734e5962 100644 --- a/packages/testing/src/mocks/mcp-oauth.mock.ts +++ b/packages/testing/src/mocks/mcp-oauth.mock.ts @@ -17,6 +17,7 @@ export const mcpOauthMockFns = { mockMcpAuthGuarded: vi.fn(), mockGetOrCreateOauthRow: vi.fn(), mockLoadOauthRow: vi.fn(), + mockHasDecryptableOauthTokens: vi.fn(), mockLoadOauthRowByState: vi.fn(), mockLoadPreregisteredClient: vi.fn(), mockSetOauthRowUser: vi.fn(), @@ -78,6 +79,7 @@ export const mcpOauthMock = { mcpAuthGuarded: mcpOauthMockFns.mockMcpAuthGuarded, getOrCreateOauthRow: mcpOauthMockFns.mockGetOrCreateOauthRow, loadOauthRow: mcpOauthMockFns.mockLoadOauthRow, + hasDecryptableOauthTokens: mcpOauthMockFns.mockHasDecryptableOauthTokens, loadOauthRowByState: mcpOauthMockFns.mockLoadOauthRowByState, loadPreregisteredClient: mcpOauthMockFns.mockLoadPreregisteredClient, setOauthRowUser: mcpOauthMockFns.mockSetOauthRowUser,