Skip to content

Commit eae86a6

Browse files
authored
improvement(search): run Search MCP calls whose arguments have one reading, and log rejected ones (#9814)
* improvement(search): run Search MCP calls whose arguments have one reading, and log rejected ones * test(search): use the default organization owner mock in the Search MCP protocol test * fix(search): look up tool schemas by own name and map only real provider IDs to source
1 parent 7cacf2b commit eae86a6

3 files changed

Lines changed: 312 additions & 1 deletion

File tree

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
import { member } from '@sim/db/schema'
2+
import { queueTableRows, resetDbChainMock } from '@sim/testing'
3+
import { createPersonalApiKeyPrincipal } from '@sim/testing/factories/principal.factory'
4+
import { createRouteContext } from '@sim/testing/helpers/http'
5+
import { credentialGroupsServiceMock } from '@sim/testing/mocks/credential-groups-service.mock'
6+
import {
7+
knowledgeAvailabilityMock,
8+
knowledgeAvailabilityMockFns,
9+
} from '@sim/testing/mocks/knowledge-availability.mock'
10+
import { knowledgeBaseUseCasesMock } from '@sim/testing/mocks/knowledge-base-use-cases.mock'
11+
import { knowledgeContextsMock } from '@sim/testing/mocks/knowledge-contexts.mock'
12+
import { knowledgeEmbeddingsMock } from '@sim/testing/mocks/knowledge-embeddings.mock'
13+
import { knowledgeServiceMock } from '@sim/testing/mocks/knowledge-service.mock'
14+
import {
15+
permissionGroupsResolveMock,
16+
permissionGroupsResolveMockFns,
17+
} from '@sim/testing/mocks/permission-groups-resolve.mock'
18+
import { rateLimiterMock } from '@sim/testing/mocks/rate-limiter.mock'
19+
import { createMockRequest } from '@sim/testing/mocks/request.mock'
20+
import { simSearchConnectorsMock } from '@sim/testing/mocks/sim-search-connectors.mock'
21+
import { urlsMockFns } from '@sim/testing/mocks/urls.mock'
22+
import { v2ApiKeyAuthModuleMock, v2RouteMocks } from '@sim/testing/mocks/v2-route.mock'
23+
import { beforeEach, describe, expect, it, vi } from 'vitest'
24+
25+
const mocks = vi.hoisted(() => ({ index: vi.fn(), liveSearch: vi.fn() }))
26+
vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => v2ApiKeyAuthModuleMock)
27+
vi.mock('@/lib/core/rate-limiter', () => rateLimiterMock)
28+
vi.mock('@/lib/core/utils/after-response', () => ({ afterResponse: vi.fn() }))
29+
vi.mock('@/lib/knowledge/mcp/activity', () => ({ recordOrganizationSearchMcpActivity: vi.fn() }))
30+
vi.mock('@/lib/knowledge/application/contexts', () => knowledgeContextsMock)
31+
vi.mock('@/lib/knowledge/application/billing', () => ({
32+
resolveKnowledgeAttributedUserId: (principal: { userId: string }) => principal.userId,
33+
}))
34+
vi.mock('@/lib/knowledge/search/search-index', () => ({ findSearchIndex: mocks.index }))
35+
vi.mock('@/lib/permission-groups/resolve.server', () => permissionGroupsResolveMock)
36+
vi.mock('@/lib/credential-groups/service', () => credentialGroupsServiceMock)
37+
vi.mock('@/lib/knowledge/service', () => knowledgeServiceMock)
38+
vi.mock('@/lib/knowledge/embeddings', () => knowledgeEmbeddingsMock)
39+
vi.mock('@/lib/knowledge/application/knowledge-bases', () => knowledgeBaseUseCasesMock)
40+
vi.mock('@/lib/knowledge/application/connectors', () => ({
41+
createKnowledgeConnector: { execute: vi.fn() },
42+
}))
43+
vi.mock('@/lib/knowledge/application/connector-access', () => ({
44+
startKnowledgeConnectorMemberEnrollment: { execute: vi.fn() },
45+
}))
46+
vi.mock('@/lib/knowledge/access/availability', () => knowledgeAvailabilityMock)
47+
vi.mock('@/connectors/registry', () => ({ CONNECTOR_META_REGISTRY: {} }))
48+
vi.mock('@/lib/sim-search/connectors', () => simSearchConnectorsMock)
49+
vi.mock('@/lib/sim-search/live/application', async (importOriginal) => ({
50+
...(await importOriginal<typeof import('@/lib/sim-search/live/application')>()),
51+
searchLiveKnowledge: { execute: mocks.liveSearch },
52+
}))
53+
54+
import { createKnowledgeMcpHandlers } from '@/lib/knowledge/mcp/route-handler'
55+
56+
urlsMockFns.mockGetBaseUrl.mockReturnValue('http://localhost')
57+
58+
const handlers = createKnowledgeMcpHandlers()
59+
const principal = createPersonalApiKeyPrincipal({ userId: 'person-1' })
60+
61+
/** Calls a tool over the real route, SDK transport, and Search MCP server. */
62+
async function callTool(name: string, args: Record<string, unknown>) {
63+
const response = await handlers.POST(
64+
createMockRequest({
65+
method: 'POST',
66+
url: 'http://localhost/api/mcp/search/organizations/org-1',
67+
headers: {
68+
'content-type': 'application/json',
69+
accept: 'application/json, text/event-stream',
70+
'x-forwarded-for': '127.0.0.1',
71+
'x-api-key': 'personal-key',
72+
},
73+
body: {
74+
jsonrpc: '2.0',
75+
id: 1,
76+
method: 'tools/call',
77+
params: { name, arguments: args },
78+
},
79+
}),
80+
createRouteContext({ organizationId: 'org-1' })
81+
)
82+
return { status: response.status, body: await response.json() }
83+
}
84+
85+
function callSearch(args: Record<string, unknown>) {
86+
return callTool('search', args)
87+
}
88+
89+
function searchedInput() {
90+
const call = mocks.liveSearch.mock.lastCall?.[0]
91+
if (!call) throw new Error('Expected the search use case to run')
92+
return call.input
93+
}
94+
95+
beforeEach(() => {
96+
resetDbChainMock()
97+
v2RouteMocks.authenticate.mockResolvedValue({
98+
principal,
99+
keyType: 'personal',
100+
keyExpiresAt: null,
101+
rateLimitSubjectIds: ['person-1'],
102+
rateLimitSubscription: null,
103+
})
104+
permissionGroupsResolveMockFns.mockGetUserPermissionConfigForOrganization.mockResolvedValue(null)
105+
knowledgeAvailabilityMockFns.mockRequireOrganizationSearchAvailable.mockResolvedValue(undefined)
106+
mocks.index.mockResolvedValue({ id: 'index-1' })
107+
mocks.liveSearch.mockResolvedValue({
108+
results: [],
109+
retrieval: { status: 'complete', timedOutLegs: [] },
110+
live: { backend: 'live', accounts: [], guidance: '' },
111+
})
112+
queueTableRows(member, [{ role: 'member' }])
113+
})
114+
115+
describe('Search MCP tool arguments with one possible reading', () => {
116+
it('searches without native queries when the client sends an empty nativeQueries list', async () => {
117+
const { body } = await callSearch({ query: 'release notes', nativeQueries: [] })
118+
expect(body.result?.isError).not.toBe(true)
119+
expect(searchedInput().nativeQueries).toBeUndefined()
120+
})
121+
122+
it('treats a null optional native query field as omitted', async () => {
123+
const { body } = await callSearch({
124+
query: 'release',
125+
nativeQueries: [{ provider: 'jira', query: 'text ~ "release"', kind: null }],
126+
})
127+
expect(body.result?.isError).not.toBe(true)
128+
expect(searchedInput().nativeQueries).toEqual([{ provider: 'jira', query: 'text ~ "release"' }])
129+
})
130+
131+
it('applies filters every native query carries alike as the search filters', async () => {
132+
const { body } = await callSearch({
133+
query: 'release',
134+
nativeQueries: [
135+
{
136+
provider: 'gmail',
137+
query: 'subject:release',
138+
startDate: '2026-10-01T00:00:00-07:00',
139+
sortBy: 'newest',
140+
},
141+
{
142+
provider: 'slack',
143+
query: 'release',
144+
startDate: '2026-10-01T00:00:00-07:00',
145+
sortBy: 'newest',
146+
},
147+
],
148+
})
149+
expect(body.result?.isError).not.toBe(true)
150+
const input = searchedInput()
151+
expect(input.filters).toMatchObject({
152+
startDate: '2026-10-01T00:00:00-07:00',
153+
sortBy: 'newest',
154+
})
155+
expect(input.nativeQueries).toEqual([
156+
{ provider: 'gmail', query: 'subject:release' },
157+
{ provider: 'slack', query: 'release' },
158+
])
159+
})
160+
161+
it('reads a top-level provider as the source filter', async () => {
162+
const { body } = await callSearch({ query: 'standup', provider: 'google_calendar' })
163+
expect(body.result?.isError).not.toBe(true)
164+
expect(searchedInput().filters).toMatchObject({ source: 'google_calendar' })
165+
})
166+
167+
it.each([
168+
['carry different values', '2026-09-01T00:00:00Z'],
169+
['do not all carry it', undefined],
170+
])('still rejects native queries whose dates %s, rather than choosing one', async (_, other) => {
171+
const { body } = await callSearch({
172+
query: 'release',
173+
nativeQueries: [
174+
{ provider: 'gmail', query: 'release', startDate: '2026-10-01T00:00:00Z' },
175+
{ provider: 'slack', query: 'release', ...(other ? { startDate: other } : {}) },
176+
],
177+
})
178+
expect(JSON.stringify(body)).toContain('startDate')
179+
expect(mocks.liveSearch).not.toHaveBeenCalled()
180+
})
181+
182+
it('still rejects a top-level provider that names no provider', async () => {
183+
const { body } = await callSearch({ query: 'standup', provider: 'Google Calendar' })
184+
expect(JSON.stringify(body)).toContain('provider')
185+
expect(mocks.liveSearch).not.toHaveBeenCalled()
186+
})
187+
188+
it.each(['constructor', 'toString', '__proto__'])(
189+
'answers a call to the unknown tool %s with an MCP error, not a server error',
190+
async (name) => {
191+
const { status, body } = await callTool(name, {})
192+
expect(status).toBe(200)
193+
expect(JSON.stringify(body)).toContain(name)
194+
}
195+
)
196+
})

‎apps/sim/lib/knowledge/mcp/route-handler.ts‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { createLogger } from '@sim/logger'
12
import type { NextRequest } from 'next/server'
23
import { organizationKnowledgeMcpContract } from '@/lib/api/contracts/knowledge/mcp'
34
import { parseRequest } from '@/lib/api/server'
@@ -13,9 +14,12 @@ import { knowledgeOperations } from '@/lib/knowledge/application/operations'
1314
import { readSearchIndex } from '@/lib/knowledge/application/sim-search'
1415
import { withSearchMcpAuthChallenge } from '@/lib/knowledge/mcp/oauth-metadata'
1516
import { createKnowledgeMcpServer } from '@/lib/knowledge/mcp/server'
17+
import { repairToolCallMessage, toolCallInputIssues } from '@/lib/knowledge/mcp/tool-arguments'
1618
import { getSearchMcpUrl } from '@/lib/knowledge/mcp/urls'
1719
import { v2CaughtOrchestrationError, v2Error } from '@/app/api/v2/lib/response'
1820

21+
const logger = createLogger('KnowledgeMcp')
22+
1923
function mcpAuth(resource: string) {
2024
return mcpCredentialAuth({ resource, allowUnboundApiTokens: true })
2125
}
@@ -47,12 +51,19 @@ export function createKnowledgeMcpHandlers() {
4751
input: parsed.data.params,
4852
request,
4953
})
54+
const { organizationId } = parsed.data.params
55+
const { message, repairs } = repairToolCallMessage(parsed.data.body)
56+
if (repairs.length)
57+
logger.info('Knowledge MCP tool input repaired', { organizationId, repairs })
58+
const rejected = toolCallInputIssues(message)
59+
if (rejected)
60+
logger.warn('Knowledge MCP tool input rejected', { organizationId, ...rejected })
5061
const server = createKnowledgeMcpServer({
5162
request,
5263
auth: admission.auth,
5364
...parsed.data.params,
5465
})
55-
return await serveStatelessMcp(server, request, parsed.data.body)
66+
return await serveStatelessMcp(server, request, message)
5667
} catch (error) {
5768
const response = v2CaughtOrchestrationError(error)
5869
if (response) return withSearchMcpAuthChallenge(response, resource)
Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
import { isRecordLike } from '@sim/utils/object'
2+
import type { z } from 'zod'
3+
import {
4+
chatSearchMcpSchema,
5+
liveSearchMcpSchema,
6+
readLiveDocumentMcpSchema,
7+
} from '@/lib/api/contracts/knowledge/mcp'
8+
import { mcpToolCallParamsSchema } from '@/lib/api/contracts/mcp'
9+
import { liveSearchProviderSchema } from '@/lib/api/contracts/mothership-assistant-tools'
10+
11+
/** Search filters a client sometimes writes on each native query instead of on the search. */
12+
const SEARCH_FILTERS_ON_NATIVE_QUERIES = [
13+
'startDate',
14+
'endDate',
15+
'modifiedAfter',
16+
'modifiedBefore',
17+
'sortBy',
18+
'topK',
19+
] as const
20+
21+
/** Keyed by client-supplied tool names, so a name such as `constructor` finds nothing. */
22+
const TOOL_INPUT_SCHEMAS = new Map<string, z.ZodType>([
23+
['search', liveSearchMcpSchema],
24+
['read_document', readLiveDocumentMcpSchema],
25+
['chat', chatSearchMcpSchema],
26+
])
27+
28+
function readToolCall(message: Record<string, unknown>) {
29+
if (message.method !== 'tools/call') return undefined
30+
const params = mcpToolCallParamsSchema.safeParse(message.params)
31+
return params.success ? params.data : undefined
32+
}
33+
34+
/**
35+
* Search arguments rewritten where the schema rejects them but they have exactly one reading: an
36+
* empty `nativeQueries` list, a null optional native query field, a top-level `provider` naming a
37+
* provider (the `source` filter), and a search filter every native query carries with the same value. Native
38+
* queries that disagree on a filter are left for the schema to reject.
39+
*/
40+
function repairSearchArguments(args: Record<string, unknown>) {
41+
const { nativeQueries: queries, provider, ...rest } = args
42+
const repaired: Record<string, unknown> = rest
43+
const repairs = new Set<string>()
44+
const providerId = liveSearchProviderSchema.safeParse(provider)
45+
if (providerId.success && (rest.source === undefined || rest.source === providerId.data)) {
46+
repaired.source = providerId.data
47+
repairs.add('providerAsSource')
48+
} else if (provider !== undefined) repaired.provider = provider
49+
if (Array.isArray(queries) && queries.length === 0) repairs.add('emptyNativeQueries')
50+
else if (queries !== undefined) repaired.nativeQueries = queries
51+
if (!Array.isArray(queries) || !queries.length || !queries.every(isRecordLike))
52+
return { args: repaired, repairs }
53+
let items = queries.map((query) =>
54+
Object.fromEntries(
55+
Object.entries(query).filter(([key, value]) => {
56+
if (value !== null || key === 'provider' || key === 'query') return true
57+
repairs.add('nullNativeQueryField')
58+
return false
59+
})
60+
)
61+
)
62+
for (const filter of SEARCH_FILTERS_ON_NATIVE_QUERIES) {
63+
if (!items.some((item) => filter in item)) continue
64+
/** A query without the filter adds `undefined`, so a filter not on every query differs too. */
65+
const values = new Set(items.map((item) => JSON.stringify(item[filter])))
66+
const [value] = values
67+
if (values.size !== 1 || value === undefined) continue
68+
if (filter in repaired && JSON.stringify(repaired[filter]) !== value) continue
69+
repaired[filter] = JSON.parse(value)
70+
items = items.map(({ [filter]: _, ...rest }) => rest)
71+
repairs.add('nativeQueryFilterAsSearchFilter')
72+
}
73+
repaired.nativeQueries = items
74+
return { args: repaired, repairs }
75+
}
76+
77+
/**
78+
* The JSON-RPC message with a search call's arguments repaired before the SDK validates them, so
79+
* a call the client clearly meant runs instead of costing it another round. Arguments the schema
80+
* already accepts, other tools, and tools/list come back unchanged.
81+
*/
82+
export function repairToolCallMessage(message: Record<string, unknown>) {
83+
const call = readToolCall(message)
84+
if (call?.name !== 'search' || !call.arguments) return { message, repairs: [] }
85+
const { args, repairs } = repairSearchArguments(call.arguments)
86+
if (!repairs.size) return { message, repairs: [] }
87+
return { message: { ...message, params: { ...call, arguments: args } }, repairs: [...repairs] }
88+
}
89+
90+
/**
91+
* Where a tool call's arguments fail its input schema, as field paths and issue codes without
92+
* values. The SDK rejects such a call before the tool runs, so nothing else records it.
93+
*/
94+
export function toolCallInputIssues(message: Record<string, unknown>) {
95+
const call = readToolCall(message)
96+
const schema = call ? TOOL_INPUT_SCHEMAS.get(call.name) : undefined
97+
if (!call || !schema) return undefined
98+
const parsed = schema.safeParse(call.arguments ?? {})
99+
if (parsed.success) return undefined
100+
return {
101+
toolName: call.name,
102+
issues: parsed.error.issues.map(({ code, path }) => ({ code, path: path.join('.') })),
103+
}
104+
}

0 commit comments

Comments
 (0)