pi-doc-injector 0.2.0 → 0.3.0

package/README.md

@@ -101,11 +101,18 @@ Injection is also skipped if the current context usage exceeds 80% of the token
 
 The extension uses a per-session injection model:
 
-- On `session_start`, the registry is rebuilt from scratch, resetting all `injected` flags.
+- On `session_start`, the registry scans `docs/` and indexes all valid documents.
 - Within a session, once a document is injected, it won't be re-injected automatically.
 - Use `/doc-inject reset` to manually reset all flags and allow docs to be injected again.
 - Use `/doc-inject list` to see which docs have been injected (✅) and which are pending (⬜).
 
+### Injection Timing
+
+- **User messages**: matched via the `input` event, injected before the assistant
+  responds — **same turn**, no delay.
+- **Assistant streaming**: if the assistant mentions a NEW keyword mid-response,
+  generation is aborted and restarted with the doc injected immediately.
+
 ### System Prompt Lifecycle
 
 Pi **reconstructs the system prompt from source files each turn** (verified against pi v0.70.6).

package/cache.ts

@@ -0,0 +1,79 @@
+/**
+ * Keyword cache persistence — load/save the `.pi/doc-injector-cache.json` file.
+ *
+ * Cache format:
+ *   { version: 1, files: { [relativePath]: { mtimeMs: number, keywords: string[] } } }
+ *
+ * Invalid files (wrong version, bad JSON, ENOENT) result in an empty cache.
+ */
+import { mkdir, readFile, writeFile } from "node:fs/promises";
+import { dirname, join } from "node:path";
+import type { KeywordCache } from "./types";
+
+const CACHE_FILENAME = ".pi/doc-injector-cache.json";
+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.
+ */
+export async function loadCache(cwd: string): Promise<KeywordCache> {
+  const cachePath = join(cwd, CACHE_FILENAME);
+
+  try {
+    const raw = await readFile(cachePath, "utf-8");
+    const parsed: unknown = JSON.parse(raw);
+
+    if (!isValidCache(parsed)) {
+      console.warn(
+        `[doc-injector] Invalid cache format or version at ${cachePath}, resetting.`,
+      );
+      return emptyCache();
+    }
+
+    return parsed;
+  } 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),
+      );
+    }
+    return emptyCache();
+  }
+}
+
+/**
+ * Save the keyword cache to disk.
+ * Creates parent directories if needed.
+ */
+export async function saveCache(
+  cwd: string,
+  cache: KeywordCache,
+): Promise<void> {
+  const cachePath = join(cwd, CACHE_FILENAME);
+
+  try {
+    await mkdir(dirname(cachePath), { recursive: true });
+  } catch {
+    // Ignore — directory may already exist
+  }
+
+  await writeFile(cachePath, JSON.stringify(cache, null, 2), "utf-8");
+}
+
+/** Check that a parsed value matches the KeywordCache shape. */
+function isValidCache(value: unknown): value is KeywordCache {
+  if (!value || typeof value !== "object") return false;
+  const c = value as Record<string, unknown>;
+  if (c.version !== CACHE_VERSION) return false;
+  if (!c.files || typeof c.files !== "object") return false;
+  return true;
+}
+
+/** Return a fresh empty cache. */
+function emptyCache(): KeywordCache {
+  return { version: CACHE_VERSION, files: {} };
+}

package/commands.ts

@@ -3,14 +3,26 @@
  */
 import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
 import type { DocRegistry } from "./registry";
+import type { DocInjectorConfig } from "./types";
 
+/** Dependencies injected into the command registrar. */
 export interface CommandDeps {
   getRegistry: () => DocRegistry | null;
   getEnabled: () => boolean;
   setEnabled: (v: boolean) => void;
   reloadRegistry: () => Promise<number>;
+  getConfig: () => DocInjectorConfig;
+  generateKeywordsLLM: (files: Array<{ path: string; snippet: string; existingKeywords: string[] }>) => Promise<void>;
 }
 
+/**
+ * Register all doc-injector slash commands on the given ExtensionAPI.
+ *
+ * Commands:
+ * - `/doc-inject [on|off|toggle|list|reset|status]` — manage injection state
+ * - `/doc-reload` — re-scan docs folder
+ * - `/doc-keywords-gen [path]` — generate LLM keywords for keyword-less files
+ */
 export function registerCommands(pi: ExtensionAPI, deps: CommandDeps): void {
   const cmd = (name: string, desc: string, handler: (args: string, ctx: ExtensionContext) => Promise<void>) => {
     pi.registerCommand(name, { description: desc, handler });
@@ -49,7 +61,8 @@ export function registerCommands(pi: ExtensionAPI, deps: CommandDeps): void {
       }
       const lines = entries.map((e) => {
         const status = e.injected ? "✅" : "⬜";
-        return `${status} ${e.relativePath}: "${e.title}" — keywords: [${e.keywords.join(", ")}]`;
+        const sourceTag = `[${e.keywordSource}]`;
+        return `${status} ${sourceTag} ${e.relativePath}: "${e.title}" — keywords: [${e.keywords.join(", ")}]`;
       });
       ctx.ui.notify(`📄 Registered docs:\n${lines.join("\n")}`, "info");
     } else {
@@ -81,4 +94,58 @@ export function registerCommands(pi: ExtensionAPI, deps: CommandDeps): void {
       ctx.ui.notify(`📄 Reload failed: ${err instanceof Error ? err.message : String(err)}`, "error");
     }
   });
+
+  cmd("doc-keywords-gen", "Generate LLM keywords: /doc-keywords-gen [path] — no arg = all keyword-less files", async (args, ctx) => {
+    const reg = deps.getRegistry();
+    if (!reg) {
+      ctx.ui.notify("📄 No registry loaded", "warning");
+      return;
+    }
+
+    const config = deps.getConfig();
+    if (!config.llmKeywords) {
+      ctx.ui.notify("📄 LLM keyword generation is disabled (llmKeywords: false in config)", "warning");
+      return;
+    }
+
+    const targetPath = args.trim();
+
+    // Filter to keyword-less entries (keywordSource !== "frontmatter", "cache", or "llm")
+    let candidates = reg.getEntries().filter((e) => {
+      if (e.keywordSource === "frontmatter") return false;
+      if (e.keywordSource === "cache") return false;
+      if (e.keywordSource === "llm") return false; // already LLM-generated
+      return true;
+    });
+
+    if (targetPath) {
+      candidates = candidates.filter((e) => e.relativePath.includes(targetPath));
+      if (candidates.length === 0) {
+        ctx.ui.notify(`📄 No keyword-less files matching "${targetPath}"`, "info");
+        return;
+      }
+    }
+
+    if (candidates.length === 0) {
+      ctx.ui.notify("📄 All files already have keywords", "info");
+      return;
+    }
+
+    const batchSize = config.llmBatchSize;
+    const batches: Array<Array<{ path: string; snippet: string; existingKeywords: string[] }>> = [];
+    for (let i = 0; i < candidates.length; i += batchSize) {
+      const batch = candidates.slice(i, i + batchSize).map((e) => ({
+        path: e.relativePath,
+        snippet: e.content.slice(0, 500),
+        existingKeywords: e.keywords,
+      }));
+      batches.push(batch);
+    }
+
+    ctx.ui.notify(`📄 Sending ${batches.length} keyword-generation batch(es) for ${candidates.length} file(s)...`, "info");
+
+    for (const batch of batches) {
+      await deps.generateKeywordsLLM(batch);
+    }
+  });
 }

package/config.ts

@@ -2,50 +2,85 @@
  * Configuration loader for the Doc Injector extension.
  * Reads from `.pi/doc-injector.json` with fallback to defaults.
  */
-import { existsSync, readFileSync } from "node:fs";
+import { readFile } from "node:fs/promises";
 import { join } from "node:path";
 import { DEFAULT_CONFIG, type DocInjectorConfig } from "./types";
 
 /**
- * Load config from `.pi/doc-injector.json` relative to the given cwd.
- * Falls back to DEFAULT_CONFIG if file doesn't exist or is invalid.
+ * Clamp an integer value to [min, max] range.
+ * Warns and clamps if out of range. Returns the default if not a number.
  */
-export function loadConfig(cwd: string): DocInjectorConfig {
-  const configPath = join(cwd, ".pi", "doc-injector.json");
+function clampInt(
+  value: unknown,
+  defaultVal: number,
+  min: number,
+  max: number,
+  fieldName: string,
+): 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 (!existsSync(configPath)) {
-    return { ...DEFAULT_CONFIG };
+/**
+ * Validate a glob pattern array.
+ * Rejects non-array or entries that aren't strings. Returns default on error.
+ */
+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)}`);
+    }
+  }
+  return result.length > 0 ? result : [...defaultVal];
+}
+
+/**
+ * Load config from `.pi/doc-injector.json` relative to the given cwd.
+ * Now async — uses readFile from fs/promises.
+ * 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");
 
   try {
-    const raw = readFileSync(configPath, "utf-8");
+    const raw = await readFile(configPath, "utf-8");
     const parsed = JSON.parse(raw) as Partial<DocInjectorConfig>;
 
-    // Clamp contextThreshold to 0-100 range
-    let contextThreshold = parsed.contextThreshold ?? DEFAULT_CONFIG.contextThreshold;
-    if (typeof contextThreshold === "number" && (contextThreshold < 0 || contextThreshold > 100)) {
-      console.warn(`[doc-injector] contextThreshold must be 0-100, got ${contextThreshold}. Clamping.`);
-      contextThreshold = Math.max(0, Math.min(100, contextThreshold));
-    }
-
-    // Clamp matchThreshold to positive integers
-    let matchThreshold = parsed.matchThreshold ?? DEFAULT_CONFIG.matchThreshold;
-    if (typeof matchThreshold === "number" && matchThreshold < 1) {
-      console.warn(`[doc-injector] matchThreshold must be >= 1, got ${matchThreshold}. Using 1.`);
-      matchThreshold = 1;
-    }
-
     return {
       docsPath: parsed.docsPath ?? DEFAULT_CONFIG.docsPath,
-      matchThreshold,
-      contextThreshold,
+      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) {
-    console.warn(
-      `[doc-injector] Failed to parse config at ${configPath}:`,
-      err instanceof Error ? err.message : String(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 { ...DEFAULT_CONFIG };
   }
 }

package/globber.ts

@@ -0,0 +1,48 @@
+/**
+ * Glob filter for include/exclude pattern matching.
+ * Uses picomatch (0 deps, ~18 KB) to compile patterns once for O(1) matching.
+ */
+import picomatch from "picomatch";
+import type { GlobFilter } from "./types";
+
+/**
+ * Create a glob filter from include and exclude patterns.
+ *
+ * A path matches if it matches at least one `include` pattern AND
+ * does not match any `exclude` pattern.
+ *
+ * When `include` is empty, all files are considered included
+ * (subject to exclude filtering).
+ *
+ * @param include - Glob patterns for files to include
+ * @param exclude - Glob patterns for files/dirs to exclude
+ * @returns A GlobFilter with a `match` method
+ */
+export function createGlobFilter(
+  include: string[],
+  exclude: string[],
+): GlobFilter {
+  const includeMatcher =
+    include.length > 0
+      ? picomatch(include, { dot: true })
+      : null;
+
+  const excludeMatcher =
+    exclude.length > 0
+      ? picomatch(exclude, { dot: true })
+      : null;
+
+  return {
+    match(relativePath: string): boolean {
+      // If include patterns are specified, path must match at least one
+      if (includeMatcher && !includeMatcher(relativePath)) {
+        return false;
+      }
+      // Path must not match any exclude pattern
+      if (excludeMatcher && excludeMatcher(relativePath)) {
+        return false;
+      }
+      return true;
+    },
+  };
+}

package/index.ts

@@ -53,35 +53,72 @@
  * is cleared after injection, and `markInjected()` operates on the registry's
  * current entries, not the stale array.
  */
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+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 { buildKeywordGenPrompt } from "./keyword-llm";
 import { extractText, KeywordMatcher } from "./matcher";
 import { DocRegistry } from "./registry";
-import { DEFAULT_MATCHER_OPTIONS, type DocEntry, type MatchResult } from "./types";
+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 = loadConfig(process.cwd());
+  let config = await loadConfig(process.cwd());
   let registry: DocRegistry | null = null;
+  let initRegistryPromise: Promise<void> | null = null;
   let enabled = true;
   let textBuffer = "";
   let pendingMatches = new Map<string, string[]>(); // filePath → matchedKeywords
   let abortingForInjection = false; // guard against cascading aborts
 
+  // P5.4b — Guard flags for LLM keyword generation
+  let keywordGenInFlight = false;
+  let llmBatchesCompleted = 0;
+  let llmTotalFiles = 0;
+  let cache: KeywordCache = { version: 1, files: {} };
+
   // ---- Helpers ----
   const getRegistry = () => registry;
   const getEnabled = () => enabled;
   const setEnabled = (v: boolean) => {
     enabled = v;
   };
+  const getConfig = () => config;
+
+  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 mergedCache: KeywordCache = { version: 1, files: {} };
+
+    // Start with fresh (disk) entries — includes any LLM writes during scan
+    for (const [key, entry] of Object.entries(freshCache.files)) {
+      mergedCache.files[key] = entry;
+    }
+
+    // Overlay dirty entries from this scan (scan results take precedence)
+    for (const [key, entry] of Object.entries(dirtyEntries)) {
+      mergedCache.files[key] = entry;
+    }
+
+    await saveCache(cwd, mergedCache);
+  };
 
   const initRegistry = async (cwd: string) => {
-    config = loadConfig(cwd);
+    config = await loadConfig(cwd);
     const docsPath = resolve(cwd, config.docsPath);
-    registry = await DocRegistry.create(docsPath, config.recursive);
+    cache = await loadCache(cwd);
+    registry = await DocRegistry.create(docsPath, config, cache);
+
+    const dirty = registry.getDirtyCache();
+    if (Object.keys(dirty).length > 0) {
+      await safeSaveCache(cwd, dirty);
+    }
+
     const count = registry.getEntries().length;
     if (count > 0) {
       console.log(`[doc-injector] Loaded ${count} documents from ${docsPath}`);
@@ -98,30 +135,109 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     );
   };
 
-  let lastInitTime = 0;
+  // P5.4f — generateKeywordsLLM: sets keywordGenInFlight and sends a user message
+  // with the prompt built by buildKeywordGenPrompt. The LLM will respond by
+  // calling the _doc_injector_keywords tool.
+  const generateKeywordsLLM = async (
+    files: Array<{ path: string; snippet: string; existingKeywords: string[] }>,
+  ) => {
+    keywordGenInFlight = true;
+    const prompt = buildKeywordGenPrompt(files);
+    pi.sendUserMessage(prompt, { deliverAs: "followUp" });
+  };
+
+  // P5.4a — Inline tool registration (BLOCKER-2 fix).
+  // Registered inside the factory for closure access to cache, cwd, saveCache,
+  // and llmBatchesCompleted. Uses real mtime from stat().
+  pi.registerTool({
+    name: "_doc_injector_keywords",
+    label: "Doc Injector Keywords",
+    description:
+      "Save LLM-generated keywords for documentation files. Call this tool with the keywords array after analyzing file snippets.",
+    parameters: Type.Object({
+      keywords: Type.Array(
+        Type.Object({
+          path: Type.String(),
+          keywords: Type.Array(Type.String()),
+        }),
+      ),
+    }),
+    execute: async (_id, params, _signal, _onUpdate, ctx) => {
+      const generated = params.keywords as Array<{ path: string; keywords: string[] }>;
+      const { stat } = await import("node:fs/promises");
+      let saved = 0;
+      for (const item of generated) {
+        const absPath = resolve(ctx.cwd, config.docsPath, item.path);
+        const fileStat = await stat(absPath).catch(() => null);
+        if (!fileStat) {
+          console.warn(`[doc-injector] Skipping keyword save for ${item.path}: file not found`);
+          continue;
+        }
+        cache.files[item.path] = {
+          mtimeMs: fileStat.mtimeMs,
+          keywords: item.keywords.map((k) => k.toLowerCase()).slice(0, 20),
+        };
+        saved++;
+      }
+      await saveCache(ctx.cwd, cache);
+      llmBatchesCompleted++;
+      llmTotalFiles += saved;
+      return {
+        content: [{ type: "text" as const, text: `Keywords saved for ${saved} files.` }],
+        details: undefined as never,
+      };
+    },
+  });
 
   // ---- Event: session_start ----
-  // Pi fires session_start twice on startup (both with reason "startup").
-  // Use a 2-second dedup window to skip the duplicate. Real session changes
-  // (/new, /resume, /fork) happen well outside this window.
-  pi.on("session_start", async (_event, ctx) => {
-    const now = Date.now();
-    if (now - lastInitTime < 100) return;
-    lastInitTime = now;
-    await initRegistry(ctx.cwd);
+  // Pi emits session_start for startup, reload, and real session transitions.
+  // Skip the reload variant because resources_discover will rebuild docs right
+  // after it, and deduplicate any overlapping non-reload inits.
+  pi.on("session_start", async (event, ctx) => {
+    // P5.4d — Safety unbind: clear all LLM keyword gen state on session start
+    keywordGenInFlight = false;
+    llmBatchesCompleted = 0;
+    llmTotalFiles = 0;
+
+    if (event.reason === "reload") return;
+
+    if (initRegistryPromise) {
+      await initRegistryPromise;
+      return;
+    }
+
+    initRegistryPromise = initRegistry(ctx.cwd);
+    try {
+      await initRegistryPromise;
+    } finally {
+      initRegistryPromise = null;
+    }
   });
 
-  const reloadRegistry = async (): Promise<number> => {
+  const reloadRegistry = async (cwd?: string): Promise<number> => {
     if (!registry) throw new Error("No registry loaded");
+    const effectiveCwd = cwd ?? process.cwd();
+
+    // Reload cache from disk to pick up LLM-generated entries
+    const freshCache = await loadCache(effectiveCwd);
+    cache = freshCache;
+    registry.updateCache(cache);
+
     await registry.rebuild();
+
+    const dirty = registry.getDirtyCache();
+    if (Object.keys(dirty).length > 0) {
+      await safeSaveCache(effectiveCwd, dirty);
+    }
+
     const count = registry.getEntries().length;
     console.log(`[doc-injector] Reloaded: ${count} documents`);
     return count;
   };
 
   // ---- Event: resources_discover (reload) ----
-  pi.on("resources_discover", async (_event, _ctx) => {
-    await reloadRegistry();
+  pi.on("resources_discover", async (_event, ctx) => {
+    await reloadRegistry(ctx.cwd);
   });
 
   // ---- Event: input (user message matching) ----
@@ -130,6 +246,17 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   // BEFORE before_agent_start fires, so docs are injected in time for
   // the assistant's immediate response.
   pi.on("input", async (event, _ctx) => {
+    // P5.4d — Safety unbind: if the user is typing interactively, clear all
+    // LLM keyword gen state (they may have aborted the generation).
+    if (event.source === "interactive") {
+      keywordGenInFlight = false;
+      llmBatchesCompleted = 0;
+      llmTotalFiles = 0;
+    }
+
+    // P5.4b — Guard: skip keyword matching during LLM keyword generation
+    if (keywordGenInFlight) return;
+
     if (!enabled || !registry) return;
     if (!event.text) return;
 
@@ -147,6 +274,9 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   // non-injected docs, abort the current generation and restart with the
   // injected context — no waiting for the next turn.
   pi.on("message_update", async (event, ctx) => {
+    // P5.4b — Guard: skip auto-abort logic during LLM keyword generation
+    if (keywordGenInFlight) return;
+
     if (!enabled || !registry) return;
 
     const msg = event.message;
@@ -187,6 +317,9 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
 
   // ---- Event: before_agent_start (inject into system prompt) ----
   pi.on("before_agent_start", async (event, ctx) => {
+    // P5.4b — Guard: skip injection during LLM keyword generation
+    if (keywordGenInFlight) return;
+
     if (!enabled || !registry || pendingMatches.size === 0) return;
 
     const matchedEntries: DocEntry[] = [];
@@ -227,13 +360,27 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     };
   });
 
-  // ---- Event: agent_end (restart after auto-abort) ----
-  pi.on("agent_end", async () => {
+  // ---- Event: agent_end (restart after auto-abort + LLM batch summary) ----
+  pi.on("agent_end", async (event, ctx) => {
+    // P5.4c — Summary notification from agent_end (BLOCKER-3)
+    keywordGenInFlight = false;
+    if (llmBatchesCompleted > 0) {
+      await ctx.ui.notify(
+        `Doc keywords: ${llmTotalFiles} files across ${llmBatchesCompleted} batch(es)`,
+        "info",
+      );
+      llmBatchesCompleted = 0;
+      llmTotalFiles = 0;
+    }
+
     if (abortingForInjection) {
       abortingForInjection = false;
-      // Send a follow-up message to restart the turn.
-      // before_agent_start will inject the matched docs into context.
-      pi.sendUserMessage("continue", { deliverAs: "followUp" });
+      // Defer sendUserMessage to next tick to avoid re-entrancy issues.
+      setTimeout(() => {
+        pi.sendUserMessage("continue");
+      }, 0);
+    } else {
+      console.log('[doc-injector] agent_end: abortingForInjection is false, skipping');
     }
   });
 
@@ -243,5 +390,7 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     getEnabled,
     setEnabled,
     reloadRegistry,
+    getConfig,
+    generateKeywordsLLM,
   });
 }

package/injector.ts

@@ -13,6 +13,21 @@ export interface NotifyCapability {
   notify: (msg: string, type?: "info" | "warning" | "error") => void;
 }
 
+/**
+ * Sanitize keywords for safe injection into the system prompt.
+ *
+ * - Strips \n and \r (replaces with space) to prevent prompt injection
+ * - Caps each keyword at 100 characters
+ * - Enforces a hard limit of 20 keywords
+ */
+function sanitizeKeywords(keywords: string[]): string[] {
+  return keywords
+    .map((k) => k.replace(/[\n\r]/g, " ").trim())
+    .filter((k) => k.length > 0)
+    .map((k) => (k.length > 100 ? k.slice(0, 100) : k))
+    .slice(0, 20);
+}
+
 /**
  * Build a system prompt append string from matched documents.
  */
@@ -29,7 +44,9 @@ export function buildSystemPromptAppend(
   ];
 
   for (const entry of entries) {
-    const keywords = matchedKeywords.get(entry.filePath) ?? [];
+    // Sanitize keywords before display to prevent prompt injection
+    const rawKeywords = matchedKeywords.get(entry.filePath) ?? [];
+    const keywords = sanitizeKeywords(rawKeywords);
     sections.push(`### ${entry.title}`);
     sections.push(`Source: \`${entry.relativePath}\``);
     if (keywords.length > 0) {