fix(search): look up tool schemas by own name and map only real provider IDs to source

Author: waleedlatif1

apps/sim/lib/knowledge/mcp/route-handler.protocol.test.ts

@@ -58,8 +58,8 @@ urlsMockFns.mockGetBaseUrl.mockReturnValue('http://localhost')
 const handlers = createKnowledgeMcpHandlers()
 const principal = createPersonalApiKeyPrincipal({ userId: 'person-1' })
 
-/** Calls the search tool over the real route, SDK transport, and Search MCP server. */
-async function callSearch(args: Record<string, unknown>) {
+/** 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',
@@ -74,14 +74,18 @@ async function callSearch(args: Record<string, unknown>) {
         jsonrpc: '2.0',
         id: 1,
         method: 'tools/call',
-        params: { name: 'search', arguments: args },
+        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')
@@ -174,4 +178,19 @@ describe('Search MCP tool arguments with one possible reading', () => {
     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/tool-arguments.ts

@@ -6,6 +6,7 @@ import {
   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 = [
@@ -17,11 +18,12 @@ const SEARCH_FILTERS_ON_NATIVE_QUERIES = [
   'topK',
 ] as const
 
-const TOOL_INPUT_SCHEMAS: Record<string, z.ZodType> = {
-  search: liveSearchMcpSchema,
-  read_document: readLiveDocumentMcpSchema,
-  chat: chatSearchMcpSchema,
-}
+/** 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
@@ -31,16 +33,17 @@ function readToolCall(message: Record<string, unknown>) {
 
 /**
  * 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` (the
- * `source` filter), and a search filter every native query carries with the same value. Native
+ * 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>()
-  if (typeof provider === 'string' && (rest.source === undefined || rest.source === provider)) {
-    repaired.source = provider
+  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')
@@ -90,7 +93,7 @@ export function repairToolCallMessage(message: Record<string, unknown>) {
  */
 export function toolCallInputIssues(message: Record<string, unknown>) {
   const call = readToolCall(message)
-  const schema = call ? TOOL_INPUT_SCHEMAS[call.name] : undefined
+  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