import { Logger } from '@nestjs/common'; import { AiChatService, AiChatRunHooks } from './ai-chat.service'; import { AiChatRunService } from './ai-chat-run.service'; import type { User, Workspace } from '@docmost/db/types/entity.types'; /** * Lifecycle unit tests for AiChatService.onModuleInit (#183 crash-recovery * sweep). The sweep is BEST-EFFORT: a failure must be logged (warn) but must * NEVER throw out of onModuleInit and block server startup. Exercised with a * hand-rolled mock repo — no Nest graph, no DB. Only `aiChatMessageRepo` is * touched by onModuleInit, so the other constructor deps are stubbed as never. */ describe('AiChatService.onModuleInit (startup sweep)', () => { function makeService(sweepStreaming: jest.Mock) { const aiChatMessageRepo = { sweepStreaming }; const service = new AiChatService( {} as never, // ai {} as never, // aiChatRepo aiChatMessageRepo as never, {} as never, // aiChatPageSnapshotRepo {} as never, // aiSettings {} as never, // tools {} as never, // mcpClients {} as never, // aiAgentRoleRepo {} as never, // pageRepo {} as never, // pageAccess {} as never, // environment ); return { service, aiChatMessageRepo }; } afterEach(() => jest.restoreAllMocks()); it('happy path: calls sweepStreaming and resolves', async () => { const sweepStreaming = jest.fn().mockResolvedValue(0); const { service } = makeService(sweepStreaming); await expect(service.onModuleInit()).resolves.toBeUndefined(); expect(sweepStreaming).toHaveBeenCalledTimes(1); }); it('logs how many rows were swept when > 0', async () => { const sweepStreaming = jest.fn().mockResolvedValue(3); const logSpy = jest .spyOn(Logger.prototype, 'log') .mockImplementation(() => undefined); const { service } = makeService(sweepStreaming); await service.onModuleInit(); expect(logSpy).toHaveBeenCalledTimes(1); expect(String(logSpy.mock.calls[0][0])).toContain('3'); }); it('sweepStreaming throws -> onModuleInit resolves (does NOT throw) and warns', async () => { const sweepStreaming = jest .fn() .mockRejectedValue(new Error('db unavailable')); const warnSpy = jest .spyOn(Logger.prototype, 'warn') .mockImplementation(() => undefined); const { service } = makeService(sweepStreaming); // Must not throw — a sweep failure may never block startup. await expect(service.onModuleInit()).resolves.toBeUndefined(); expect(warnSpy).toHaveBeenCalledTimes(1); expect(String(warnSpy.mock.calls[0][0])).toContain('db unavailable'); }); }); /** * #184 CRITICAL run-lifecycle safety net (review fix). A transient failure * AFTER a successful beginRun but BEFORE streamText's terminal callbacks own the * lifecycle must STILL settle the run — otherwise the run row is stuck 'running' * forever (sweepRunning only runs at startup) and the partial unique index + the * controller pre-check 409 every future turn in that chat until a restart. Here * we model the very first bare await after beginRun (the user-message insert) * throwing, wiring the run hooks to a REAL AiChatRunService (mock repo) exactly * as the controller does, and assert the run is settled to 'error' and its * in-memory entry dropped (so a follow-up turn would NOT be 409'd). */ describe('AiChatService.stream run-lifecycle safety net (#184)', () => { const user = { id: 'u1' } as User; const workspace = { id: 'ws1' } as Workspace; afterEach(() => jest.restoreAllMocks()); it('an exception after beginRun settles the run to error and drops the in-memory entry', async () => { jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined); // Real run service over a mock repo, so finalizeRun's in-memory bookkeeping // (active.delete) is exercised for real. const runRepo = { insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }), update: jest.fn().mockResolvedValue({ id: 'run-1' }), }; const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never); // The user-message insert (the first bare await after beginRun) throws. const aiChatMessageRepo = { insert: jest.fn().mockRejectedValue(new Error('insert boom')), }; const aiChatRepo = { // Existing chat -> chatId stays, no new-chat insert path. findById: jest.fn().mockResolvedValue({ id: 'chat-1', creatorId: 'u1' }), }; const service = new AiChatService( {} as never, // ai aiChatRepo as never, aiChatMessageRepo as never, {} as never, // aiChatPageSnapshotRepo {} as never, // aiSettings {} as never, // tools {} as never, // mcpClients {} as never, // aiAgentRoleRepo {} as never, // pageRepo {} as never, // pageAccess {} as never, // environment ); const runHooks: AiChatRunHooks = { begin: (chatId) => runService.beginRun({ chatId, workspaceId: workspace.id, userId: user.id, trigger: 'user', }), onSettled: (runId, status, error) => runService.finalizeRun(runId, workspace.id, status, error), }; await expect( service.stream({ user, workspace, sessionId: 'sess', body: { chatId: 'chat-1', messages: [ { id: 'm', role: 'user', parts: [{ type: 'text', text: 'hi' }] }, ], }, res: {} as never, signal: new AbortController().signal, model: {} as never, role: null, runHooks, }), ).rejects.toThrow('insert boom'); // The run was begun... expect(runRepo.insert).toHaveBeenCalledTimes(1); // ...then settled to a terminal FAILED status by the safety net... expect(runRepo.update).toHaveBeenCalledTimes(1); expect(runRepo.update).toHaveBeenCalledWith( 'run-1', 'ws1', expect.objectContaining({ status: 'failed' }), ); // ...and the in-memory entry is gone, so a follow-up turn is NOT 409'd. expect(runService.isLocallyActive('run-1')).toBe(false); }); });