diff --git a/apps/sim/lib/knowledge/mcp/route-handler.protocol.test.ts b/apps/sim/lib/knowledge/mcp/route-handler.protocol.test.ts new file mode 100644 index 00000000000..a6be8a5ee7a --- /dev/null +++ b/apps/sim/lib/knowledge/mcp/route-handler.protocol.test.ts @@ -0,0 +1,196 @@ +import { member } from '@sim/db/schema' +import { queueTableRows, resetDbChainMock } from '@sim/testing' +import { createPersonalApiKeyPrincipal } from '@sim/testing/factories/principal.factory' +import { createRouteContext } from '@sim/testing/helpers/http' +import { credentialGroupsServiceMock } from '@sim/testing/mocks/credential-groups-service.mock' +import { + knowledgeAvailabilityMock, + knowledgeAvailabilityMockFns, +} from '@sim/testing/mocks/knowledge-availability.mock' +import { knowledgeBaseUseCasesMock } from '@sim/testing/mocks/knowledge-base-use-cases.mock' +import { knowledgeContextsMock } from '@sim/testing/mocks/knowledge-contexts.mock' +import { knowledgeEmbeddingsMock } from '@sim/testing/mocks/knowledge-embeddings.mock' +import { knowledgeServiceMock } from '@sim/testing/mocks/knowledge-service.mock' +import { + permissionGroupsResolveMock, + permissionGroupsResolveMockFns, +} from '@sim/testing/mocks/permission-groups-resolve.mock' +import { rateLimiterMock } from '@sim/testing/mocks/rate-limiter.mock' +import { createMockRequest } from '@sim/testing/mocks/request.mock' +import { simSearchConnectorsMock } from '@sim/testing/mocks/sim-search-connectors.mock' +import { urlsMockFns } from '@sim/testing/mocks/urls.mock' +import { v2ApiKeyAuthModuleMock, v2RouteMocks } from '@sim/testing/mocks/v2-route.mock' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ index: vi.fn(), liveSearch: vi.fn() })) +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => v2ApiKeyAuthModuleMock) +vi.mock('@/lib/core/rate-limiter', () => rateLimiterMock) +vi.mock('@/lib/core/utils/after-response', () => ({ afterResponse: vi.fn() })) +vi.mock('@/lib/knowledge/mcp/activity', () => ({ recordOrganizationSearchMcpActivity: vi.fn() })) +vi.mock('@/lib/knowledge/application/contexts', () => knowledgeContextsMock) +vi.mock('@/lib/knowledge/application/billing', () => ({ + resolveKnowledgeAttributedUserId: (principal: { userId: string }) => principal.userId, +})) +vi.mock('@/lib/knowledge/search/search-index', () => ({ findSearchIndex: mocks.index })) +vi.mock('@/lib/permission-groups/resolve.server', () => permissionGroupsResolveMock) +vi.mock('@/lib/credential-groups/service', () => credentialGroupsServiceMock) +vi.mock('@/lib/knowledge/service', () => knowledgeServiceMock) +vi.mock('@/lib/knowledge/embeddings', () => knowledgeEmbeddingsMock) +vi.mock('@/lib/knowledge/application/knowledge-bases', () => knowledgeBaseUseCasesMock) +vi.mock('@/lib/knowledge/application/connectors', () => ({ + createKnowledgeConnector: { execute: vi.fn() }, +})) +vi.mock('@/lib/knowledge/application/connector-access', () => ({ + startKnowledgeConnectorMemberEnrollment: { execute: vi.fn() }, +})) +vi.mock('@/lib/knowledge/access/availability', () => knowledgeAvailabilityMock) +vi.mock('@/connectors/registry', () => ({ CONNECTOR_META_REGISTRY: {} })) +vi.mock('@/lib/sim-search/connectors', () => simSearchConnectorsMock) +vi.mock('@/lib/sim-search/live/application', async (importOriginal) => ({ + ...(await importOriginal()), + searchLiveKnowledge: { execute: mocks.liveSearch }, +})) + +import { createKnowledgeMcpHandlers } from '@/lib/knowledge/mcp/route-handler' + +urlsMockFns.mockGetBaseUrl.mockReturnValue('http://localhost') + +const handlers = createKnowledgeMcpHandlers() +const principal = createPersonalApiKeyPrincipal({ userId: 'person-1' }) + +/** Calls a tool over the real route, SDK transport, and Search MCP server. */ +async function callTool(name: string, args: Record) { + const response = await handlers.POST( + createMockRequest({ + method: 'POST', + url: 'http://localhost/api/mcp/search/organizations/org-1', + headers: { + 'content-type': 'application/json', + accept: 'application/json, text/event-stream', + 'x-forwarded-for': '127.0.0.1', + 'x-api-key': 'personal-key', + }, + body: { + jsonrpc: '2.0', + id: 1, + method: 'tools/call', + params: { name, arguments: args }, + }, + }), + createRouteContext({ organizationId: 'org-1' }) + ) + return { status: response.status, body: await response.json() } +} + +function callSearch(args: Record) { + return callTool('search', args) +} + +function searchedInput() { + const call = mocks.liveSearch.mock.lastCall?.[0] + if (!call) throw new Error('Expected the search use case to run') + return call.input +} + +beforeEach(() => { + resetDbChainMock() + v2RouteMocks.authenticate.mockResolvedValue({ + principal, + keyType: 'personal', + keyExpiresAt: null, + rateLimitSubjectIds: ['person-1'], + rateLimitSubscription: null, + }) + permissionGroupsResolveMockFns.mockGetUserPermissionConfigForOrganization.mockResolvedValue(null) + knowledgeAvailabilityMockFns.mockRequireOrganizationSearchAvailable.mockResolvedValue(undefined) + mocks.index.mockResolvedValue({ id: 'index-1' }) + mocks.liveSearch.mockResolvedValue({ + results: [], + retrieval: { status: 'complete', timedOutLegs: [] }, + live: { backend: 'live', accounts: [], guidance: '' }, + }) + queueTableRows(member, [{ role: 'member' }]) +}) + +describe('Search MCP tool arguments with one possible reading', () => { + it('searches without native queries when the client sends an empty nativeQueries list', async () => { + const { body } = await callSearch({ query: 'release notes', nativeQueries: [] }) + expect(body.result?.isError).not.toBe(true) + expect(searchedInput().nativeQueries).toBeUndefined() + }) + + it('treats a null optional native query field as omitted', async () => { + const { body } = await callSearch({ + query: 'release', + nativeQueries: [{ provider: 'jira', query: 'text ~ "release"', kind: null }], + }) + expect(body.result?.isError).not.toBe(true) + expect(searchedInput().nativeQueries).toEqual([{ provider: 'jira', query: 'text ~ "release"' }]) + }) + + it('applies filters every native query carries alike as the search filters', async () => { + const { body } = await callSearch({ + query: 'release', + nativeQueries: [ + { + provider: 'gmail', + query: 'subject:release', + startDate: '2026-10-01T00:00:00-07:00', + sortBy: 'newest', + }, + { + provider: 'slack', + query: 'release', + startDate: '2026-10-01T00:00:00-07:00', + sortBy: 'newest', + }, + ], + }) + expect(body.result?.isError).not.toBe(true) + const input = searchedInput() + expect(input.filters).toMatchObject({ + startDate: '2026-10-01T00:00:00-07:00', + sortBy: 'newest', + }) + expect(input.nativeQueries).toEqual([ + { provider: 'gmail', query: 'subject:release' }, + { provider: 'slack', query: 'release' }, + ]) + }) + + it('reads a top-level provider as the source filter', async () => { + const { body } = await callSearch({ query: 'standup', provider: 'google_calendar' }) + expect(body.result?.isError).not.toBe(true) + expect(searchedInput().filters).toMatchObject({ source: 'google_calendar' }) + }) + + it.each([ + ['carry different values', '2026-09-01T00:00:00Z'], + ['do not all carry it', undefined], + ])('still rejects native queries whose dates %s, rather than choosing one', async (_, other) => { + const { body } = await callSearch({ + query: 'release', + nativeQueries: [ + { provider: 'gmail', query: 'release', startDate: '2026-10-01T00:00:00Z' }, + { provider: 'slack', query: 'release', ...(other ? { startDate: other } : {}) }, + ], + }) + expect(JSON.stringify(body)).toContain('startDate') + expect(mocks.liveSearch).not.toHaveBeenCalled() + }) + + it('still rejects a top-level provider that names no provider', async () => { + const { body } = await callSearch({ query: 'standup', provider: 'Google Calendar' }) + expect(JSON.stringify(body)).toContain('provider') + expect(mocks.liveSearch).not.toHaveBeenCalled() + }) + + it.each(['constructor', 'toString', '__proto__'])( + 'answers a call to the unknown tool %s with an MCP error, not a server error', + async (name) => { + const { status, body } = await callTool(name, {}) + expect(status).toBe(200) + expect(JSON.stringify(body)).toContain(name) + } + ) +}) diff --git a/apps/sim/lib/knowledge/mcp/route-handler.ts b/apps/sim/lib/knowledge/mcp/route-handler.ts index 980672d468b..4748f7b9053 100644 --- a/apps/sim/lib/knowledge/mcp/route-handler.ts +++ b/apps/sim/lib/knowledge/mcp/route-handler.ts @@ -1,3 +1,4 @@ +import { createLogger } from '@sim/logger' import type { NextRequest } from 'next/server' import { organizationKnowledgeMcpContract } from '@/lib/api/contracts/knowledge/mcp' import { parseRequest } from '@/lib/api/server' @@ -13,9 +14,12 @@ import { knowledgeOperations } from '@/lib/knowledge/application/operations' import { readSearchIndex } from '@/lib/knowledge/application/sim-search' import { withSearchMcpAuthChallenge } from '@/lib/knowledge/mcp/oauth-metadata' import { createKnowledgeMcpServer } from '@/lib/knowledge/mcp/server' +import { repairToolCallMessage, toolCallInputIssues } from '@/lib/knowledge/mcp/tool-arguments' import { getSearchMcpUrl } from '@/lib/knowledge/mcp/urls' import { v2CaughtOrchestrationError, v2Error } from '@/app/api/v2/lib/response' +const logger = createLogger('KnowledgeMcp') + function mcpAuth(resource: string) { return mcpCredentialAuth({ resource, allowUnboundApiTokens: true }) } @@ -47,12 +51,19 @@ export function createKnowledgeMcpHandlers() { input: parsed.data.params, request, }) + const { organizationId } = parsed.data.params + const { message, repairs } = repairToolCallMessage(parsed.data.body) + if (repairs.length) + logger.info('Knowledge MCP tool input repaired', { organizationId, repairs }) + const rejected = toolCallInputIssues(message) + if (rejected) + logger.warn('Knowledge MCP tool input rejected', { organizationId, ...rejected }) const server = createKnowledgeMcpServer({ request, auth: admission.auth, ...parsed.data.params, }) - return await serveStatelessMcp(server, request, parsed.data.body) + return await serveStatelessMcp(server, request, message) } catch (error) { const response = v2CaughtOrchestrationError(error) if (response) return withSearchMcpAuthChallenge(response, resource) diff --git a/apps/sim/lib/knowledge/mcp/tool-arguments.ts b/apps/sim/lib/knowledge/mcp/tool-arguments.ts new file mode 100644 index 00000000000..0e0603ed0ae --- /dev/null +++ b/apps/sim/lib/knowledge/mcp/tool-arguments.ts @@ -0,0 +1,104 @@ +import { isRecordLike } from '@sim/utils/object' +import type { z } from 'zod' +import { + chatSearchMcpSchema, + liveSearchMcpSchema, + readLiveDocumentMcpSchema, +} from '@/lib/api/contracts/knowledge/mcp' +import { mcpToolCallParamsSchema } from '@/lib/api/contracts/mcp' +import { liveSearchProviderSchema } from '@/lib/api/contracts/mothership-assistant-tools' + +/** Search filters a client sometimes writes on each native query instead of on the search. */ +const SEARCH_FILTERS_ON_NATIVE_QUERIES = [ + 'startDate', + 'endDate', + 'modifiedAfter', + 'modifiedBefore', + 'sortBy', + 'topK', +] as const + +/** Keyed by client-supplied tool names, so a name such as `constructor` finds nothing. */ +const TOOL_INPUT_SCHEMAS = new Map([ + ['search', liveSearchMcpSchema], + ['read_document', readLiveDocumentMcpSchema], + ['chat', chatSearchMcpSchema], +]) + +function readToolCall(message: Record) { + if (message.method !== 'tools/call') return undefined + const params = mcpToolCallParamsSchema.safeParse(message.params) + return params.success ? params.data : undefined +} + +/** + * Search arguments rewritten where the schema rejects them but they have exactly one reading: an + * empty `nativeQueries` list, a null optional native query field, a top-level `provider` naming a + * provider (the `source` filter), and a search filter every native query carries with the same value. Native + * queries that disagree on a filter are left for the schema to reject. + */ +function repairSearchArguments(args: Record) { + const { nativeQueries: queries, provider, ...rest } = args + const repaired: Record = rest + const repairs = new Set() + const providerId = liveSearchProviderSchema.safeParse(provider) + if (providerId.success && (rest.source === undefined || rest.source === providerId.data)) { + repaired.source = providerId.data + repairs.add('providerAsSource') + } else if (provider !== undefined) repaired.provider = provider + if (Array.isArray(queries) && queries.length === 0) repairs.add('emptyNativeQueries') + else if (queries !== undefined) repaired.nativeQueries = queries + if (!Array.isArray(queries) || !queries.length || !queries.every(isRecordLike)) + return { args: repaired, repairs } + let items = queries.map((query) => + Object.fromEntries( + Object.entries(query).filter(([key, value]) => { + if (value !== null || key === 'provider' || key === 'query') return true + repairs.add('nullNativeQueryField') + return false + }) + ) + ) + for (const filter of SEARCH_FILTERS_ON_NATIVE_QUERIES) { + if (!items.some((item) => filter in item)) continue + /** A query without the filter adds `undefined`, so a filter not on every query differs too. */ + const values = new Set(items.map((item) => JSON.stringify(item[filter]))) + const [value] = values + if (values.size !== 1 || value === undefined) continue + if (filter in repaired && JSON.stringify(repaired[filter]) !== value) continue + repaired[filter] = JSON.parse(value) + items = items.map(({ [filter]: _, ...rest }) => rest) + repairs.add('nativeQueryFilterAsSearchFilter') + } + repaired.nativeQueries = items + return { args: repaired, repairs } +} + +/** + * The JSON-RPC message with a search call's arguments repaired before the SDK validates them, so + * a call the client clearly meant runs instead of costing it another round. Arguments the schema + * already accepts, other tools, and tools/list come back unchanged. + */ +export function repairToolCallMessage(message: Record) { + const call = readToolCall(message) + if (call?.name !== 'search' || !call.arguments) return { message, repairs: [] } + const { args, repairs } = repairSearchArguments(call.arguments) + if (!repairs.size) return { message, repairs: [] } + return { message: { ...message, params: { ...call, arguments: args } }, repairs: [...repairs] } +} + +/** + * Where a tool call's arguments fail its input schema, as field paths and issue codes without + * values. The SDK rejects such a call before the tool runs, so nothing else records it. + */ +export function toolCallInputIssues(message: Record) { + const call = readToolCall(message) + const schema = call ? TOOL_INPUT_SCHEMAS.get(call.name) : undefined + if (!call || !schema) return undefined + const parsed = schema.safeParse(call.arguments ?? {}) + if (parsed.success) return undefined + return { + toolName: call.name, + issues: parsed.error.issues.map(({ code, path }) => ({ code, path: path.join('.') })), + } +}