@vellumai/assistant 0.6.4 → 0.6.5

package/src/memory/pkb/pkb-index.test.ts

@@ -0,0 +1,368 @@
+import { mkdir, mkdtemp, rm, utimes, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+mock.module("../../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Capture calls to embedAndUpsert so we can assert on targetType + payload.
+const embedAndUpsertCalls: Array<{
+  config: unknown;
+  targetType: string;
+  targetId: string;
+  input: unknown;
+  extraPayload: unknown;
+}> = [];
+
+mock.module("../job-utils.js", () => ({
+  embedAndUpsert: async (
+    config: unknown,
+    targetType: string,
+    targetId: string,
+    input: unknown,
+    extraPayload: unknown,
+  ) => {
+    embedAndUpsertCalls.push({
+      config,
+      targetType,
+      targetId,
+      input,
+      extraPayload,
+    });
+  },
+}));
+
+// Minimal stub for getConfig — indexPkbFile forwards it opaquely to the
+// mocked embedAndUpsert, so any sentinel value works.
+mock.module("../../config/loader.js", () => ({
+  getConfig: () => ({ __stub: true }),
+}));
+
+// Track Qdrant deletes by capturing the filter the client sends.
+const qdrantDeleteCalls: Array<{
+  targetType: string;
+  path: string;
+  memoryScopeId: string;
+}> = [];
+
+// Track per-target deletes (used by write-then-cleanup in indexPkbFile).
+const qdrantDeleteByTargetCalls: Array<{
+  targetType: string;
+  targetId: string;
+}> = [];
+
+// Points the mocked scroll will return on the next call. Tests mutate this
+// to simulate pre-existing PKB chunks on disk.
+let scrollReturnPoints: Array<{
+  id: string;
+  payload: Record<string, unknown>;
+}> = [];
+const qdrantScrollCalls: Array<{
+  targetType: string;
+  memoryScopeId?: string;
+  path?: string;
+}> = [];
+
+mock.module("../qdrant-client.js", () => ({
+  getQdrantClient: () => ({
+    deleteByTargetTypeAndPath: async (
+      targetType: string,
+      path: string,
+      memoryScopeId: string,
+    ) => {
+      qdrantDeleteCalls.push({ targetType, path, memoryScopeId });
+    },
+    deleteByTarget: async (targetType: string, targetId: string) => {
+      qdrantDeleteByTargetCalls.push({ targetType, targetId });
+    },
+    scrollByTargetType: async (
+      targetType: string,
+      options?: {
+        memoryScopeId?: string;
+        path?: string;
+        batchSize?: number;
+      },
+    ) => {
+      qdrantScrollCalls.push({
+        targetType,
+        memoryScopeId: options?.memoryScopeId,
+        path: options?.path,
+      });
+      return scrollReturnPoints;
+    },
+  }),
+}));
+
+// The circuit breaker is a thin wrapper; just call the function through.
+mock.module("../qdrant-circuit-breaker.js", () => ({
+  withQdrantBreaker: async <T,>(fn: () => Promise<T>) => fn(),
+}));
+
+import {
+  chunkPkbFile,
+  deletePkbFilePoints,
+  indexPkbFile,
+  scanPkbFiles,
+} from "./pkb-index.js";
+
+describe("chunkPkbFile", () => {
+  test("returns whole-file for small inputs", () => {
+    const small = "a".repeat(500);
+    const chunks = chunkPkbFile(small);
+    expect(chunks).toEqual([small]);
+  });
+
+  test("splits on ## headings with lossless concatenation", () => {
+    const sectionA = "## Section A\n" + "a".repeat(5990) + "\n";
+    const sectionB = "## Section B\n" + "b".repeat(6010);
+    const content = sectionA + sectionB;
+    expect(content.length).toBeGreaterThanOrEqual(12000);
+
+    const chunks = chunkPkbFile(content);
+    expect(chunks).toHaveLength(2);
+    expect(chunks.join("")).toBe(content);
+    expect(chunks[0].startsWith("## Section A")).toBe(true);
+    expect(chunks[1].startsWith("## Section B")).toBe(true);
+  });
+
+  test("falls back to char-window chunks when no ## headings exist", () => {
+    const content = "x".repeat(12000);
+    const chunks = chunkPkbFile(content);
+    // 12000 / 4000 = 3 windows.
+    expect(chunks).toHaveLength(3);
+    expect(chunks.join("")).toBe(content);
+    expect(chunks[0].length).toBe(4000);
+    expect(chunks[1].length).toBe(4000);
+    expect(chunks[2].length).toBe(4000);
+  });
+});
+
+describe("scanPkbFiles", () => {
+  test("returns entries for each .md file and ignores non-markdown", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-scan-"));
+    await writeFile(join(root, "a.md"), "# A\nalpha content");
+    await writeFile(join(root, "b.md"), "# B\nbeta content");
+    await writeFile(join(root, "notes.txt"), "plain text");
+
+    // Set deterministic mtimes so we can assert them.
+    const mtimeA = new Date(1_700_000_000_000);
+    const mtimeB = new Date(1_700_000_001_000);
+    await utimes(join(root, "a.md"), mtimeA, mtimeA);
+    await utimes(join(root, "b.md"), mtimeB, mtimeB);
+
+    const entries = await scanPkbFiles(root);
+    expect(entries).not.toBeNull();
+    const byPath = new Map(entries!.map((e) => [e.path, e]));
+
+    expect(byPath.size).toBe(2);
+    expect(byPath.has("a.md")).toBe(true);
+    expect(byPath.has("b.md")).toBe(true);
+    expect(byPath.has("notes.txt")).toBe(false);
+
+    const a = byPath.get("a.md")!;
+    expect(a.mtimeMs).toBe(mtimeA.getTime());
+    expect(a.chunkIndex).toBe(0);
+    expect(a.contentHash).toHaveLength(16);
+
+    // Hash is stable across scans.
+    const entriesAgain = await scanPkbFiles(root);
+    const aAgain = entriesAgain!.find((e) => e.path === "a.md")!;
+    expect(aAgain.contentHash).toBe(a.contentHash);
+  });
+
+  test("walks nested directories", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-scan-nested-"));
+    const sub = join(root, "sub");
+    await mkdir(sub);
+    await writeFile(join(sub, "nested.md"), "# nested");
+
+    const entries = await scanPkbFiles(root);
+    expect(entries).toHaveLength(1);
+    expect(entries![0].path).toBe(join("sub", "nested.md"));
+  });
+
+  test("returns null when pkbRoot does not exist", async () => {
+    const parent = await mkdtemp(join(tmpdir(), "pkb-scan-missing-"));
+    const missing = join(parent, "does-not-exist");
+    const entries = await scanPkbFiles(missing);
+    expect(entries).toBeNull();
+  });
+
+  test("returns null when pkbRoot existed then was removed", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-scan-removed-"));
+    await writeFile(join(root, "a.md"), "# A");
+    await rm(root, { recursive: true, force: true });
+
+    const entries = await scanPkbFiles(root);
+    expect(entries).toBeNull();
+  });
+
+  test("returns [] (not null) when pkbRoot exists but is empty", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-scan-empty-"));
+    const entries = await scanPkbFiles(root);
+    expect(entries).not.toBeNull();
+    expect(entries).toEqual([]);
+  });
+
+  test("returns null when pkbRoot points at a file instead of a directory", async () => {
+    const parent = await mkdtemp(join(tmpdir(), "pkb-scan-file-"));
+    const filePath = join(parent, "not-a-dir");
+    await writeFile(filePath, "just a file");
+    const entries = await scanPkbFiles(filePath);
+    expect(entries).toBeNull();
+  });
+});
+
+describe("indexPkbFile", () => {
+  beforeEach(() => {
+    embedAndUpsertCalls.length = 0;
+    qdrantDeleteCalls.length = 0;
+    qdrantDeleteByTargetCalls.length = 0;
+    qdrantScrollCalls.length = 0;
+    scrollReturnPoints = [];
+  });
+
+  test("invokes embedAndUpsert once per chunk with pkb_file target_type", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-"));
+    const filePath = join(root, "doc.md");
+    await writeFile(filePath, "# hello\nworld");
+
+    await indexPkbFile(root, filePath, "scope-xyz");
+
+    expect(embedAndUpsertCalls).toHaveLength(1);
+    const call = embedAndUpsertCalls[0];
+    expect(call.targetType).toBe("pkb_file");
+    expect(call.targetId).toBe("scope-xyz:doc.md#0");
+    expect(call.input).toEqual({ type: "text", text: "# hello\nworld" });
+    const payload = call.extraPayload as Record<string, unknown>;
+    expect(payload.path).toBe("doc.md");
+    expect(payload.chunk_index).toBe(0);
+    expect(payload.memory_scope_id).toBe("scope-xyz");
+    expect(typeof payload.mtime_ms).toBe("number");
+    expect(typeof payload.content_hash).toBe("string");
+    expect((payload.content_hash as string).length).toBe(16);
+  });
+
+  test("emits one embedAndUpsert call per chunk for a large file", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-large-"));
+    const filePath = join(root, "big.md");
+    const content =
+      "## Section A\n" +
+      "a".repeat(5990) +
+      "\n## Section B\n" +
+      "b".repeat(5990);
+    await writeFile(filePath, content);
+
+    await indexPkbFile(root, filePath, "scope-1");
+
+    expect(embedAndUpsertCalls).toHaveLength(2);
+    expect(embedAndUpsertCalls[0].targetId).toBe("scope-1:big.md#0");
+    expect(embedAndUpsertCalls[1].targetId).toBe("scope-1:big.md#1");
+    expect(embedAndUpsertCalls.every((c) => c.targetType === "pkb_file")).toBe(
+      true,
+    );
+  });
+
+  test("scope-namespaces target ids so two scopes indexing the same path do not collide", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-scope-"));
+    const filePath = join(root, "shared.md");
+    await writeFile(filePath, "# shared");
+
+    await indexPkbFile(root, filePath, "alpha");
+    await indexPkbFile(root, filePath, "beta");
+
+    expect(embedAndUpsertCalls).toHaveLength(2);
+    const ids = embedAndUpsertCalls.map((c) => c.targetId);
+    expect(ids).toEqual(["alpha:shared.md#0", "beta:shared.md#0"]);
+  });
+
+  test("scrolls existing chunks scoped to (target_type, scope, path) before upserting", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-scroll-"));
+    const filePath = join(root, "noted.md");
+    await writeFile(filePath, "# one");
+
+    await indexPkbFile(root, filePath, "scope-xyz");
+
+    expect(qdrantScrollCalls).toHaveLength(1);
+    expect(qdrantScrollCalls[0]).toEqual({
+      targetType: "pkb_file",
+      memoryScopeId: "scope-xyz",
+      path: "noted.md",
+    });
+  });
+
+  test("deletes only stale chunks after upserting (write-then-cleanup)", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-shrink-"));
+    const filePath = join(root, "shrinking.md");
+    // New content produces a single chunk (index #0). Pre-existing chunks
+    // #0..#2 simulate a prior run over a larger file.
+    await writeFile(filePath, "# just one");
+    scrollReturnPoints = [
+      {
+        id: "point-0",
+        payload: { target_id: "scope-xyz:shrinking.md#0" },
+      },
+      {
+        id: "point-1",
+        payload: { target_id: "scope-xyz:shrinking.md#1" },
+      },
+      {
+        id: "point-2",
+        payload: { target_id: "scope-xyz:shrinking.md#2" },
+      },
+    ];
+
+    await indexPkbFile(root, filePath, "scope-xyz");
+
+    // Exactly one upsert for the surviving chunk.
+    expect(embedAndUpsertCalls).toHaveLength(1);
+    expect(embedAndUpsertCalls[0].targetId).toBe("scope-xyz:shrinking.md#0");
+
+    // The pre-delete is gone; only the two stale chunks are removed.
+    expect(qdrantDeleteCalls).toHaveLength(0);
+    const staleTargetIds = qdrantDeleteByTargetCalls.map((c) => c.targetId);
+    expect(staleTargetIds.sort()).toEqual([
+      "scope-xyz:shrinking.md#1",
+      "scope-xyz:shrinking.md#2",
+    ]);
+  });
+
+  test("does not delete points whose target_id is regenerated", async () => {
+    const root = await mkdtemp(join(tmpdir(), "pkb-index-stable-"));
+    const filePath = join(root, "stable.md");
+    await writeFile(filePath, "# same");
+    scrollReturnPoints = [
+      {
+        id: "point-0",
+        payload: { target_id: "scope-xyz:stable.md#0" },
+      },
+    ];
+
+    await indexPkbFile(root, filePath, "scope-xyz");
+
+    expect(embedAndUpsertCalls).toHaveLength(1);
+    expect(qdrantDeleteByTargetCalls).toHaveLength(0);
+  });
+});
+
+describe("deletePkbFilePoints", () => {
+  beforeEach(() => {
+    qdrantDeleteCalls.length = 0;
+  });
+
+  test("sends a filter with target_type, path, and memory_scope_id predicates", async () => {
+    await deletePkbFilePoints("notes/todo.md", "scope-xyz");
+
+    expect(qdrantDeleteCalls).toHaveLength(1);
+    expect(qdrantDeleteCalls[0]).toEqual({
+      targetType: "pkb_file",
+      path: "notes/todo.md",
+      memoryScopeId: "scope-xyz",
+    });
+  });
+});

package/src/memory/pkb/pkb-index.ts

@@ -0,0 +1,255 @@
+/**
+ * PKB (Personal Knowledge Base) filesystem indexing primitives.
+ *
+ * Provides the low-level building blocks used by the PKB job handler and
+ * startup reconciliation:
+ *   - `scanPkbFiles`: recursively walk a PKB root and emit one entry per chunk.
+ *   - `chunkPkbFile`: split a markdown file into retrieval-friendly chunks.
+ *   - `indexPkbFile`: embed each chunk and upsert it to Qdrant.
+ *   - `deletePkbFilePoints`: remove every Qdrant point for a given file.
+ *
+ * Consumers (job queue wiring, startup scan) land in later PRs.
+ */
+
+import { createHash } from "node:crypto";
+import { readdir, readFile, stat } from "node:fs/promises";
+import { join, relative } from "node:path";
+
+import { getConfig } from "../../config/loader.js";
+import { embedAndUpsert } from "../job-utils.js";
+import { withQdrantBreaker } from "../qdrant-circuit-breaker.js";
+import { getQdrantClient } from "../qdrant-client.js";
+import type { PkbIndexEntry } from "./types.js";
+import { PKB_TARGET_TYPE } from "./types.js";
+
+/** Files larger than this are split into chunks for retrieval. */
+const WHOLE_FILE_THRESHOLD = 8000;
+
+/** Character-window size when falling back for unstructured content. */
+const CHAR_WINDOW_SIZE = 4000;
+
+/**
+ * Recursively walk `pkbRoot` and return one `PkbIndexEntry` per chunk of
+ * every `*.md` file found. Paths in the returned entries are relative to
+ * `pkbRoot`; mtime is read from the filesystem and `contentHash` is the
+ * first 16 hex chars of the sha256 of the file's contents.
+ *
+ * Returns `null` if `pkbRoot` cannot be confirmed as an existing directory
+ * (missing, not a directory, or stat failed for any reason — ENOENT, EACCES,
+ * EIO, etc.). This is distinct from returning `[]` for a directory that
+ * exists but has no `*.md` files — callers that run destructive
+ * reconciliation against the result (e.g. `reconcilePkbIndex`) use the
+ * sentinel to avoid interpreting an unreadable or transiently missing root
+ * as "delete every indexed point".
+ */
+export async function scanPkbFiles(
+  pkbRoot: string,
+): Promise<PkbIndexEntry[] | null> {
+  const entries: PkbIndexEntry[] = [];
+
+  // Verify the root exists up front. Any failure to confirm the root is a
+  // directory — ENOENT, EACCES, EIO, or a path that exists but isn't a
+  // directory — returns the missing sentinel so callers that run destructive
+  // reconciliation don't interpret "couldn't read the tree" as "disk is
+  // empty, delete everything indexed".
+  try {
+    const rootStat = await stat(pkbRoot);
+    if (!rootStat.isDirectory()) {
+      return null;
+    }
+  } catch {
+    return null;
+  }
+
+  async function walk(dir: string): Promise<void> {
+    let dirents;
+    try {
+      dirents = await readdir(dir, { withFileTypes: true });
+    } catch {
+      return;
+    }
+
+    for (const dirent of dirents) {
+      const absPath = join(dir, dirent.name);
+      if (dirent.isDirectory()) {
+        await walk(absPath);
+        continue;
+      }
+      if (!dirent.isFile()) continue;
+      if (!dirent.name.toLowerCase().endsWith(".md")) continue;
+
+      let content: string;
+      let mtimeMs: number;
+      try {
+        content = await readFile(absPath, "utf8");
+        const st = await stat(absPath);
+        mtimeMs = st.mtimeMs;
+      } catch {
+        continue;
+      }
+
+      const contentHash = hashContent(content);
+      const relPath = relative(pkbRoot, absPath);
+      const chunks = chunkPkbFile(content);
+      for (let chunkIndex = 0; chunkIndex < chunks.length; chunkIndex++) {
+        entries.push({
+          path: relPath,
+          mtimeMs,
+          contentHash,
+          chunkIndex,
+        });
+      }
+    }
+  }
+
+  await walk(pkbRoot);
+  return entries;
+}
+
+/**
+ * Split markdown content into retrieval chunks.
+ *
+ * Strategy:
+ *   - If the file is small (< WHOLE_FILE_THRESHOLD chars), return the whole
+ *     file as a single chunk.
+ *   - Otherwise split on lines starting with `## `, keeping each heading
+ *     with the body of its section. Concatenation of the returned chunks is
+ *     lossless — no content is dropped or duplicated.
+ *   - If no `## ` headings are present, fall back to fixed-size character
+ *     windows.
+ */
+export function chunkPkbFile(content: string): string[] {
+  if (content.length < WHOLE_FILE_THRESHOLD) {
+    return [content];
+  }
+
+  const headingIndices: number[] = [];
+  let cursor = 0;
+  while (cursor < content.length) {
+    // Find the next line that starts with "## " (either at the very beginning
+    // of the file or immediately after a newline).
+    const atStart = cursor === 0 && content.startsWith("## ");
+    if (atStart) {
+      headingIndices.push(0);
+      cursor = 1; // advance past the match so the `indexOf` below keeps moving
+      continue;
+    }
+    const nextNewline = content.indexOf("\n## ", cursor);
+    if (nextNewline === -1) break;
+    headingIndices.push(nextNewline + 1);
+    cursor = nextNewline + 1;
+  }
+
+  if (headingIndices.length === 0) {
+    // Fallback: fixed-size char windows.
+    const chunks: string[] = [];
+    for (let i = 0; i < content.length; i += CHAR_WINDOW_SIZE) {
+      chunks.push(content.slice(i, i + CHAR_WINDOW_SIZE));
+    }
+    return chunks;
+  }
+
+  // Build chunks from heading boundaries. Preserve any preamble before the
+  // first heading so concatenation stays lossless.
+  const chunks: string[] = [];
+  if (headingIndices[0] > 0) {
+    chunks.push(content.slice(0, headingIndices[0]));
+  }
+  for (let i = 0; i < headingIndices.length; i++) {
+    const start = headingIndices[i];
+    const end = i + 1 < headingIndices.length ? headingIndices[i + 1] : content.length;
+    chunks.push(content.slice(start, end));
+  }
+  return chunks;
+}
+
+/**
+ * Read a PKB file, chunk it, and upsert each chunk to Qdrant via the shared
+ * embedding pipeline. `relPath` in the payload is computed relative to
+ * `pkbRoot`.
+ *
+ * Write-then-cleanup: enumerate the existing chunk target_ids for this
+ * (scope, path), upsert every new chunk, then delete only the stale
+ * target_ids (those the fresh content no longer produces). Upsert dedupes on
+ * (target_type, target_id), so matching chunk indexes are replaced in place
+ * and the prior index stays queryable if any embed/upsert call throws. A
+ * pre-delete would instead leave the file unsearchable on transient failure
+ * until a successful retry.
+ */
+export async function indexPkbFile(
+  pkbRoot: string,
+  absPath: string,
+  memoryScopeId: string,
+): Promise<void> {
+  const content = await readFile(absPath, "utf8");
+  const st = await stat(absPath);
+  const mtimeMs = st.mtimeMs;
+  const contentHash = hashContent(content);
+  const relPath = relative(pkbRoot, absPath);
+  const chunks = chunkPkbFile(content);
+
+  const config = getConfig();
+
+  const qdrant = getQdrantClient();
+  const existing = await withQdrantBreaker(() =>
+    qdrant.scrollByTargetType(PKB_TARGET_TYPE, {
+      memoryScopeId,
+      path: relPath,
+    }),
+  );
+
+  const newTargetIds = new Set<string>();
+  for (let chunkIndex = 0; chunkIndex < chunks.length; chunkIndex++) {
+    const chunk = chunks[chunkIndex];
+    // Scope-namespace the target_id so `qdrant.upsert` — which dedupes on
+    // (target_type, target_id) — cannot collapse distinct scopes' chunks of
+    // the same relpath into a single point. Without the scope prefix, the
+    // second scope to index a shared path would overwrite the first's vectors.
+    const targetId = `${memoryScopeId}:${relPath}#${chunkIndex}`;
+    newTargetIds.add(targetId);
+    await embedAndUpsert(
+      config,
+      PKB_TARGET_TYPE,
+      targetId,
+      { type: "text", text: chunk },
+      {
+        path: relPath,
+        mtime_ms: mtimeMs,
+        chunk_index: chunkIndex,
+        content_hash: contentHash,
+        memory_scope_id: memoryScopeId,
+      },
+    );
+  }
+
+  // All upserts succeeded — safe to remove stale chunks that the new content
+  // did not regenerate (e.g. the file shrunk from 4 chunks to 2).
+  for (const point of existing) {
+    const staleTargetId = point.payload.target_id;
+    if (typeof staleTargetId !== "string") continue;
+    if (newTargetIds.has(staleTargetId)) continue;
+    await withQdrantBreaker(() =>
+      qdrant.deleteByTarget(PKB_TARGET_TYPE, staleTargetId),
+    );
+  }
+}
+
+/**
+ * Remove every Qdrant point belonging to a given PKB file (all chunks) within
+ * a single memory scope. `relPath` must match the `path` payload written by
+ * `indexPkbFile`. The `memoryScopeId` filter is required — omitting it would
+ * wipe that relpath's chunks across every scope that indexes the same file.
+ */
+export async function deletePkbFilePoints(
+  relPath: string,
+  memoryScopeId: string,
+): Promise<void> {
+  const qdrant = getQdrantClient();
+  await withQdrantBreaker(() =>
+    qdrant.deleteByTargetTypeAndPath(PKB_TARGET_TYPE, relPath, memoryScopeId),
+  );
+}
+
+function hashContent(content: string): string {
+  return createHash("sha256").update(content).digest("hex").slice(0, 16);
+}