|
| 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 | +}) |
0 commit comments