pi-doc-injector 0.2.1 → 0.3.1

package/index.ts

@@ -53,18 +53,21 @@
  * 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;
@@ -72,23 +75,51 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   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);
-    const count = registry.getEntries().length;
-    if (count > 0) {
-      console.log(`[doc-injector] Loaded ${count} documents from ${docsPath}`);
-    } else {
-      console.warn(`[doc-injector] No documents found at ${docsPath}`);
+    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;
   };
 
   const buildMatcher = (): KeywordMatcher | null => {
@@ -99,11 +130,69 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     );
   };
 
+  // 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) {
+          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 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) {
@@ -119,17 +208,29 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     }
   });
 
-  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) ----
@@ -138,6 +239,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;
 
@@ -155,6 +267,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;
@@ -195,6 +310,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[] = [];
@@ -213,7 +331,6 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     // past the model's limit.
     const usage = ctx.getContextUsage();
     if (usage && usage.tokens && usage.tokens > 0 && usage.percent && usage.percent > config.contextThreshold) {
-      console.warn(`[doc-injector] Skipping injection: context usage > ${config.contextThreshold}%`);
       pendingMatches.clear();
       return;
     }
@@ -235,13 +352,25 @@ 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);
     }
   });
 
@@ -251,5 +380,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) {

package/keyword-gen.ts

@@ -0,0 +1,142 @@
+/**
+ * Local keyword generation — extracts keywords from filenames and content
+ * when no frontmatter is available.
+ *
+ * Extraction sources:
+ *   1. Filename parts (split on -, _, .)
+ *   2. Markdown headings (# Title, ## Title, etc.)
+ *   3. Code symbols (function, class, const, interface, type, enum)
+ *
+ * All keywords are lowercased, deduplicated, and filtered through a stop-word list.
+ * Output is capped at 20 keywords.
+ */
+
+const STOP_WORDS = new Set<string>([
+  // Articles
+  "a", "an", "the",
+  // Pronouns
+  "i", "you", "he", "she", "it", "we", "they",
+  "me", "him", "her", "us", "them",
+  "my", "your", "his", "its", "our", "their",
+  "this", "that", "these", "those",
+  "who", "whom", "whose", "which", "what",
+  // Prepositions
+  "in", "on", "at", "by", "for", "with", "about",
+  "to", "from", "of", "into", "onto", "upon",
+  "over", "under", "between", "among", "through",
+  "during", "before", "after", "above", "below",
+  "up", "down", "out", "off",
+  // Conjunctions
+  "and", "but", "or", "nor", "so", "yet", "for",
+  "if", "then", "than", "as", "when", "while",
+  "because", "since", "although", "though",
+  // Auxiliary/modal verbs
+  "is", "are", "was", "were", "be", "been", "being",
+  "have", "has", "had", "having",
+  "do", "does", "did", "doing",
+  "will", "would", "shall", "should", "can", "could",
+  "may", "might", "must",
+  // Common adverbs
+  "not", "no", "yes",
+  "just", "only", "also", "too", "very", "now", "then",
+  "here", "there", "where", "how", "why",
+  "all", "each", "every", "both", "few", "more", "most",
+  "some", "any", "other", "another", "such",
+  "much", "many", "little", "less",
+  // Common content-less words
+  "get", "set", "put", "use", "make", "see", "need",
+  "one", "two", "three", "first", "second", "third",
+  "using", "used", "into", "onto", "new",
+  "note", "notes", "example", "examples", "todo",
+]);
+
+/**
+ * Generate up to 20 keywords from a file's name and content.
+ *
+ * Sources (in order, each adds keywords until cap is reached):
+ *   1. Filename parts — split on `-`, `_`, and `.`, keep segments ≥ 3 chars
+ *   2. Markdown headings — text after `#` markers
+ *   3. Code symbols — function/class/const/interface/type/enum declarations
+ *
+ * Each candidate is lowercased, filtered through a stop-word list, deduplicated,
+ * and limited to words with ≥ 3 characters.
+ *
+ * @param fileName - The basename of the file (e.g. "api-authentication.md")
+ * @param content  - The full file content
+ * @returns Up to 20 deduplicated keyword strings
+ */
+export function generateKeywords(
+  fileName: string,
+  content: string,
+): string[] {
+  const keywords: string[] = [];
+
+  // Source 1: Filename parts
+  addFromFilename(fileName, keywords);
+
+  // Source 2: Markdown headings
+  addFromHeadings(content, keywords);
+
+  // Source 3: Code symbols
+  addFromCodeSymbols(content, keywords);
+
+  // Deduplicate while preserving order
+  const seen = new Set<string>();
+  const result: string[] = [];
+  for (const kw of keywords) {
+    const lower = kw.toLowerCase();
+    if (seen.has(lower)) continue;
+    seen.add(lower);
+    result.push(kw);
+  }
+
+  return result.slice(0, 20);
+}
+
+/** Extract keyword candidates from filename parts. */
+function addFromFilename(fileName: string, out: string[]): void {
+  // Strip extension(s)
+  const nameWithoutExt = fileName.replace(/\.[^.]+$/, "");
+
+  // Split on common delimiters
+  const parts = nameWithoutExt.split(/[-_.\s]+/);
+
+  for (const part of parts) {
+    const cleaned = part.replace(/[^a-zA-Z0-9]/g, "").toLowerCase();
+    if (cleaned.length >= 3 && !STOP_WORDS.has(cleaned)) {
+      out.push(cleaned);
+    }
+  }
+}
+
+/** Extract keyword candidates from markdown headings (#, ##, ###, etc.). */
+function addFromHeadings(content: string, out: string[]): void {
+  const headingRegex = /^#{1,6}\s+(.+)$/gm;
+  let match: RegExpExecArray | null;
+  while ((match = headingRegex.exec(content)) !== null) {
+    const headingText = match[1].trim();
+    // Split heading into words
+    const words = headingText.split(/\s+/);
+    for (const word of words) {
+      const cleaned = word.replace(/[^a-zA-Z0-9]/g, "").toLowerCase();
+      if (cleaned.length >= 3 && !STOP_WORDS.has(cleaned)) {
+        out.push(cleaned);
+      }
+    }
+  }
+}
+
+/** Extract keyword candidates from code symbol declarations. */
+function addFromCodeSymbols(content: string, out: string[]): void {
+  // Match: function name, class name, const name, interface name, type name, enum name
+  // Also: export function, export class, export const, etc.
+  const symbolRegex = /(?:export\s+)?(?:async\s+)?(?:function|class|const|interface|type|enum)\s+(\w+)/gm;
+  let match: RegExpExecArray | null;
+  while ((match = symbolRegex.exec(content)) !== null) {
+    const name = match[1];
+    const cleaned = name.toLowerCase();
+    if (cleaned.length >= 3 && !STOP_WORDS.has(cleaned)) {
+      out.push(cleaned);
+    }
+  }
+}

package/keyword-llm.ts

@@ -0,0 +1,57 @@
+/**
+ * LLM Keyword Generation — builds prompts for the LLM to generate keywords
+ * for documentation files via the _doc_injector_keywords tool.
+ */
+
+/** Input for a single file in a keyword generation batch. */
+export interface FileInput {
+  /** Path relative to cwd (e.g. "docs/api.md") */
+  path: string;
+  /** First ~500 chars of the file content as context */
+  snippet: string;
+  /** Existing keywords (from frontmatter/heuristic), so LLM augments not replaces */
+  existingKeywords: string[];
+}
+
+/**
+ * Build a user message prompt instructing the LLM to generate keywords
+ * for a batch of documentation files by calling the _doc_injector_keywords tool.
+ *
+ * The prompt asks the LLM to read each file's snippet and produce 3-10 concise,
+ * searchable keywords per file, incorporating any existing keywords.
+ */
+export function buildKeywordGenPrompt(files: FileInput[]): string {
+  if (files.length === 0) return "";
+
+  const fileDescriptions = files.map((f, i) => {
+    const existing = f.existingKeywords.length > 0
+      ? `  Existing keywords: ${f.existingKeywords.join(", ")}`
+      : "";
+    // Escape markdown special chars in path to prevent prompt injection
+    const safePath = f.path.replace(/[*`\[\]]/g, "\\$&");
+    // Escape backticks in snippet to prevent breaking code fences
+    const safeSnippet = f.snippet.replace(/```/g, "'''");
+    return `${i + 1}. **${safePath}**\n${existing}\n   Snippet:\n\`\`\`\n${safeSnippet}\n\`\`\``;
+  }).join("\n\n");
+
+  const expectedOutput = files.map((f) => {
+    const safePath = f.path.replace(/[*`\[\]]/g, "\\$&");
+    return `  - "${safePath}": keywords array incorporating relevant existing keywords [${f.existingKeywords.slice(0, 5).map(k => `"${k}"`).join(", ")}${f.existingKeywords.length > 5 ? ", ..." : ""}]`;
+  }).join("\n");
+
+  return `Generate documentation keywords for the following ${files.length} file(s). For each file, read the snippet and produce 3-10 concise, searchable keywords that someone might type when looking for this documentation.
+
+Rules:
+- Keywords should be lowercase, 3+ characters, no stop-words
+- Incorporate any existing keywords that are still relevant
+- Focus on the document's core topic, not generic terms
+- Prefer specific technical terms over vague ones
+
+Files:
+${fileDescriptions}
+
+After analysis, call the \`_doc_injector_keywords\` tool with a \`keywords\` array like:
+${expectedOutput}
+
+Do not output any other text — just call the tool with the keywords.`;
+}

package/matcher.ts

@@ -31,6 +31,10 @@ export function extractText(content: unknown): string {
 export class KeywordMatcher {
   private options: MatcherOptions;
 
+  /**
+   * @param entries - The document entries to match against
+   * @param options - Optional matcher settings (merged with defaults)
+   */
   constructor(private entries: DocEntry[], options?: Partial<MatcherOptions>) {
     this.options = { ...DEFAULT_MATCHER_OPTIONS, ...options };
   }
@@ -44,8 +48,13 @@ export class KeywordMatcher {
     for (const entry of this.entries) {
       if (entry.injected) continue;
 
+      // Skip entries with no keywords (empty array or falsy)
+      if (!entry.keywords || entry.keywords.length === 0) continue;
+
       const matchedKeywords: string[] = [];
       for (const keyword of entry.keywords) {
+        // Skip empty keywords — they'd match everything with word boundaries
+        if (!keyword || keyword.trim().length === 0) continue;
         if (this.keywordMatches(text, keyword)) {
           matchedKeywords.push(keyword);
         }
@@ -63,18 +72,13 @@ export class KeywordMatcher {
     return results;
   }
 
+  /**
+   * Check if a single keyword matches the given text.
+   * Uses simple substring inclusion (case-insensitive by default).
+   */
   private keywordMatches(text: string, keyword: string): boolean {
     const search = this.options.caseSensitive ? text : text.toLowerCase();
     const kw = this.options.caseSensitive ? keyword : keyword.toLowerCase();
-
-    if (this.options.wordBoundary) {
-      // Escape special regex chars in keyword, then apply word boundary
-      const escaped = kw.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
-      const flags = this.options.caseSensitive ? "" : "i";
-      const regex = new RegExp(`\\b${escaped}\\b`, flags);
-      return regex.test(search);
-    }
-
     return search.includes(kw);
   }
-}
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-doc-injector",
-  "version": "0.2.1",
+  "version": "0.3.1",
   "description": "Auto-inject relevant project documentation into Pi's LLM context based on keyword matching",
   "type": "module",
   "main": "./index.ts",
@@ -9,6 +9,7 @@
   },
   "files": [
     "*.ts",
+    "*.d.ts",
     "docs/**/*.md",
     "README.md"
   ],
@@ -33,6 +34,9 @@
       "./index.ts"
     ]
   },
+  "dependencies": {
+    "picomatch": "^4.0.2"
+  },
   "peerDependencies": {
     "@mariozechner/pi-coding-agent": "*"
   },

package/picomatch.d.ts

@@ -0,0 +1,11 @@
+declare module "picomatch" {
+  interface PicomatchOptions {
+    dot?: boolean;
+  }
+
+  type Matcher = (input: string) => boolean;
+
+  function picomatch(patterns: string | string[], options?: PicomatchOptions): Matcher;
+
+  export default picomatch;
+}