import { afterEach, describe, expect, it, vi } from 'vitest'; import { MindDB, type EmbeddingProviderInstance, type EmbeddingProviderStatus, } from '@waggle/core'; import { VectorEnrichmentService, type VectorEnrichmentRunResult, } from '../../src/local/services/vector-enrichment-service.js'; import type { VectorBackfillResult } from '../../src/local/vector-backfill.js'; function backfillResult(overrides: Partial = {}): VectorBackfillResult { return { skipped: null, vectorsRepaired: false, framesReembedded: 0, chunksCreated: 0, framesProcessed: 0, hasMore: false, errors: [], ...overrides, }; } function deferred() { let resolve!: (value: T) => void; const promise = new Promise((res) => { resolve = res; }); return { promise, resolve }; } const provider = {} as EmbeddingProviderInstance; describe('VectorEnrichmentService', () => { const dbs: MindDB[] = []; afterEach(() => { vi.useRealTimers(); for (const db of dbs.splice(0)) db.close(); }); it('runs immediately, every five minutes, and coalesces overlapping requests', async () => { vi.useFakeTimers(); const personal = new MindDB(':memory:'); dbs.push(personal); const firstPass = deferred(); const runPass = vi.fn() .mockImplementationOnce(() => firstPass.promise) .mockResolvedValue(backfillResult()); const service = new VectorEnrichmentService({ personalMind: personal, embeddingProvider: provider, listWorkspaceIds: () => [], acquireWorkspaceMind: () => { throw new Error('not used'); }, releaseWorkspaceMind: () => undefined, runPass, }, { maxPassesPerMind: 2 }); service.start(); await vi.advanceTimersByTimeAsync(0); const overlapping = service.runNow(); expect(runPass).toHaveBeenCalledTimes(1); firstPass.resolve(backfillResult({ framesProcessed: 1, hasMore: true })); await overlapping; expect(runPass).toHaveBeenCalledTimes(2); // bounded second pass await vi.advanceTimersByTimeAsync(5 * 60 * 1000); expect(runPass).toHaveBeenCalledTimes(3); await service.stop(); }); it('pins each workspace for all of its passes and releases it even after a failure', async () => { const personal = new MindDB(':memory:'); const ws1 = new MindDB(':memory:'); const ws2 = new MindDB(':memory:'); dbs.push(personal, ws1, ws2); const acquire = vi.fn((id: string) => id === 'ws-1' ? ws1 : ws2); const release = vi.fn(); const perMindCalls = new Map(); const runPass = vi.fn(async (db: MindDB) => { perMindCalls.set(db, (perMindCalls.get(db) ?? 0) + 1); if (db === ws2) throw new Error('workspace pass failed'); return backfillResult({ hasMore: (perMindCalls.get(db) ?? 0) === 1 }); }); const service = new VectorEnrichmentService({ personalMind: personal, embeddingProvider: provider, listWorkspaceIds: () => ['ws-1', 'ws-2'], acquireWorkspaceMind: acquire, releaseWorkspaceMind: release, runPass, }, { maxPassesPerMind: 2 }); const result = await service.runNow(); expect(acquire.mock.calls).toEqual([['ws-1'], ['ws-2']]); expect(release.mock.calls).toEqual([['ws-1'], ['ws-2']]); expect(perMindCalls.get(ws1)).toBe(2); expect(result.errors).toEqual([expect.stringContaining('workspace ws-2')]); }); it('waits for an in-flight pass before shutdown resolves', async () => { const personal = new MindDB(':memory:'); dbs.push(personal); const pass = deferred(); const service = new VectorEnrichmentService({ personalMind: personal, embeddingProvider: provider, listWorkspaceIds: () => [], acquireWorkspaceMind: () => { throw new Error('not used'); }, releaseWorkspaceMind: () => undefined, runPass: () => pass.promise, }); const run = service.runNow(); let stopped = false; const stopping = service.stop().then(() => { stopped = true; }); await Promise.resolve(); expect(stopped).toBe(false); pass.resolve(backfillResult()); await expect(stopping).resolves.toBeUndefined(); await expect(run).resolves.toMatchObject>({ mindsVisited: 1 }); }); it('does not let a stalled provider reprobe block shutdown for an empty mind', async () => { const personal = new MindDB(':memory:'); dbs.push(personal); const status: EmbeddingProviderStatus = { activeProvider: 'mock', availableProviders: ['mock'], dimensions: 1024, modelName: 'deterministic-mock', probeTimestamp: new Date(0).toISOString(), }; const reprobe = vi.fn(() => new Promise(() => undefined)); const stalledProvider = { dimensions: 1024, getActiveProvider: () => 'mock', getStatus: () => status, reprobe, } as unknown as EmbeddingProviderInstance; const service = new VectorEnrichmentService({ personalMind: personal, embeddingProvider: stalledProvider, listWorkspaceIds: () => [], acquireWorkspaceMind: () => { throw new Error('not used'); }, releaseWorkspaceMind: () => undefined, }); const run = service.runNow(); await Promise.resolve(); expect(reprobe).not.toHaveBeenCalled(); await expect(run).resolves.toMatchObject>({ mindsVisited: 1, passes: 1, }); await expect(service.stop()).resolves.toBeUndefined(); }); });