SRCmail/lib/ai/retrieval/mail-embeddings.ts

// The `server` embedding leg of retrieval (docs/AI-ASSISTANT-CONCEPT.md §7
// step 2, §8.1's mail RetrieverAdapter) — real JMAP fetch, real embeddings
// via Ollama's /api/embed, real cosine similarity. No mocked vectors
// anywhere in this file.
//
// Persistence: in-memory only, per server process, keyed by accountId, with
// a TTL that triggers a full re-embed. A real persistent vector store with
// incremental updates (the doc's own P4 phase) is real follow-up work, not
// a same-night stretch goal on top of everything else built tonight — this
// is honest about that rather than pretending otherwise.
import { fetchJmapSession, postJmap, rebaseApiUrl } from '@/lib/stalwart/jmap-api';
import type { Chunk, Scored, SourceRef } from './types';

const EMBED_MODEL = process.env.AI_EMBED_MODEL || 'nomic-embed-text';
const MAX_EMAILS = 200;
const CACHE_TTL_MS = 5 * 60 * 1000;
const MAX_CHUNK_CHARS = 1000;

interface CachedEntry {
  ref: SourceRef;
  title: string;
  text: string;
  vector: number[];
}

interface CacheRecord {
  builtAt: number;
  entries: CachedEntry[];
}

// globalThis-stashed like entitlement.ts/config-manager.ts, so dev HMR
// doesn't silently start re-embedding on every hot reload.
const CACHE_KEY = Symbol.for('vncmail.ai.mail-embeddings-cache');
type GlobalWithCache = typeof globalThis & { [CACHE_KEY]?: Map<string, CacheRecord> };

function getCache(): Map<string, CacheRecord> {
  const g = globalThis as GlobalWithCache;
  if (!g[CACHE_KEY]) g[CACHE_KEY] = new Map();
  return g[CACHE_KEY];
}

async function embed(texts: string[]): Promise<number[][]> {
  const baseUrl = process.env.AI_SERVER_BASE_URL;
  if (!baseUrl) throw new Error('AI_SERVER_BASE_URL is not configured');
  const res = await fetch(`${baseUrl.replace(/\/+$/, '')}/api/embed`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ model: EMBED_MODEL, input: texts }),
  });
  if (!res.ok) throw new Error(`embedding runtime returned ${res.status}`);
  const body = (await res.json()) as { embeddings?: number[][] };
  if (!body.embeddings || body.embeddings.length !== texts.length) {
    throw new Error('embedding runtime returned an unexpected shape');
  }
  return body.embeddings;
}

function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0;
  let normA = 0;
  let normB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  const denom = Math.sqrt(normA) * Math.sqrt(normB);
  return denom === 0 ? 0 : dot / denom;
}

interface JmapEmail {
  id: string;
  mailboxIds?: Record<string, boolean>;
  subject?: string;
  preview?: string;
  receivedAt?: string;
}

async function fetchRecentMail(serverUrl: string, authHeader: string): Promise<{ accountId: string; emails: JmapEmail[] }> {
  const session = await fetchJmapSession(serverUrl, authHeader);
  if (!session) throw new Error('no JMAP session');
  const accountId = session.primaryAccounts?.['urn:ietf:params:jmap:mail'];
  if (!accountId) throw new Error('no primary mail account');
  const apiUrl = rebaseApiUrl(session, serverUrl);
  if (!apiUrl) throw new Error('session advertises no usable apiUrl');

  // Back-reference: Email/get's #ids resolves against the previous call's
  // result within the same request, one round trip instead of two.
  const res = await postJmap(apiUrl, authHeader, JSON.stringify({
    using: ['urn:ietf:params:jmap:core', 'urn:ietf:params:jmap:mail'],
    methodCalls: [
      ['Email/query', {
        accountId,
        sort: [{ property: 'receivedAt', isAscending: false }],
        limit: MAX_EMAILS,
      }, '0'],
      ['Email/get', {
        accountId,
        '#ids': { resultOf: '0', name: 'Email/query', path: '/ids' },
        properties: ['id', 'mailboxIds', 'subject', 'preview', 'receivedAt'],
      }, '1'],
    ],
  }));
  if (!res.ok) throw new Error(`Email/get returned ${res.status}`);

  const payload = await res.json() as {
    methodResponses?: [string, { list?: JmapEmail[] }, string][];
  };
  const getResult = payload.methodResponses?.find((r) => r[2] === '1');
  if (!getResult || getResult[0] !== 'Email/get') throw new Error('Email/get failed');
  return { accountId, emails: getResult[1]?.list ?? [] };
}

/**
 * Rebuild (or return the cached) embedding index for this account. Real
 * work only happens on a cache miss/expiry — repeated questions in the same
 * session don't re-embed everything.
 */
async function getOrBuildCache(accountId: string, serverUrl: string, authHeader: string): Promise<CacheRecord> {
  const cache = getCache();
  const existing = cache.get(accountId);
  if (existing && Date.now() - existing.builtAt < CACHE_TTL_MS) return existing;

  const { emails } = await fetchRecentMail(serverUrl, authHeader);
  const candidates = emails
    .map((email) => {
      const text = `${email.subject ?? ''}\n${email.preview ?? ''}`.slice(0, MAX_CHUNK_CHARS).trim();
      const collectionId = Object.keys(email.mailboxIds ?? {})[0] ?? 'unknown';
      return { email, text, collectionId };
    })
    .filter((c) => c.text.length > 0);

  if (candidates.length === 0) {
    const empty: CacheRecord = { builtAt: Date.now(), entries: [] };
    cache.set(accountId, empty);
    return empty;
  }

  const vectors = await embed(candidates.map((c) => c.text));
  const entries: CachedEntry[] = candidates.map((c, i) => ({
    ref: { product: 'mail', accountId, collectionId: c.collectionId, itemId: c.email.id, chunkIx: 0 },
    title: c.email.subject || '(no subject)',
    text: c.text,
    vector: vectors[i],
  }));

  const record: CacheRecord = { builtAt: Date.now(), entries };
  cache.set(accountId, record);
  return record;
}

export async function serverSearchMail(
  serverUrl: string,
  authHeader: string,
  query: string,
  limit: number,
): Promise<Scored<SourceRef>[]> {
  const session = await fetchJmapSession(serverUrl, authHeader);
  const accountId = session?.primaryAccounts?.['urn:ietf:params:jmap:mail'];
  if (!accountId) throw new Error('no primary mail account');

  const record = await getOrBuildCache(accountId, serverUrl, authHeader);
  if (record.entries.length === 0) return [];

  const [queryVector] = await embed([query]);
  return record.entries
    .map((entry) => ({ ref: entry.ref, score: cosineSimilarity(queryVector, entry.vector) }))
    .sort((a, b) => b.score - a.score)
    .slice(0, limit);
}

export async function hydrateMailRefs(
  serverUrl: string,
  authHeader: string,
  refs: SourceRef[],
): Promise<Chunk[]> {
  const session = await fetchJmapSession(serverUrl, authHeader);
  const accountId = session?.primaryAccounts?.['urn:ietf:params:jmap:mail'];
  if (!accountId) return [];

  const record = getCache().get(accountId);
  if (!record) return [];

  const byItemId = new Map(record.entries.map((e) => [e.ref.itemId, e]));
  return refs
    .map((ref) => byItemId.get(ref.itemId))
    .filter((e): e is CachedEntry => !!e)
    .map((e) => ({ ref: e.ref, text: e.text, title: e.title }));
}