@vellumai/assistant 0.6.4 → 0.6.6

package/src/__tests__/estimator-calibration-integration.test.ts

@@ -0,0 +1,208 @@
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import {
+  getCalibrationSnapshot,
+  getCorrection,
+  recordEstimate,
+  resetCalibrations,
+} from "../context/estimator-calibration.js";
+import {
+  estimatePromptTokens,
+  getCalibrationProviderKey,
+} from "../context/token-estimator.js";
+import type { Message, Provider } from "../providers/types.js";
+
+/**
+ * Integration-style tests that exercise the full self-calibration loop end
+ * to end:
+ *   1. Estimate is recorded for a `(provider, model)` pair via
+ *      `handleUsage` (the record side still threads the provider-echoed
+ *      model through `recordEstimate`).
+ *   2. A subsequent `estimatePromptTokens` lookup picks up the learned
+ *      correction via the per-provider aggregate key `(provider, "")`.
+ *      Lookup always uses the aggregate — model-specific keys are only
+ *      read as a fallback inside `getCorrection`.
+ *
+ * Since the `modelId` lookup option has been removed from the public
+ * token-estimator API, the lookup side always converges to the aggregate.
+ * `recordEstimate` still updates both the specific `(provider, model)`
+ * key AND the `(provider, "")` aggregate on every sample, so the
+ * aggregate stays accurate even as per-model data accumulates.
+ */
+describe("estimator calibration — end-to-end recording → lookup", () => {
+  beforeEach(() => {
+    resetCalibrations();
+  });
+
+  /**
+   * Build a representative message history with enough content to clear the
+   * MIN_SAMPLE_MAGNITUDE floor (500 tokens). Each message repeats a block of
+   * text large enough to make the heuristic estimator produce a substantial
+   * token count so the calibration machinery actually runs.
+   */
+  function largeHistory(): Message[] {
+    const body = "lorem ipsum dolor sit amet ".repeat(500);
+    return [
+      { role: "user", content: [{ type: "text", text: body }] },
+      { role: "assistant", content: [{ type: "text", text: body }] },
+      { role: "user", content: [{ type: "text", text: body }] },
+    ];
+  }
+
+  test("subsequent estimate picks up the aggregate-key correction", () => {
+    const provider: Provider = {
+      name: "anthropic",
+      async sendMessage() {
+        throw new Error("not used in this test");
+      },
+    };
+    const model = "claude-sonnet-4-5";
+    const history = largeHistory();
+
+    // 1. Raw estimate (what agent/loop.ts computes pre-send).
+    const preSend = estimatePromptTokens(history, "system", {
+      providerName: getCalibrationProviderKey(provider),
+    });
+    expect(preSend).toBeGreaterThan(0);
+
+    // Baseline: no correction recorded yet.
+    expect(getCorrection("anthropic", "")).toBe(1.0);
+
+    // 2. Provider returns ground truth (simulating `handleUsage`, which
+    //    still records under (provider, event.model) and folds into the
+    //    aggregate). Simulate a systematic 30% underestimate.
+    const groundTruth = Math.ceil(preSend * 1.3);
+    recordEstimate(
+      getCalibrationProviderKey(provider),
+      model,
+      preSend,
+      groundTruth,
+    );
+
+    // 3. Lookup under the aggregate key now returns the learned ratio.
+    expect(getCorrection("anthropic", "")).toBeCloseTo(1.3, 3);
+
+    // And the corrected estimate moves toward the ground truth.
+    const corrected = estimatePromptTokens(history, "system", {
+      providerName: getCalibrationProviderKey(provider),
+    });
+    // With correction factor ≈1.3, corrected estimate is within 1 token of
+    // the ground truth (Math.ceil rounding).
+    expect(corrected).toBeGreaterThan(preSend);
+    expect(Math.abs(corrected - groundTruth)).toBeLessThanOrEqual(1);
+  });
+
+  test("record with model writes both the specific and aggregate keys", () => {
+    // Simulate a preflight site that records against (anthropic, sonnet).
+    // `recordEstimate` also folds the sample into the `(anthropic, "")`
+    // aggregate so aggregate-key callers see the correction.
+    const provider: Provider = {
+      name: "anthropic",
+      async sendMessage() {
+        throw new Error("not used");
+      },
+    };
+    const history = largeHistory();
+
+    const preSend = estimatePromptTokens(history, "system", {
+      providerName: getCalibrationProviderKey(provider),
+    });
+    const groundTruth = Math.ceil(preSend * 1.25);
+
+    recordEstimate(
+      getCalibrationProviderKey(provider),
+      "claude-sonnet-4-5",
+      preSend,
+      groundTruth,
+    );
+
+    // A subsequent lookup via the token-estimator uses the per-provider
+    // aggregate (the only key the public API reads).
+    const correctedAggregate = estimatePromptTokens(history, "system", {
+      providerName: getCalibrationProviderKey(provider),
+    });
+    // Aggregate ratio ≈ 1.25 (first sample snaps to exact ratio).
+    expect(correctedAggregate).toBe(Math.ceil(preSend * 1.25));
+  });
+
+  test("wrapper provider (OpenRouter → Anthropic) uses the canonical key on both sides", () => {
+    // This is the Devin scenario: OpenRouter wraps Anthropic. If the record
+    // site used `name` ("openrouter") and the lookup site used
+    // `tokenEstimationProvider` ("anthropic"), the data would be scattered
+    // across mismatched keys and calibration would silently fail.
+    // `getCalibrationProviderKey` gives us one source of truth.
+    const openrouter: Provider = {
+      name: "openrouter",
+      tokenEstimationProvider: "anthropic",
+      async sendMessage() {
+        throw new Error("not used");
+      },
+    };
+    const model = "anthropic/claude-sonnet-4-5";
+    const history = largeHistory();
+
+    // Pre-send estimate via the canonical key.
+    const preSend = estimatePromptTokens(history, "system", {
+      providerName: getCalibrationProviderKey(openrouter),
+    });
+    expect(preSend).toBeGreaterThan(0);
+
+    // Provider returns ground truth. `handleUsage` uses the same helper
+    // to pick the calibration key, so the record and lookup sides agree.
+    const groundTruth = Math.ceil(preSend * 1.2);
+    recordEstimate(
+      getCalibrationProviderKey(openrouter),
+      model,
+      preSend,
+      groundTruth,
+    );
+
+    // Lookup under "anthropic" — the canonical upstream key — returns the
+    // ratio. See note above about precision=3.
+    expect(getCorrection("anthropic", model)).toBeCloseTo(1.2, 3);
+    // Aggregate under the canonical upstream key is also populated.
+    expect(getCorrection("anthropic", "")).toBeCloseTo(1.2, 3);
+    // And under the bare wrapper name stays at the default, because NOTHING
+    // was recorded under "openrouter".
+    expect(getCorrection("openrouter", "")).toBe(1.0);
+
+    // The snapshot reflects a single (provider, model) key + aggregate under
+    // the canonical upstream key — never under the wrapper name.
+    const keys = getCalibrationSnapshot().map(
+      (e) => `${e.provider}::${e.model}`,
+    );
+    expect(keys).toContain(`anthropic::${model}`);
+    expect(keys).toContain("anthropic::");
+    expect(keys).not.toContain(`openrouter::${model}`);
+  });
+
+  test("a run of consistent samples pulls the estimate toward ground truth", () => {
+    // The EWMA should converge quickly. After five consistent 1.3 samples
+    // the correction should be within 1% of 1.3, and the corrected estimate
+    // should be within 1% of the ground truth.
+    const model = "claude-sonnet-4-5";
+    const history = largeHistory();
+
+    const preSend = estimatePromptTokens(history, "system", {
+      providerName: "anthropic",
+    });
+    const groundTruth = Math.ceil(preSend * 1.3);
+
+    for (let i = 0; i < 5; i++) {
+      recordEstimate("anthropic", model, preSend, groundTruth);
+    }
+
+    const finalCorrection = getCorrection("anthropic", "");
+    // EWMA with alpha=0.2 on constant 1.3 stays at 1.3 from the first sample
+    // onward (all deltas are 0 after the initial snap). `precision=3` gives
+    // us ~0.0005 tolerance which covers the Math.ceil rounding noise.
+    expect(finalCorrection).toBeCloseTo(1.3, 3);
+
+    const corrected = estimatePromptTokens(history, "system", {
+      providerName: "anthropic",
+    });
+    // Corrected should be very close to the ground truth (within 1 token
+    // because of the Math.ceil rounding at the end of estimatePromptTokens).
+    expect(Math.abs(corrected - groundTruth)).toBeLessThanOrEqual(1);
+  });
+});

package/src/__tests__/estimator-calibration.test.ts

@@ -0,0 +1,213 @@
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import {
+  getCalibrationSnapshot,
+  getCorrection,
+  recordEstimate,
+  resetCalibrations,
+} from "../context/estimator-calibration.js";
+
+describe("estimator calibration", () => {
+  beforeEach(() => {
+    resetCalibrations();
+  });
+
+  test("default correction is 1.0 for unseen (provider, model)", () => {
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+    expect(getCorrection("openai", "gpt-5")).toBe(1.0);
+  });
+
+  test("first sample yields the exact ratio", () => {
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    const ratio = getCorrection("anthropic", "claude-sonnet-4-5");
+    expect(ratio).toBeCloseTo(1.3, 5);
+  });
+
+  test("EWMA converges to the target ratio given consistent samples", () => {
+    // Seed with a ratio far from the target so the first sample is off,
+    // then hammer with consistent 1.3 samples and watch EWMA close the gap.
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 100_000);
+    for (let i = 0; i < 20; i++) {
+      recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    }
+    const ratio = getCorrection("anthropic", "claude-sonnet-4-5");
+    expect(ratio).toBeGreaterThan(1.25);
+    expect(ratio).toBeLessThan(1.35);
+  });
+
+  test("ten consistent 1.3 samples land within 0.05 of 1.3", () => {
+    for (let i = 0; i < 10; i++) {
+      recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    }
+    const ratio = getCorrection("anthropic", "claude-sonnet-4-5");
+    expect(Math.abs(ratio - 1.3)).toBeLessThan(0.05);
+  });
+
+  test("values below MIN_SAMPLE_MAGNITUDE are ignored", () => {
+    // Both below the floor
+    recordEstimate("anthropic", "claude-sonnet-4-5", 200, 400);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+
+    // Estimated below the floor
+    recordEstimate("anthropic", "claude-sonnet-4-5", 200, 100_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+
+    // Actual below the floor
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 200);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+  });
+
+  test("ratios outside [1/3, 3] are discarded as outliers", () => {
+    // 4x too high
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 400_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+
+    // 4x too low
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 25_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+
+    // Just above the 3x edge — still discarded
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 300_001);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+
+    // Exactly at the 3x edge — accepted (ratio === 3, not > 3)
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 300_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(3, 5);
+  });
+
+  test("resetCalibrations clears all state", () => {
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    recordEstimate("openai", "gpt-5", 100_000, 90_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(1.3);
+    expect(getCorrection("openai", "gpt-5")).toBeCloseTo(0.9);
+
+    resetCalibrations();
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBe(1.0);
+    expect(getCorrection("openai", "gpt-5")).toBe(1.0);
+    expect(getCalibrationSnapshot()).toHaveLength(0);
+  });
+
+  test("distinct (provider, model) keys are independent", () => {
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    recordEstimate("anthropic", "claude-opus-4-7", 100_000, 110_000);
+    recordEstimate("openai", "gpt-5", 100_000, 90_000);
+    recordEstimate("openai", "gpt-5", 100_000, 90_000);
+    recordEstimate("openai", "gpt-5", 100_000, 90_000);
+
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(1.3, 5);
+    expect(getCorrection("anthropic", "claude-opus-4-7")).toBeCloseTo(1.1, 5);
+    // openai::gpt-5: after 3 EWMA steps of ratio 0.9, still exactly 0.9
+    // because the first sample snaps to ratio and subsequent deltas are 0.
+    expect(getCorrection("openai", "gpt-5")).toBeCloseTo(0.9, 5);
+
+    // Model separation: ensure a sample to one (provider, model) pair does
+    // not pollute another model under the same provider.
+    expect(getCorrection("anthropic", "claude-opus-4-7")).not.toBe(
+      getCorrection("anthropic", "claude-sonnet-4-5"),
+    );
+  });
+
+  test("snapshot reports current calibrations", () => {
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    recordEstimate("openai", "gpt-5", 100_000, 90_000);
+
+    const snap = getCalibrationSnapshot();
+    // Every non-empty-model sample also updates the per-provider aggregate
+    // key (provider, ""). Two providers × {specific model, aggregate} = 4.
+    expect(snap).toHaveLength(4);
+
+    const anthropicEntry = snap.find(
+      (e) => e.provider === "anthropic" && e.model === "claude-sonnet-4-5",
+    );
+    expect(anthropicEntry).toBeDefined();
+    expect(anthropicEntry?.samples).toBe(2);
+    expect(anthropicEntry?.ratio).toBeCloseTo(1.3, 5);
+
+    const openaiEntry = snap.find(
+      (e) => e.provider === "openai" && e.model === "gpt-5",
+    );
+    expect(openaiEntry).toBeDefined();
+    expect(openaiEntry?.samples).toBe(1);
+    expect(openaiEntry?.ratio).toBeCloseTo(0.9, 5);
+
+    // Per-provider aggregates are present and track the specific samples.
+    const anthropicAggregate = snap.find(
+      (e) => e.provider === "anthropic" && e.model === "",
+    );
+    expect(anthropicAggregate?.samples).toBe(2);
+    expect(anthropicAggregate?.ratio).toBeCloseTo(1.3, 5);
+
+    const openaiAggregate = snap.find(
+      (e) => e.provider === "openai" && e.model === "",
+    );
+    expect(openaiAggregate?.samples).toBe(1);
+    expect(openaiAggregate?.ratio).toBeCloseTo(0.9, 5);
+  });
+
+  test("specific-model lookup falls back to the per-provider aggregate when no model-specific samples exist", () => {
+    // Writing only to the aggregate simulates the case where callers without
+    // a resolved modelId have been recording, and then a model-specific
+    // caller asks for its correction.
+    recordEstimate("anthropic", "", 100_000, 130_000);
+    expect(getCorrection("anthropic", "")).toBeCloseTo(1.3, 5);
+    // The specific model has no own samples, but gets the aggregate.
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(1.3, 5);
+    // An unseen provider still defaults to 1.0.
+    expect(getCorrection("gemini", "gemini-2.5-pro")).toBe(1.0);
+  });
+
+  test("specific-model samples take precedence over the aggregate", () => {
+    recordEstimate("anthropic", "", 100_000, 130_000); // aggregate = 1.3
+    recordEstimate("anthropic", "claude-opus-4-7", 100_000, 110_000); // specific = 1.1
+    expect(getCorrection("anthropic", "claude-opus-4-7")).not.toBeCloseTo(
+      1.3,
+      5,
+    );
+    // A different model under the same provider still falls back to the
+    // aggregate (which now reflects both samples via EWMA).
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).not.toBe(1.0);
+  });
+
+  test("recording with a specific model also updates the per-provider aggregate", () => {
+    // The calibration writes `(provider, model)` and every caller of
+    // `estimatePromptTokens` that cannot resolve a model falls back to
+    // `(provider, "")`. Without the aggregate update, those callers would
+    // stay at the 1.0 default forever, defeating the whole mechanism.
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 130_000);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(1.3, 5);
+    expect(getCorrection("anthropic", "")).toBeCloseTo(1.3, 5);
+    // A different provider is unaffected.
+    expect(getCorrection("openai", "")).toBe(1.0);
+  });
+
+  test("the per-provider aggregate tracks a rolling EWMA across models", () => {
+    // Samples across two models for one provider should each feed the
+    // aggregate EWMA, so a generic lookup reflects recent activity from
+    // any model on that provider.
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 120_000);
+    recordEstimate("anthropic", "claude-opus-4-7", 100_000, 120_000);
+    recordEstimate("anthropic", "claude-sonnet-4-5", 100_000, 120_000);
+    recordEstimate("anthropic", "claude-opus-4-7", 100_000, 120_000);
+
+    // After four consistent 1.2 samples (folded into both the model-specific
+    // and the aggregate keys) the EWMA for each sits very close to 1.2.
+    expect(getCorrection("anthropic", "")).toBeCloseTo(1.2, 1);
+    expect(getCorrection("anthropic", "claude-sonnet-4-5")).toBeCloseTo(1.2, 1);
+    expect(getCorrection("anthropic", "claude-opus-4-7")).toBeCloseTo(1.2, 1);
+  });
+
+  test("explicit empty-model recording does not double-count the aggregate", () => {
+    // When a caller passes an empty model string (degenerate case), the
+    // aggregate update path must not run — otherwise the sample would be
+    // folded into the EWMA twice and the ratio would be wrong.
+    recordEstimate("anthropic", "", 100_000, 120_000);
+    // Exactly one sample applied.
+    const snap = getCalibrationSnapshot().filter(
+      (e) => e.provider === "anthropic" && e.model === "",
+    );
+    expect(snap).toHaveLength(1);
+    expect(snap[0].samples).toBe(1);
+    expect(snap[0].ratio).toBeCloseTo(1.2, 5);
+  });
+});

package/src/__tests__/extension-id-sync-guard.test.ts

@@ -16,7 +16,7 @@ import { homedir } from "node:os";
 import { join, resolve } from "node:path";
 import { describe, expect, test } from "bun:test";
 
-import { ALLOWED_EXTENSION_ORIGINS } from "../runtime/routes/browser-extension-pair-routes.js";
+import { getAllowedExtensionOrigins } from "../runtime/routes/browser-extension-pair-routes.js";
 
 const repoRoot = resolve(__dirname, "..", "..", "..");
 const CANONICAL_CONFIG_REL_PATH =
@@ -41,11 +41,15 @@ function parseCanonicalConfig(): AllowlistConfig {
   const parsed = JSON.parse(raw) as Partial<AllowlistConfig>;
 
   if (!Number.isInteger(parsed.version) || (parsed.version ?? 0) <= 0) {
-    throw new Error("Invalid canonical config: version must be a positive integer");
+    throw new Error(
+      "Invalid canonical config: version must be a positive integer",
+    );
   }
 
   if (!Array.isArray(parsed.allowedExtensionIds)) {
-    throw new Error("Invalid canonical config: allowedExtensionIds must be an array");
+    throw new Error(
+      "Invalid canonical config: allowedExtensionIds must be an array",
+    );
   }
 
   if (parsed.allowedExtensionIds.length === 0) {
@@ -95,6 +99,7 @@ function listTextFilesRecursively(root: string): string[] {
     ".turbo",
     ".idea",
     ".vscode",
+    ".codex-worktrees",
   ]);
 
   const allowedExtensions = new Set([
@@ -171,7 +176,7 @@ describe("Chrome extension allowlist guard", () => {
     const config = parseCanonicalConfig();
     for (const id of config.allowedExtensionIds) {
       const origin = `chrome-extension://${id}/`;
-      expect(ALLOWED_EXTENSION_ORIGINS.has(origin)).toBe(true);
+      expect(getAllowedExtensionOrigins().has(origin)).toBe(true);
     }
   });
 
@@ -193,15 +198,23 @@ describe("Chrome extension allowlist guard", () => {
     const expectedOrigins = new Set(
       config.allowedExtensionIds.map((id) => `chrome-extension://${id}/`),
     );
-    expect(new Set(ALLOWED_EXTENSION_ORIGINS)).toEqual(expectedOrigins);
+    expect(new Set(getAllowedExtensionOrigins())).toEqual(expectedOrigins);
   });
 
-  test("concrete extension IDs appear only in canonical config", () => {
+  test("concrete extension IDs appear only in canonical config or CWS URLs", () => {
+    // The canonical extension ID may also appear in Chrome Web Store URLs
+    // (e.g. chromewebstore.google.com/detail/.../ID) or in documentation
+    // referencing the published extension. Those are acceptable — they
+    // reference the published extension, not duplicate config. We flag
+    // files where the ID appears outside of a CWS URL context.
     const config = parseCanonicalConfig();
     const allFiles = listTextFilesRecursively(repoRoot);
 
+    const CWS_URL_PATTERN =
+      /chromewebstore\.google\.com\/detail\/[^/]+\/[a-p]{32}/g;
+
     for (const extensionId of config.allowedExtensionIds) {
-      const matches: string[] = [];
+      const unexpectedMatches: string[] = [];
       for (const absPath of allFiles) {
         const relPath = absPath.replace(`${repoRoot}/`, "");
         let content: string;
@@ -212,11 +225,17 @@ describe("Chrome extension allowlist guard", () => {
           if ((err as NodeJS.ErrnoException).code === "ENOENT") continue;
           throw err;
         }
-        if (content.includes(extensionId)) {
-          matches.push(relPath);
+        if (!content.includes(extensionId)) continue;
+        if (relPath === CANONICAL_CONFIG_REL_PATH) continue;
+
+        // Strip all CWS URLs and check if the ID still appears — if it
+        // does, the file is using the ID as a standalone config value.
+        const stripped = content.replace(CWS_URL_PATTERN, "");
+        if (stripped.includes(extensionId)) {
+          unexpectedMatches.push(relPath);
         }
       }
-      expect(matches).toEqual([CANONICAL_CONFIG_REL_PATH]);
+      expect(unexpectedMatches).toEqual([]);
     }
   });
 });

package/src/__tests__/file-write-tool.test.ts

@@ -1,5 +1,6 @@
 import {
   existsSync,
+  mkdirSync,
   mkdtempSync,
   readFileSync,
   realpathSync,
@@ -8,8 +9,47 @@ import {
 } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { afterEach, beforeAll, describe, expect, test } from "bun:test";
+import {
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  mock,
+  test,
+} from "bun:test";
+
+// Track calls to enqueuePkbIndexJob across tests. Captured via mock.module
+// below; individual tests clear and inspect the array.
+const enqueueCalls: Array<{
+  pkbRoot: string;
+  absPath: string;
+  memoryScopeId: string;
+}> = [];
+let enqueueThrows = false;
 
+mock.module("../memory/jobs/embed-pkb-file.js", () => ({
+  enqueuePkbIndexJob: (input: {
+    pkbRoot: string;
+    absPath: string;
+    memoryScopeId: string;
+  }) => {
+    if (enqueueThrows) {
+      throw new Error("simulated enqueue failure");
+    }
+    enqueueCalls.push(input);
+    return "job-id";
+  },
+}));
+
+// Override workspace dir via VELLUM_WORKSPACE_DIR so PKB-root detection
+// targets a temp directory without having to mock platform.js wholesale
+// (which would destabilize the rest of the tool registry's dependency tree).
+function setWorkspaceDir(dir: string): void {
+  process.env.VELLUM_WORKSPACE_DIR = dir;
+}
+
+import { PKB_WORKSPACE_SCOPE } from "../memory/pkb/types.js";
 import { getTool } from "../tools/registry.js";
 import type { Tool, ToolContext } from "../tools/types.js";
 
@@ -35,6 +75,24 @@ afterEach(() => {
   }
 });
 
+const originalWorkspaceDirEnv = process.env.VELLUM_WORKSPACE_DIR;
+
+beforeEach(() => {
+  enqueueCalls.length = 0;
+  enqueueThrows = false;
+  // Reset to a stable tmp path so the sandbox tests (which don't use pkb/)
+  // deterministically land outside any configured PKB root.
+  process.env.VELLUM_WORKSPACE_DIR = tmpdir();
+});
+
+afterEach(() => {
+  if (originalWorkspaceDirEnv === undefined) {
+    delete process.env.VELLUM_WORKSPACE_DIR;
+  } else {
+    process.env.VELLUM_WORKSPACE_DIR = originalWorkspaceDirEnv;
+  }
+});
+
 function makeTempDir(): string {
   const dir = realpathSync(mkdtempSync(join(tmpdir(), "file-write-test-")));
   testDirs.push(dir);
@@ -118,3 +176,95 @@ describe("file_write tool (sandbox)", () => {
     expect(result.content).toContain("exceeds");
   });
 });
+
+describe("file_write tool PKB re-index hook", () => {
+  test("enqueues a PKB re-index job when writing under pkb/", async () => {
+    const workingDir = makeTempDir();
+    setWorkspaceDir(workingDir);
+    mkdirSync(join(workingDir, "pkb"), { recursive: true });
+
+    const result = await fileWriteTool.execute(
+      { path: "pkb/note.md", content: "# hello\nworld\n" },
+      makeContext(workingDir),
+    );
+
+    expect(result.isError).toBe(false);
+    expect(enqueueCalls).toHaveLength(1);
+    expect(enqueueCalls[0]).toEqual({
+      pkbRoot: join(workingDir, "pkb"),
+      absPath: join(workingDir, "pkb", "note.md"),
+      memoryScopeId: PKB_WORKSPACE_SCOPE,
+    });
+  });
+
+  test("always uses PKB_WORKSPACE_SCOPE regardless of context.memoryScopeId", async () => {
+    const workingDir = makeTempDir();
+    setWorkspaceDir(workingDir);
+    mkdirSync(join(workingDir, "pkb"), { recursive: true });
+
+    const ctx: ToolContext = {
+      ...makeContext(workingDir),
+      memoryScopeId: "private:abc123",
+    };
+
+    const result = await fileWriteTool.execute(
+      { path: "pkb/private.md", content: "secret\n" },
+      ctx,
+    );
+
+    expect(result.isError).toBe(false);
+    expect(enqueueCalls).toHaveLength(1);
+    // PKB files are workspace-level — the per-conversation scopeId is NOT
+    // threaded through so different conversations can't overwrite each
+    // other's Qdrant points via target_id upsert deduplication.
+    expect(enqueueCalls[0]?.memoryScopeId).toBe(PKB_WORKSPACE_SCOPE);
+  });
+
+  test("does NOT enqueue when writing outside pkb/", async () => {
+    const workingDir = makeTempDir();
+    setWorkspaceDir(workingDir);
+
+    const result = await fileWriteTool.execute(
+      { path: "notes.md", content: "# not pkb\n" },
+      makeContext(workingDir),
+    );
+
+    expect(result.isError).toBe(false);
+    expect(enqueueCalls).toHaveLength(0);
+  });
+
+  test("does NOT enqueue for a sibling directory whose name is a pkb prefix", async () => {
+    // Guard against `<root>/pkbsomethingelse` being treated as inside `<root>/pkb`.
+    const workingDir = makeTempDir();
+    setWorkspaceDir(workingDir);
+    mkdirSync(join(workingDir, "pkbsibling"), { recursive: true });
+
+    const result = await fileWriteTool.execute(
+      { path: "pkbsibling/file.md", content: "not pkb\n" },
+      makeContext(workingDir),
+    );
+
+    expect(result.isError).toBe(false);
+    expect(enqueueCalls).toHaveLength(0);
+  });
+
+  test("enqueue failure is swallowed and write result stays successful", async () => {
+    const workingDir = makeTempDir();
+    setWorkspaceDir(workingDir);
+    mkdirSync(join(workingDir, "pkb"), { recursive: true });
+    enqueueThrows = true;
+
+    const result = await fileWriteTool.execute(
+      { path: "pkb/oops.md", content: "still writes\n" },
+      makeContext(workingDir),
+    );
+
+    expect(result.isError).toBe(false);
+    // The mock throws, so nothing gets pushed to enqueueCalls. The critical
+    // behavior is that the thrown error never surfaces through execute().
+    expect(enqueueCalls).toHaveLength(0);
+    expect(
+      existsSync(join(workingDir, "pkb", "oops.md")),
+    ).toBe(true);
+  });
+});