Skip to content

Commit b5d32a2

Browse files
committed
improvement(search): run Search MCP calls whose arguments have one reading, and log rejected ones
1 parent 343a5cc commit b5d32a2

3 files changed

Lines changed: 300 additions & 1 deletion

File tree

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

‎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: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
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+
10+
/** Search filters a client sometimes writes on each native query instead of on the search. */
11+
const SEARCH_FILTERS_ON_NATIVE_QUERIES = [
12+
'startDate',
13+
'endDate',
14+
'modifiedAfter',
15+
'modifiedBefore',
16+
'sortBy',
17+
'topK',
18+
] as const
19+
20+
const TOOL_INPUT_SCHEMAS: Record<string, z.ZodType> = {
21+
search: liveSearchMcpSchema,
22+
read_document: readLiveDocumentMcpSchema,
23+
chat: chatSearchMcpSchema,
24+
}
25+
26+
function readToolCall(message: Record<string, unknown>) {
27+
if (message.method !== 'tools/call') return undefined
28+
const params = mcpToolCallParamsSchema.safeParse(message.params)
29+
return params.success ? params.data : undefined
30+
}
31+
32+
/**
33+
* Search arguments rewritten where the schema rejects them but they have exactly one reading: an
34+
* empty `nativeQueries` list, a null optional native query field, a top-level `provider` (the
35+
* `source` filter), and a search filter every native query carries with the same value. Native
36+
* queries that disagree on a filter are left for the schema to reject.
37+
*/
38+
function repairSearchArguments(args: Record<string, unknown>) {
39+
const { nativeQueries: queries, provider, ...rest } = args
40+
const repaired: Record<string, unknown> = rest
41+
const repairs = new Set<string>()
42+
if (typeof provider === 'string' && (rest.source === undefined || rest.source === provider)) {
43+
repaired.source = provider
44+
repairs.add('providerAsSource')
45+
} else if (provider !== undefined) repaired.provider = provider
46+
if (Array.isArray(queries) && queries.length === 0) repairs.add('emptyNativeQueries')
47+
else if (queries !== undefined) repaired.nativeQueries = queries
48+
if (!Array.isArray(queries) || !queries.length || !queries.every(isRecordLike))
49+
return { args: repaired, repairs }
50+
let items = queries.map((query) =>
51+
Object.fromEntries(
52+
Object.entries(query).filter(([key, value]) => {
53+
if (value !== null || key === 'provider' || key === 'query') return true
54+
repairs.add('nullNativeQueryField')
55+
return false
56+
})
57+
)
58+
)
59+
for (const filter of SEARCH_FILTERS_ON_NATIVE_QUERIES) {
60+
if (!items.some((item) => filter in item)) continue
61+
/** A query without the filter adds `undefined`, so a filter not on every query differs too. */
62+
const values = new Set(items.map((item) => JSON.stringify(item[filter])))
63+
const [value] = values
64+
if (values.size !== 1 || value === undefined) continue
65+
if (filter in repaired && JSON.stringify(repaired[filter]) !== value) continue
66+
repaired[filter] = JSON.parse(value)
67+
items = items.map(({ [filter]: _, ...rest }) => rest)
68+
repairs.add('nativeQueryFilterAsSearchFilter')
69+
}
70+
repaired.nativeQueries = items
71+
return { args: repaired, repairs }
72+
}
73+
74+
/**
75+
* The JSON-RPC message with a search call's arguments repaired before the SDK validates them, so
76+
* a call the client clearly meant runs instead of costing it another round. Arguments the schema
77+
* already accepts, other tools, and tools/list come back unchanged.
78+
*/
79+
export function repairToolCallMessage(message: Record<string, unknown>) {
80+
const call = readToolCall(message)
81+
if (call?.name !== 'search' || !call.arguments) return { message, repairs: [] }
82+
const { args, repairs } = repairSearchArguments(call.arguments)
83+
if (!repairs.size) return { message, repairs: [] }
84+
return { message: { ...message, params: { ...call, arguments: args } }, repairs: [...repairs] }
85+
}
86+
87+
/**
88+
* Where a tool call's arguments fail its input schema, as field paths and issue codes without
89+
* values. The SDK rejects such a call before the tool runs, so nothing else records it.
90+
*/
91+
export function toolCallInputIssues(message: Record<string, unknown>) {
92+
const call = readToolCall(message)
93+
const schema = call ? TOOL_INPUT_SCHEMAS[call.name] : undefined
94+
if (!call || !schema) return undefined
95+
const parsed = schema.safeParse(call.arguments ?? {})
96+
if (parsed.success) return undefined
97+
return {
98+
toolName: call.name,
99+
issues: parsed.error.issues.map(({ code, path }) => ({ code, path: path.join('.') })),
100+
}
101+
}

0 commit comments

Comments
 (0)