pi-doc-injector 0.1.4 → 0.2.1

package/README.md

@@ -54,7 +54,10 @@ keywords:
 ```
 
 3. Start Pi. The extension scans `docs/` on session start.
-4. When the LLM mentions keywords from your docs, the relevant document is injected into the next turn's system prompt.
+4. When the user mentions a keyword, the matching doc is injected into the
+   system prompt **before the assistant responds** — no one-turn delay.
+5. If the assistant mentions a NEW keyword mid-response, generation is
+   automatically aborted and restarted with the doc injected immediately.
 
 ## Configuration
 
@@ -98,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/index.ts

@@ -66,9 +66,11 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   // ---- State ----
   let config = 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
 
   // ---- Helpers ----
   const getRegistry = () => registry;
@@ -98,8 +100,23 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
   };
 
   // ---- Event: session_start ----
-  pi.on("session_start", async (_event, ctx) => {
-    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) => {
+    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> => {
@@ -115,54 +132,65 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     await reloadRegistry();
   });
 
-  // ---- Event: message_update (streaming detection) ----
-  // NOTE: Pi's message_update event sends the full accumulated content of the
-  // assistant message on each update, not just the delta. We therefore REPLACE
-  // (not append to) the text buffer on each event, ensuring we always match
-  // against the complete message text.
-  pi.on("message_update", async (event, _ctx) => {
+  // ---- Event: input (user message matching) ----
+  // message_update only fires for assistant streaming messages, not user
+  // messages. We use the input event instead to populate pendingMatches
+  // BEFORE before_agent_start fires, so docs are injected in time for
+  // the assistant's immediate response.
+  pi.on("input", async (event, _ctx) => {
+    if (!enabled || !registry) return;
+    if (!event.text) return;
+
+    const matcher = buildMatcher();
+    if (!matcher) return;
+
+    const results = matcher.match(event.text);
+    for (const result of results) {
+      pendingMatches.set(result.entry.filePath, result.matchedKeywords);
+    }
+  });
+
+  // ---- Event: message_update (assistant streaming) ----
+  // For assistant streaming messages: if we detect NEW keyword matches for
+  // 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) => {
     if (!enabled || !registry) return;
 
-    // Only process assistant messages
     const msg = event.message;
     if (msg.role !== "assistant") return;
 
-    // Replace buffer with full message text (message_update contains full content)
     const content = (msg as unknown as { content: unknown }).content;
     textBuffer = extractText(content);
     if (!textBuffer) return;
 
-    // Run matcher
     const matcher = buildMatcher();
     if (!matcher) return;
 
     const results = matcher.match(textBuffer);
 
-    // Store matches (dedup by filePath)
+    let hasNew = false;
     for (const result of results) {
+      if (!pendingMatches.has(result.entry.filePath)) {
+        hasNew = true;
+      }
       pendingMatches.set(result.entry.filePath, result.matchedKeywords);
     }
+
+    if (hasNew && !ctx.isIdle() && !abortingForInjection) {
+      abortingForInjection = true;
+      ctx.abort();
+    }
   });
 
   // ---- Event: message_end (finalize matches) ----
-  pi.on("message_end", async (event, ctx) => {
+  // Notification moved to before_agent_start so it fires for both user-triggered
+  // and auto-abort-triggered injections. message_end now just resets state.
+  pi.on("message_end", async (event, _ctx) => {
     if (!enabled || !registry) return;
-
     const msg = event.message;
     if (msg.role !== "assistant") return;
-
-    // Clear buffer
     textBuffer = "";
-
-    // Notify user about pending injections
-    if (pendingMatches.size > 0) {
-      const matchedEntries: DocEntry[] = [];
-      for (const [filePath] of pendingMatches) {
-        const entry = registry.getEntries().find((e) => e.filePath === filePath);
-        if (entry) matchedEntries.push(entry);
-      }
-      notifyInjection(ctx.ui, matchedEntries, pendingMatches);
-    }
   });
 
   // ---- Event: before_agent_start (inject into system prompt) ----
@@ -194,6 +222,12 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
 
     // Mark as injected only after confirming injection will happen
     registry.markInjected(matchedEntries.map((e) => e.filePath));
+
+    // Notify user about injection (moved here from message_end so it fires
+    // even when matches come from user messages, which get cleared before
+    // the assistant's message_end)
+    notifyInjection(ctx.ui, matchedEntries, pendingMatches);
+
     pendingMatches.clear();
 
     return {
@@ -201,6 +235,16 @@ export default async function docInjectorExtension(pi: ExtensionAPI) {
     };
   });
 
+  // ---- Event: agent_end (restart after auto-abort) ----
+  pi.on("agent_end", async () => {
+    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" });
+    }
+  });
+
   // ---- Commands ----
   registerCommands(pi, {
     getRegistry,

package/package.json

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