import { expect } from 'chai'; import { beforeEach, describe, it } from 'mocha'; import proxyquire from 'proxyquire'; import sinon from 'sinon'; const License = { hasModule: sinon.stub(), }; const cachedSettings = { get: sinon.stub(), }; const Messages = { findVisibleByIds: sinon.stub(), }; const Rooms = { findByIds: sinon.stub(), findOneByNameOrFname: sinon.stub(), }; const Subscriptions = { findByUserId: sinon.stub(), findByUserIdAndRoomIds: sinon.stub(), }; const Users = { findOneById: sinon.stub(), }; const serverFetch = sinon.stub(); const { AISearchService } = proxyquire.noCallThru().load('../../../../../server/services/ai-search/service', { '@rocket.chat/core-services': { License, ServiceClass: class { protected name = ''; }, }, '../../settings': { settings: cachedSettings, }, '@rocket.chat/logger': { Logger: class { debug = sinon.stub(); warn = sinon.stub(); }, }, '@rocket.chat/models': { Messages, Rooms, Subscriptions, Users, }, '@rocket.chat/server-fetch': { serverFetch, }, }); type CursorResult = { toArray(): Promise; map(callback: (item: T) => U): CursorResult; }; const cursor = (items: T[]): CursorResult => ({ toArray: async () => items, map: (callback) => cursor(items.map(callback)), }); const settings: Record = { AI_Intelligent_Search_Enabled: true, AI_Intelligent_Search_Pipeline_Base_URL: 'https://pipeline.example.com', AI_Intelligent_Search_Pipeline_ID: 'workspace', AI_Intelligent_Search_API_Key: 'key', AI_Intelligent_Search_API_Key_Secret: 'secret', AI_Intelligent_Search_Query_Template: '', AI_Intelligent_Search_Min_Similarity_Percent: 61, AI_Intelligent_Search_Answer_Enabled: true, AI_LLM_OpenAI_Base_URL: 'https://llm.example.com', AI_LLM_OpenAI_API_Key: 'llm-key', AI_LLM_OpenAI_Model: 'gpt-test', AI_Intelligent_Search_Answer_System_Prompt: 'Use sources only.', SSRF_Allowlist: 'pipeline.example.com,llm.example.com', }; const createService = (): InstanceType => new AISearchService(); describe('AISearchService', () => { beforeEach(() => { License.hasModule.reset(); cachedSettings.get.reset(); Messages.findVisibleByIds.reset(); Rooms.findByIds.reset(); Rooms.findOneByNameOrFname.reset(); Subscriptions.findByUserId.reset(); Subscriptions.findByUserIdAndRoomIds.reset(); Users.findOneById.reset(); serverFetch.reset(); License.hasModule.resolves(true); cachedSettings.get.callsFake((key: string) => settings[key]); Users.findOneById.resolves({ roles: ['admin'] }); Rooms.findByIds.callsFake((roomIds: string[]) => cursor( roomIds.map((roomId) => ({ _id: roomId, t: 'c', name: roomId === 'allowed' ? 'general' : roomId, fname: roomId === 'allowed' ? 'General' : roomId, })), ), ); Subscriptions.findByUserIdAndRoomIds.callsFake((_userId: string, roomIds: string[]) => cursor(roomIds.filter((roomId) => roomId === 'allowed' || roomId === 'room-general').map((rid) => ({ rid }))), ); Messages.findVisibleByIds.callsFake((msgIds: string[]) => cursor( msgIds.map((msgId) => ({ _id: msgId, rid: msgId === 'blocked-msg' ? 'blocked' : 'allowed', msg: `${msgId} from db`, ts: new Date('2026-01-05T12:00:00.000Z'), u: { username: 'alice', name: 'Alice' }, })), ), ); }); describe('status', () => { it('reports availability from license, settings, pipeline, and LLM configuration', async () => { expect(await createService().status()).to.deep.equal({ hasIntelligentSearchLicense: true, intelligentSearchEnabled: true, intelligentSearchConfigured: true, answerGenerationConfigured: true, }); }); it('marks answer generation unavailable when the answer setting is off', async () => { cachedSettings.get.callsFake((key: string) => (key === 'AI_Intelligent_Search_Answer_Enabled' ? false : settings[key])); expect(await createService().status()).to.include({ answerGenerationConfigured: false, }); }); it('marks answer generation unavailable when the pipeline is not configured', async () => { cachedSettings.get.callsFake((key: string) => (key === 'AI_Intelligent_Search_Pipeline_Base_URL' ? '' : settings[key])); expect(await createService().status()).to.include({ intelligentSearchConfigured: false, answerGenerationConfigured: false, }); }); }); describe('search', () => { it('does not call the pipeline when license, setting, or configuration is unavailable', async () => { License.hasModule.resolves(false); expect(await createService().search({ query: 'fruit', userId: 'user-id' })).to.deep.equal([]); expect(serverFetch.called).to.be.false; }); it('sends subscribed room scope to the pipeline for broad searches', async () => { const subscribedRoomIds = [...Array.from({ length: 1001 }, (_, index) => `room-${index}`), 'allowed']; Subscriptions.findByUserId.returns(cursor([...subscribedRoomIds.map((rid) => ({ rid })), { rid: 'banned-room', status: 'BANNED' }])); serverFetch.resolves({ ok: true, status: 200, json: async () => ({ results: [ { metadata: { room_id: 'blocked', msg_id: 'blocked-msg' }, text: 'blocked pipeline text', score: 0.1 }, { metadata: { room_id: 'allowed', msg_id: 'allowed-msg' }, text: 'allowed pipeline text', score: 0.39 }, ], }), text: async () => '', }); const results = await createService().search({ query: 'fruit', userId: 'user-id', limit: 5 }); expect(results).to.deep.equal([ { _id: 'allowed-msg', rid: 'allowed', msgId: 'allowed-msg', text: 'allowed-msg from db', ts: '2026-01-05T12:00:00.000Z', u: { username: 'alice', name: 'Alice' }, score: 0.61, room: { _id: 'allowed', t: 'c', name: 'general', fname: 'General' }, }, ]); const [, options] = serverFetch.firstCall.args; const body = JSON.parse(options.body); expect(body.params.k).to.equal(5); expect(body.filters).to.deep.equal({ room_id: { $in: subscribedRoomIds }, }); expect(options).to.include({ ignoreSsrfValidation: false, allowList: 'pipeline.example.com,llm.example.com', }); expect( Subscriptions.findByUserId.calledWith('user-id', { projection: { rid: 1, status: 1 }, }), ).to.be.true; expect( Subscriptions.findByUserIdAndRoomIds.calledWith('user-id', ['blocked', 'allowed'], { projection: { rid: 1, status: 1 }, }), ).to.be.true; }); it('resolves room-name filters before querying the pipeline', async () => { Rooms.findOneByNameOrFname.resolves({ _id: 'room-general' }); Subscriptions.findByUserId.returns(cursor([{ rid: 'room-general' }])); serverFetch.resolves({ ok: true, status: 200, json: async () => ({ results: [ { metadata: { room_id: 'room-general', msg_id: 'general-msg' }, text: 'general pipeline text', similarity: 0.8 }, { metadata: { room_id: 'blocked', msg_id: 'blocked-msg' }, text: 'blocked pipeline text', similarity: 0.99 }, ], }), text: async () => '', }); Messages.findVisibleByIds.returns( cursor([ { _id: 'general-msg', rid: 'room-general', msg: 'general message from db', ts: new Date('2026-01-05T12:00:00.000Z'), u: { username: 'alice', name: 'Alice' }, }, ]), ); Rooms.findByIds.returns(cursor([{ _id: 'room-general', t: 'c', name: 'general', fname: 'General' }])); const results = await createService().search({ query: 'fruit', userId: 'user-id', filters: { roomNames: ['general'], fromUsernames: ['alice'], startDate: '2026-01-01T00:00:00.000Z', endDate: '2026-01-31T00:00:00.000Z', }, limit: 5, }); expect(results).to.have.lengthOf(1); expect(results[0]).to.include({ _id: 'general-msg', rid: 'room-general', score: 0.8 }); expect( Messages.findVisibleByIds.calledWith(['general-msg', 'blocked-msg'], { projection: { _id: 1, rid: 1, msg: 1, ts: 1, u: 1 }, }), ).to.be.true; const [, options] = serverFetch.firstCall.args; expect(JSON.parse(options.body).filters).to.deep.equal({ room_id: { $eq: 'room-general' }, username: { $eq: 'alice' }, timestamp: { $ge: '2026-01-01T00:00:00.000Z', $le: '2026-01-31T00:00:00.000Z', }, }); }); it('returns no results when room-name filters do not resolve to subscribed rooms', async () => { Rooms.findOneByNameOrFname.resolves(null); expect( await createService().search({ query: 'fruit', userId: 'user-id', filters: { roomNames: ['unknown-room'] }, limit: 5 }), ).to.deep.equal([]); expect(serverFetch.called).to.be.false; }); it('drops pipeline hits that do not reference a message', async () => { Subscriptions.findByUserId.returns(cursor([{ rid: 'allowed' }])); serverFetch.resolves({ ok: true, status: 200, json: async () => ({ results: [{ metadata: { room_id: 'allowed' }, text: 'index text without a message reference', score: 0.1 }], }), text: async () => '', }); expect(await createService().search({ query: 'fruit', userId: 'user-id', limit: 5 })).to.deep.equal([]); }); it('drops pipeline hits when the referenced message is not visible', async () => { Subscriptions.findByUserId.returns(cursor([{ rid: 'allowed' }])); serverFetch.resolves({ ok: true, status: 200, json: async () => ({ results: [{ metadata: { room_id: 'allowed', msg_id: 'hidden-msg' }, text: 'pipeline text should not leak', score: 0.1 }], }), text: async () => '', }); Messages.findVisibleByIds.returns(cursor([])); expect(await createService().search({ query: 'fruit', userId: 'user-id', limit: 5 })).to.deep.equal([]); }); it('drops pipeline hits from rooms where the user subscription is banned', async () => { Subscriptions.findByUserId.returns(cursor([{ rid: 'allowed' }])); Subscriptions.findByUserIdAndRoomIds.returns(cursor([{ rid: 'allowed', status: 'BANNED' }])); serverFetch.resolves({ ok: true, status: 200, json: async () => ({ results: [{ metadata: { room_id: 'allowed', msg_id: 'allowed-msg' }, text: 'pipeline text', similarity: 0.8 }], }), text: async () => '', }); expect(await createService().search({ query: 'fruit', userId: 'user-id', limit: 5 })).to.deep.equal([]); }); }); describe('answer', () => { it('rejects answer generation when AI Search or answer generation is unavailable', async () => { cachedSettings.get.callsFake((key: string) => (key === 'AI_Intelligent_Search_Answer_Enabled' ? false : settings[key])); try { await createService().answer({ query: 'fruit', messages: [{ text: 'oranges are green' }] }); throw new Error('Expected answer generation to fail'); } catch (error) { expect((error as Error).message).to.equal('error-ai-not-enabled'); } expect(serverFetch.called).to.be.false; }); it('generates an answer from source messages with the configured LLM provider', async () => { serverFetch.resolves({ ok: true, status: 200, json: async () => ({ choices: [{ message: { content: 'Oranges are green.' } }] }), text: async () => '', }); expect( await createService().answer({ query: 'fruit colors', messages: [{ text: 'oranges are green', username: 'alice', roomName: 'general', score: 0.61 }], }), ).to.deep.equal({ answer: 'Oranges are green.', provider: { name: 'OpenAI compatible', model: 'gpt-test' }, }); const [url, options] = serverFetch.firstCall.args; expect(url).to.equal('https://llm.example.com/chat/completions'); expect(options.headers.Authorization).to.equal('Bearer llm-key'); expect(JSON.parse(options.body).messages[0]).to.deep.equal({ role: 'system', content: 'Use sources only.' }); }); }); describe('models', () => { it('does not contact the provider without the AI add-on', async () => { License.hasModule.resolves(false); expect(await createService().models()).to.deep.equal([]); expect(serverFetch.called).to.be.false; }); }); });