improvement(search): run Search MCP calls whose arguments have one reading, and log rejected ones

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<typeof import('@/lib/sim-search/live/application')>()),
+  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<string, unknown>) {
+  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<string, unknown>) {
+  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)
+    }
+  )
+})

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)

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<string, z.ZodType>([
+  ['search', liveSearchMcpSchema],
+  ['read_document', readLiveDocumentMcpSchema],
+  ['chat', chatSearchMcpSchema],
+])
+
+function readToolCall(message: Record<string, unknown>) {
+  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<string, unknown>) {
+  const { nativeQueries: queries, provider, ...rest } = args
+  const repaired: Record<string, unknown> = rest
+  const repairs = new Set<string>()
+  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<string, unknown>) {
+  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<string, unknown>) {
+  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('.') })),
+  }
+}