waggle-os/packages/server/tests/local/vector-enrichment-service.test.ts

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> = {}): VectorBackfillResult {
  return {
    skipped: null,
    vectorsRepaired: false,
    framesReembedded: 0,
    chunksCreated: 0,
    framesProcessed: 0,
    hasMore: false,
    errors: [],
    ...overrides,
  };
}

function deferred<T>() {
  let resolve!: (value: T) => void;
  const promise = new Promise<T>((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<VectorBackfillResult>();
    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<MindDB, number>();
    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<VectorBackfillResult>();
    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<Partial<VectorEnrichmentRunResult>>({ 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<EmbeddingProviderStatus>(() => 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<Partial<VectorEnrichmentRunResult>>({
      mindsVisited: 1,
      passes: 1,
    });
    await expect(service.stop()).resolves.toBeUndefined();
  });
});