pi-doc-injector 0.4.0 → 0.5.1

package/README.md

@@ -1,6 +1,6 @@
 # Pi Doc Injector
 
-A [Pi](https://pi.dev) extension that automatically injects relevant project documentation into the LLM system prompt by monitoring streaming output for keyword matches.
+A [Pi](https://pi.dev) extension that automatically injects relevant project documentation into the LLM context by monitoring streaming output for keyword matches. Docs are delivered as a `CustomMessage` so the system prompt stays untouched and the provider's prompt cache stays warm.
 
 ## Installation
 
@@ -29,8 +29,9 @@ git clone https://github.com/yourname/pi-doc-injector.git .pi/extensions/doc-inj
 1. Create a `docs/` folder in your project root.
 2. Add markdown files with frontmatter (`title` + `keywords`). See [Document Format](#document-format) for supported formats.
 3. Start Pi. The extension scans `docs/` on session start.
-4. When the user mentions a keyword, the matching doc is injected into the
-   system prompt **before the assistant responds** — no one-turn delay.
+4. When the user mentions a keyword, the matching doc is injected as a
+   `CustomMessage` into the conversation **before the assistant responds** —
+   no one-turn delay. The system prompt is never modified.
 5. If the assistant mentions a NEW keyword mid-response, generation is
    automatically aborted and restarted with the doc injected immediately.
 
@@ -208,18 +209,41 @@ The extension uses a per-session injection model:
 - **Assistant streaming**: if the assistant mentions a NEW keyword mid-response,
   generation is aborted and restarted with the doc injected immediately.
 
-### System Prompt Lifecycle
+### Injection Mechanism
 
-Pi **reconstructs the system prompt from source files each turn** (verified against pi v0.70.6).
+On match, the extension returns a `message` field from `before_agent_start`
+with `customType: "doc-injector"`. Pi appends this to the session and sends
+it to the LLM as part of the conversation. The system prompt is **never**
+mutated.
 
-When `before_agent_start` fires, the `systemPrompt` passed to the extension is a freshly rebuilt prompt from `AGENTS.md`, `SYSTEM.md`, skills, and tool snippets. It is **not** accumulated from previous turns.
+#### Why a CustomMessage, not the system prompt?
 
-This means:
+- The system prompt is the highest-value prompt-cache slot. Each unique
+  system prompt text breaks the cache (5-min TTL by default). Appending
+  per-turn doc content there would invalidate the cache on every first
+  injection.
+- A `CustomMessage` only adds to the conversation prefix, leaving the
+  system prompt byte-identical across turns and the cache warm.
 
-- Injections apply to the **current turn only** and do not persist in subsequent turns.
-- There is no risk of duplicate injection sections stacking up over time.
-- The `injected` flag alone is sufficient to prevent re-injection — no additional deduplication or marker-based stripping is needed.
+#### Double-injection prevention
 
+Two independent guards make duplicate injection impossible in a session:
+
+1. **Matcher guard** — `buildMatcher()` only includes non-injected entries
+   (via `getNonInjectedEntries()`), so already-injected docs cannot be
+   re-matched.
+2. **Mark guard** — `markInjected()` runs inside `before_agent_start` before
+   the LLM call, so even if the matcher ever produced a duplicate, the
+   mark would still prevent a second send.
+
+In practice, the matcher guard is the primary defense; the mark guard is
+defense-in-depth for race conditions (e.g. if `resources_discover` rebuilds
+the registry mid-injection).
+
+The `injected` flag is per-session: it's reset on `session_start` and can
+be manually cleared with `/doc-inject reset`.
+
+For the full source-level verification, see the JSDoc block in `index.ts`.
 For the full source-level verification, see the JSDoc block in `index.ts`.
 
 ## Development

package/cache.ts

@@ -9,6 +9,7 @@
 import { mkdir, readFile, writeFile } from "node:fs/promises";
 import { dirname, join } from "node:path";
 import type { KeywordCache } from "./types";
+import type { Notifier } from "./notifier";
 
 const CACHE_FILENAME = ".pi/doc-injector-cache.json";
 const CACHE_VERSION = 1;
@@ -17,8 +18,11 @@ const CACHE_VERSION = 1;
  * Load the keyword cache from disk.
  * Returns an empty cache (version 1, no files) if the file doesn't exist,
  * has wrong version, or is corrupted.
+ *
+ * Recoverable issues (corrupt JSON, wrong version) emit a warning via the
+ * `notifier`. ENOENT (no cache file yet) is silent.
  */
-export async function loadCache(cwd: string): Promise<KeywordCache> {
+export async function loadCache(cwd: string, notifier: Notifier): Promise<KeywordCache> {
   const cachePath = join(cwd, CACHE_FILENAME);
 
   try {
@@ -26,7 +30,7 @@ export async function loadCache(cwd: string): Promise<KeywordCache> {
     const parsed: unknown = JSON.parse(raw);
 
     if (!isValidCache(parsed)) {
-      console.warn(
+      notifier.warn(
         `[doc-injector] Invalid cache format or version at ${cachePath}, resetting.`,
       );
       return emptyCache();
@@ -36,10 +40,8 @@ export async function loadCache(cwd: string): Promise<KeywordCache> {
   } catch (err) {
     // ENOENT = no cache file yet, that's fine
     if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
-      console.warn(
-        `[doc-injector] Failed to read cache at ${cachePath}:`,
-        err instanceof Error ? err.message : String(err),
-      );
+      const detail = err instanceof Error ? err.message : String(err);
+      notifier.warn(`[doc-injector] Failed to read cache at ${cachePath}: ${detail}`);
     }
     return emptyCache();
   }

package/config.ts

@@ -5,47 +5,56 @@
 import { readFile } from "node:fs/promises";
 import { join } from "node:path";
 import { DEFAULT_CONFIG, type DocInjectorConfig } from "./types";
+import type { Notifier } from "./notifier";
 
 /**
  * Clamp an integer value to [min, max] range.
- * Warns and clamps if out of range. Returns the default if not a number.
+ * Warns via the `notifier` and clamps if out of range. Returns the default
+ * if not a number.
  */
 function clampInt(
-  value: unknown,
-  defaultVal: number,
-  min: number,
-  max: number,
-  fieldName: string,
+    value: unknown,
+    defaultVal: number,
+    min: number,
+    max: number,
+    fieldName: string,
+    notifier: Notifier,
 ): number {
-  if (typeof value !== "number" || Number.isNaN(value)) {
-    return defaultVal;
-  }
-  const intVal = Math.trunc(value);
-  if (intVal < min || intVal > max) {
-    const clamped = Math.max(min, Math.min(max, intVal));
-    console.warn(`[doc-injector] ${fieldName} must be ${min}-${max}, got ${intVal}. Clamping to ${clamped}.`);
-    return clamped;
-  }
-  return intVal;
+    if (typeof value !== "number" || Number.isNaN(value)) {
+        return defaultVal;
+    }
+    const intVal = Math.trunc(value);
+    if (intVal < min || intVal > max) {
+        const clamped = Math.max(min, Math.min(max, intVal));
+        notifier.warn(`[doc-injector] ${fieldName} must be ${min}-${max}, got ${intVal}. Clamping to ${clamped}.`);
+        return clamped;
+    }
+    return intVal;
 }
 
+/**
 /**
  * Validate a glob pattern array.
  * Rejects non-array or entries that aren't strings. Returns default on error.
+ * Warns via the `notifier` for non-string entries.
  */
-function validateGlobArray(value: unknown, defaultVal: string[]): string[] {
-  if (!Array.isArray(value)) {
-    return [...defaultVal];
-  }
-  const result: string[] = [];
-  for (const item of value) {
-    if (typeof item === "string") {
-      result.push(item);
-    } else {
-      console.warn(`[doc-injector] Non-string entry in glob array ignored: ${String(item)}`);
+function validateGlobArray(
+    value: unknown,
+    defaultVal: string[],
+    notifier: Notifier,
+): string[] {
+    if (!Array.isArray(value)) {
+        return [...defaultVal];
+    }
+    const result: string[] = [];
+    for (const item of value) {
+        if (typeof item === "string") {
+            result.push(item);
+        } else {
+            notifier.warn(`[doc-injector] Non-string entry in glob array ignored: ${String(item)}`);
+        }
     }
-  }
-  return result.length > 0 ? result : [...defaultVal];
+    return result.length > 0 ? result : [...defaultVal];
 }
 
 /**
@@ -54,33 +63,37 @@ function validateGlobArray(value: unknown, defaultVal: string[]): string[] {
  * Validates and clamps all numeric fields. Falls back to DEFAULT_CONFIG
  * if file doesn't exist or is invalid.
  */
-export async function loadConfig(cwd: string): Promise<DocInjectorConfig> {
-  const configPath = join(cwd, ".pi", "doc-injector.json");
+/**
+ * Load config from `.pi/doc-injector.json` relative to the given cwd.
+ * Async — uses readFile from fs/promises. Validates and clamps all numeric
+ * fields. Falls back to DEFAULT_CONFIG if the file doesn't exist or is
+ * invalid. Warnings (clamping, invalid entries) go through the `notifier`.
+ */
+export async function loadConfig(cwd: string, notifier: Notifier): Promise<DocInjectorConfig> {
+    const configPath = join(cwd, ".pi", "doc-injector.json");
 
-  try {
-    const raw = await readFile(configPath, "utf-8");
-    const parsed = JSON.parse(raw) as Partial<DocInjectorConfig>;
+    try {
+        const raw = await readFile(configPath, "utf-8");
+        const parsed = JSON.parse(raw) as Partial<DocInjectorConfig>;
 
-    return {
-      docsPath: parsed.docsPath ?? DEFAULT_CONFIG.docsPath,
-      matchThreshold: clampInt(parsed.matchThreshold, DEFAULT_CONFIG.matchThreshold, 1, Infinity, "matchThreshold"),
-      contextThreshold: clampInt(parsed.contextThreshold, DEFAULT_CONFIG.contextThreshold, 0, 100, "contextThreshold"),
-      recursive: parsed.recursive ?? DEFAULT_CONFIG.recursive,
-      include: validateGlobArray(parsed.include, DEFAULT_CONFIG.include),
-      exclude: validateGlobArray(parsed.exclude, DEFAULT_CONFIG.exclude),
-      maxFileSize: clampInt(parsed.maxFileSize, DEFAULT_CONFIG.maxFileSize, 1024, 10 * 1024 * 1024, "maxFileSize"),
-      autoKeywords: parsed.autoKeywords ?? DEFAULT_CONFIG.autoKeywords,
-      llmKeywords: parsed.llmKeywords ?? DEFAULT_CONFIG.llmKeywords,
-      maxConcurrent: clampInt(parsed.maxConcurrent, DEFAULT_CONFIG.maxConcurrent, 1, 100, "maxConcurrent"),
-      llmBatchSize: clampInt(parsed.llmBatchSize, DEFAULT_CONFIG.llmBatchSize, 1, 100, "llmBatchSize"),
-    };
-  } catch (err) {
-    if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
-      console.warn(
-        `[doc-injector] Failed to parse config at ${configPath}:`,
-        err instanceof Error ? err.message : String(err),
-      );
+        return {
+            docsPath: parsed.docsPath ?? DEFAULT_CONFIG.docsPath,
+            matchThreshold: clampInt(parsed.matchThreshold, DEFAULT_CONFIG.matchThreshold, 1, Infinity, "matchThreshold", notifier),
+            contextThreshold: clampInt(parsed.contextThreshold, DEFAULT_CONFIG.contextThreshold, 0, 100, "contextThreshold", notifier),
+            recursive: parsed.recursive ?? DEFAULT_CONFIG.recursive,
+            include: validateGlobArray(parsed.include, DEFAULT_CONFIG.include, notifier),
+            exclude: validateGlobArray(parsed.exclude, DEFAULT_CONFIG.exclude, notifier),
+            maxFileSize: clampInt(parsed.maxFileSize, DEFAULT_CONFIG.maxFileSize, 1024, 10 * 1024 * 1024, "maxFileSize", notifier),
+            autoKeywords: parsed.autoKeywords ?? DEFAULT_CONFIG.autoKeywords,
+            llmKeywords: parsed.llmKeywords ?? DEFAULT_CONFIG.llmKeywords,
+            maxConcurrent: clampInt(parsed.maxConcurrent, DEFAULT_CONFIG.maxConcurrent, 1, 100, "maxConcurrent", notifier),
+            llmBatchSize: clampInt(parsed.llmBatchSize, DEFAULT_CONFIG.llmBatchSize, 1, 100, "llmBatchSize", notifier),
+        };
+    } catch (err) {
+        if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
+            const detail = err instanceof Error ? err.message : String(err);
+            notifier.warn(`[doc-injector] Failed to parse config at ${configPath}: ${detail}`);
+        }
+        return { ...DEFAULT_CONFIG };
     }
-    return { ...DEFAULT_CONFIG };
-  }
 }

package/index.ts

@@ -4,6 +4,21 @@
  * Automatically injects relevant project documentation into the LLM context
  * by monitoring streaming output for keyword matches.
  *
+ * ## Injection Model: CustomMessage (NOT system prompt)
+ *
+ * On match, the extension returns a `message` field from `before_agent_start`
+ * (a `CustomMessage` with `customType: "doc-injector"`). Pi appends this to the
+ * session and sends it to the LLM as part of the conversation — the system
+ * prompt is NEVER mutated.
+ *
+ * Why a message and not the system prompt:
+ * - The system prompt is the highest-value Anthropic prompt-cache slot. Each
+ *   unique system prompt text breaks the cache (5-min TTL by default).
+ *   Appending per-turn doc content there would invalidate the cache on every
+ *   first injection.
+ * - A `message` only adds to the conversation prefix, leaving the system
+ *   prompt cache warm across turns.
+ *
  * ## Streaming Model
  *
  * This extension relies on Pi's streaming event contract:
@@ -13,7 +28,7 @@
  * - `message_end`: Fires once when the assistant's response is complete.
  *   The extension finalizes matches and notifies the user.
  * - `before_agent_start`: Fires before the next agent turn. The extension
- *   injects matched docs into the system prompt, then marks them as injected.
+ *   returns a `message` carrying the matched docs and marks them as injected.
  *
  * ## Injection Lifecycle
  *
@@ -22,28 +37,23 @@
  * session, once a doc is injected, it won't be re-injected unless the user
  * manually runs `/doc-inject reset`.
  *
- * ## System Prompt Lifecycle (verified against pi v0.70.6)
+ * ## Double-Injection Prevention
  *
- * Pi **reconstructs the system prompt from source files each turn**. Here is
- * the exact flow, verified via source-code review of dist/core/agent-session.js
- * and dist/core/extensions/runner.js (v0.70.6):
+ * Two independent guards make duplicate injection impossible in a session:
  *
- * 1. Before each agent turn, pi calls `this._rebuildSystemPrompt(toolNames)`.
- *    This builds the prompt from `AGENTS.md`, `SYSTEM.md`, skills, enabled
- *    tool snippets — never from a previously modified (injected) prompt.
- * 2. The rebuilt prompt is stored in `this._baseSystemPrompt`.
- * 3. `emitBeforeAgentStart(..., this._baseSystemPrompt, ...)` passes this
- *    *fresh* base prompt to every extension handler.
- * 4. Extension handlers can return a modified `systemPrompt` for the current
- *    turn. Pi uses the modified prompt **only for this turn**.
- * 5. When no extension modifies the prompt, pi explicitly resets to
- *    `this._baseSystemPrompt` (comment in source: "Ensure we're using the
- *    base prompt (in case previous turn had modifications)").
+ * 1. **Matcher-level guard**: `buildMatcher()` calls `getNonInjectedEntries()`,
+ *    so already-injected docs are excluded from the candidate set. The
+ *    `pendingMatches` map is only populated from the matcher's output, so
+ *    once a doc is injected, the next `input` event cannot re-match it.
  *
- * **Therefore**: Previous injections from `before_agent_start` do NOT persist
- * across turns. Duplicate sections cannot accumulate in the system prompt.
- * The `injected` flag alone is sufficient to prevent re-injection — no
- * marker-based stripping or deduplication is needed.
+ * 2. **Mark guard**: `markInjected()` is called inside `before_agent_start`
+ *    AFTER the build step but BEFORE the return value is processed. This
+ *    means the flag flips synchronously with the LLM call — even if the
+ *    session is reloaded mid-turn, the next `buildMatcher()` won't see the
+ *    doc as a candidate.
+ *
+ * The two guards are redundant by design: if matcher exclusion ever fails
+ * (e.g. a race), the mark step still prevents the doc from being sent twice.
  *
  * ## Race Condition Note
  *
@@ -58,16 +68,22 @@ import { Type } from "@sinclair/typebox";
 import { resolve } from "node:path";
 import { loadCache, saveCache } from "./cache";
 import { loadConfig } from "./config";
-import { buildSystemPromptAppend, notifyInjection } from "./injector";
+import { buildInjectionContent, notifyInjection } from "./injector";
 import { buildKeywordGenPrompt } from "./keyword-llm";
 import { extractText, KeywordMatcher } from "./matcher";
+import { ExtensionNotifier, type Notifier } from "./notifier";
 import { DocRegistry } from "./registry";
 import { DEFAULT_MATCHER_OPTIONS, type DocEntry, type MatchResult, type KeywordCache, type CacheEntry } from "./types";
 import { registerCommands } from "./commands";
 
 export default async function docInjectorExtension(pi: ExtensionAPI) {
   // ---- State ----
-  let config = await loadConfig(process.cwd());
+  // The notifier buffers warnings emitted during startup (loadConfig,
+  // loadCache, initRegistry) and flushes them via ctx.ui.notify() in
+  // session_start. The notifier is bound to the extension lifecycle so
+  // startup messages aren't lost.
+  const notifier: Notifier = new ExtensionNotifier();
+  let config = await loadConfig(process.cwd(), notifier);
   let registry: DocRegistry | null = null;
   let initRegistryPromise: Promise<void> | null = null;
   let enabled = true;
@@ -92,7 +108,7 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   const safeSaveCache = async (cwd: string, dirtyEntries: Record<string, CacheEntry>) => {
     // MAJOR-2 fix: before saveCache, re-read cache from disk to merge
     // LLM-written entries that may have landed during the scan.
-    const freshCache = await loadCache(cwd);
+    const freshCache = await loadCache(cwd, notifier);
     const mergedCache: KeywordCache = { version: 1, files: {} };
 
     // Start with fresh (disk) entries — includes any LLM writes during scan
@@ -109,10 +125,10 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   };
 
   const initRegistry = async (cwd: string) => {
-    config = await loadConfig(cwd);
+    config = await loadConfig(cwd, notifier);
     const docsPath = resolve(cwd, config.docsPath);
-    cache = await loadCache(cwd);
-    registry = await DocRegistry.create(docsPath, config, cache);
+    cache = await loadCache(cwd, notifier);
+    registry = await DocRegistry.create(docsPath, config, cache, notifier);
 
     const dirty = registry.getDirtyCache();
     if (Object.keys(dirty).length > 0) {
@@ -193,6 +209,12 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     llmBatchesCompleted = 0;
     llmTotalFiles = 0;
 
+    // Bind the notifier to the live context FIRST so any warnings emitted
+    // during initRegistry below go directly to the TUI instead of being
+    // buffered. Messages buffered from earlier (e.g. the factory-body
+    // loadConfig call) are flushed here in arrival order.
+    notifier.setContext(ctx);
+
     if (event.reason === "reload") return;
 
     if (initRegistryPromise) {
@@ -213,7 +235,7 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     const effectiveCwd = cwd ?? process.cwd();
 
     // Reload cache from disk to pick up LLM-generated entries
-    const freshCache = await loadCache(effectiveCwd);
+    const freshCache = await loadCache(effectiveCwd, notifier);
     cache = freshCache;
     registry.updateCache(cache);
 
@@ -308,7 +330,11 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     textBuffer = "";
   });
 
-  // ---- Event: before_agent_start (inject into system prompt) ----
+  // ---- Event: before_agent_start (inject as CustomMessage) ----
+  // Returns a `message` (CustomMessage with customType: "doc-injector") rather
+  // than mutating `systemPrompt`. The system prompt stays byte-identical across
+  // turns, preserving the prompt cache. The CustomMessage is appended to the
+  // session and sent to the LLM as part of the conversation.
   pi.on("before_agent_start", async (event, ctx) => {
     // P5.4b — Guard: skip injection during LLM keyword generation
     if (keywordGenInFlight) return;
@@ -335,9 +361,12 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
       return;
     }
 
-    const append = buildSystemPromptAppend(matchedEntries, pendingMatches);
+    const content = buildInjectionContent(matchedEntries, pendingMatches);
 
-    // Mark as injected only after confirming injection will happen
+    // Mark as injected only after confirming injection will happen.
+    // This is the second half of the double-injection guard: even if the
+    // matcher ever produced a duplicate match, markInjected prevents a
+    // second send.
     registry.markInjected(matchedEntries.map((e) => e.filePath));
 
     // Notify user about injection (moved here from message_end so it fires
@@ -348,7 +377,11 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     pendingMatches.clear();
 
     return {
-      systemPrompt: (event.systemPrompt || "") + "\n\n" + append,
+      message: {
+        customType: "doc-injector",
+        content,
+        display: true,
+      },
     };
   });
 

package/injector.ts

@@ -1,8 +1,12 @@
 /**
- * Context Injector — formats matched docs into system prompt append
- * and sends TUI notifications.
+ * Context Injector — formats matched docs into a content string suitable for
+ * injection as a `CustomMessage` (returned from `before_agent_start`) and
+ * sends TUI notifications.
+ *
+ * The produced content is delivered to the LLM as a `CustomMessage` rather
+ * than appended to the system prompt. This keeps the system prompt
+ * byte-identical across turns so the provider's prompt cache stays warm.
  */
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import type { DocEntry } from "./types";
 
 /**
@@ -14,7 +18,7 @@ export interface NotifyCapability {
 }
 
 /**
- * Sanitize keywords for safe injection into the system prompt.
+ * Sanitize keywords for safe display in the injection content.
  *
  * - Strips \n and \r (replaces with space) to prevent prompt injection
  * - Caps each keyword at 100 characters
@@ -29,11 +33,13 @@ function sanitizeKeywords(keywords: string[]): string[] {
 }
 
 /**
- * Build a system prompt append string from matched documents.
+ * Build the content string for a `CustomMessage` injection from matched
+ * documents. This is the payload that gets returned in
+ * `before_agent_start`'s `message.content` and sent to the LLM.
  */
-export function buildSystemPromptAppend(
-  entries: DocEntry[],
-  matchedKeywords: Map<string, string[]>,
+export function buildInjectionContent(
+    entries: DocEntry[],
+    matchedKeywords: Map<string, string[]>,
 ): string {
   if (entries.length === 0) return "";
 

package/notifier.ts

@@ -0,0 +1,71 @@
+/**
+ * Notifier — thin wrapper around Pi's `ctx.ui.notify()` that buffers
+ * messages until a context is available.
+ *
+ * ## Why a buffer?
+ *
+ * Several warnings fire at startup (during `loadConfig` and `initRegistry`),
+ * before any `ExtensionContext` exists — extensions are constructed first,
+ * events fire later. The `Notifier` interface accepts messages at any time:
+ *
+ * - If a context has been set, messages are forwarded to `ctx.ui.notify()`.
+ * - If not, messages are buffered in memory and flushed on the next
+ *   `setContext()` call (typically from `session_start`).
+ *
+ * Production code uses `ExtensionNotifier`. Tests inject a plain object
+ * satisfying the `Notifier` interface (or a `vi.fn()` spy) — no real
+ * extension context is needed.
+ */
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+
+export type NotifierLevel = "info" | "warning" | "error";
+
+export interface Notifier {
+    /** Show an informational message. */
+    info(message: string): void;
+    /** Show a warning. */
+    warn(message: string): void;
+    /** Show an error. */
+    error(message: string): void;
+    /**
+     * Bind a context. Flushes any buffered messages via `ctx.ui.notify()`
+     * in arrival order. Idempotent: re-calling replaces the context and
+     * clears the buffer (already-flushed messages are not re-sent).
+     */
+    setContext(ctx: ExtensionContext): void;
+}
+
+/** Production notifier. Buffers until a context is bound. */
+export class ExtensionNotifier implements Notifier {
+    private ctx: ExtensionContext | null = null;
+    private buffer: Array<{ level: NotifierLevel; message: string }> = [];
+
+    setContext(ctx: ExtensionContext): void {
+        this.ctx = ctx;
+        const pending = this.buffer;
+        this.buffer = [];
+        for (const { level, message } of pending) {
+            ctx.ui.notify(message, level);
+        }
+    }
+
+    info(message: string): void {
+        this.emit("info", message);
+    }
+
+    warn(message: string): void {
+        this.emit("warning", message);
+    }
+
+    error(message: string): void {
+        this.emit("error", message);
+    }
+
+    private emit(level: NotifierLevel, message: string): void {
+        if (this.ctx) {
+            this.ctx.ui.notify(message, level);
+        } else {
+            this.buffer.push({ level, message });
+        }
+    }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-doc-injector",
-  "version": "0.4.0",
+  "version": "0.5.1",
   "description": "Auto-inject relevant project documentation into Pi's LLM context based on keyword matching",
   "type": "module",
   "main": "./index.ts",

package/registry.ts

@@ -9,6 +9,7 @@ import type { Dirent } from "node:fs";
 import { readdir, readFile, stat } from "node:fs/promises";
 import { basename, extname, join, relative, resolve } from "node:path";
 import type { CacheEntry, DocEntry, DocInjectorConfig, KeywordCache } from "./types";
+import type { Notifier } from "./notifier";
 import { createGlobFilter } from "./globber";
 import { generateKeywords } from "./keyword-gen";
 
@@ -225,28 +226,36 @@ class PromisePool {
  * Document Registry class. Scans a docs folder and maintains an index of DocEntry.
  */
 export class DocRegistry {
-  private entries: DocEntry[] = [];
-  private docsPath: string;
-  private config: DocInjectorConfig;
-  private cache: KeywordCache | null = null;
-  private dirtyCache: KeywordCache = { version: 1, files: {} };
-
-  private constructor(docsPath: string, config: DocInjectorConfig, cache?: KeywordCache) {
-    this.docsPath = docsPath;
-    this.config = config;
-    this.cache = cache ?? null;
-  }
+    private entries: DocEntry[] = [];
+    private docsPath: string;
+    private config: DocInjectorConfig;
+    private cache: KeywordCache | null = null;
+    private dirtyCache: KeywordCache = { version: 1, files: {} };
+    private notifier: Notifier;
+
+    private constructor(
+        docsPath: string,
+        config: DocInjectorConfig,
+        cache: KeywordCache | undefined,
+        notifier: Notifier,
+    ) {
+        this.docsPath = docsPath;
+        this.config = config;
+        this.cache = cache ?? null;
+        this.notifier = notifier;
+    }
 
-  /** Create a registry by scanning the docs folder. */
-  static async create(
-    docsPath: string,
-    config: DocInjectorConfig,
-    cache?: KeywordCache,
-  ): Promise<DocRegistry> {
-    const registry = new DocRegistry(docsPath, config, cache);
-    await registry.rebuild();
-    return registry;
-  }
+    /** Create a registry by scanning the docs folder. */
+    static async create(
+        docsPath: string,
+        config: DocInjectorConfig,
+        cache: KeywordCache | undefined,
+        notifier: Notifier,
+    ): Promise<DocRegistry> {
+        const registry = new DocRegistry(docsPath, config, cache, notifier);
+        await registry.rebuild();
+        return registry;
+    }
 
   /** Re-scan the docs folder and rebuild the index. */
   async rebuild(): Promise<void> {
@@ -274,7 +283,7 @@ export class DocRegistry {
       const results = await pool.all(tasks);
       this.entries = results.filter((e): e is DocEntry => e !== null);
     } catch {
-      console.warn(`[doc-injector] Docs folder not found: ${resolved}`);
+      this.notifier.warn(`[doc-injector] Docs folder not found: ${resolved}`);
       this.entries = [];
     }
   }
@@ -295,7 +304,7 @@ export class DocRegistry {
 
       // Step 2: Skip files exceeding maxFileSize
       if (fileStat.size > this.config.maxFileSize) {
-        console.warn(
+        this.notifier.warn(
           `[doc-injector] Skipping ${relativePath}: size ${fileStat.size} > max ${this.config.maxFileSize}`,
         );
         return null;
@@ -346,7 +355,7 @@ export class DocRegistry {
         keywordSource = "heuristic";
       } else {
         // Step 11: No frontmatter and autoKeywords disabled — skip
-        console.warn(
+        this.notifier.warn(
           `[doc-injector] Skipping ${relativePath}: no valid frontmatter with keywords`,
         );
         return null;
@@ -373,7 +382,7 @@ export class DocRegistry {
     } catch (err) {
       // Only warn for unexpected errors, not ENOENT (file deleted/moved after scan)
       if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
-        console.warn(`[doc-injector] Error reading ${relativePath}:`, err);
+        this.notifier.warn(`[doc-injector] Error reading ${relativePath}: ${err instanceof Error ? err.message : String(err)}`);
       }
       return null;
     }