pi-context-engine 0.2.0 → 0.3.0

package/README.md

@@ -88,6 +88,19 @@ and the index lifecycle.
 Re-indexing only touches changed files. In the run above, a second pass over an
 unchanged repo took 237 ms.
 
+### Automatic indexing
+
+- **Auto-index at session start**: ~2 s after a session opens, an incremental
+  reindex runs in the background, so the index is warm before the agent needs
+  it. You get a notification only when something changed.
+- **Lazy fallback**: if a search still lands on an empty index, indexing kicks
+  off on the spot and the tool asks you to retry in a few seconds.
+- **One switch for both**: `indexer.auto_index: false` disables both behaviors;
+  build the index manually with `/context-index` instead.
+- **Why it matters**: incremental runs are cheap (a few hundred ms on an
+  unchanged repo), so this removes the classic cold-start failure of searching
+  an empty index, without ever blocking the UI.
+
 ## Requirements
 
 - **Node 22+** on macOS, Linux or Windows.
@@ -153,6 +166,7 @@ indexer:
   context_headers: true  # prepend `// file: ... | unit: ... | imports: ...` to chunk text
   chunking: legacy        # legacy (default) | ast-density — cAST-style budgeted split-merge (experimental)
   chunk_budget_chars: 1200 # non-whitespace budget per ast-density chunk
+  auto_index: true        # incremental background reindex shortly after session start
 memory:
   redact_pii: true
 pricing:

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-context-engine",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "Semantic code context for pi: tree-sitter index + hybrid vector/BM25 retrieval + cross-session memory, all local. Port of code-context-engine (MIT, elara-labs) as a native pi extension.",
   "keywords": [
     "pi-package",

package/src/config.ts

@@ -93,6 +93,8 @@ export interface Config {
 	indexerIgnore: string[];
 	/** Prepend the R1 context header (path/lang/unit/imports) to chunk content. */
 	indexerContextHeaders: boolean;
+	/** Run an incremental reindex in the background shortly after session start. */
+	indexerAutoIndex: boolean;
 	indexerRedactSecrets: boolean;
 	memoryRedactPii: boolean;
 	pricingModel: string;
@@ -125,6 +127,7 @@ export function defaultConfig(): Config {
 		indexerChunkBudgetChars: 1200,
 		indexerIgnore: [...DEFAULT_IGNORE],
 		indexerContextHeaders: true,
+		indexerAutoIndex: true,
 		indexerRedactSecrets: true,
 		memoryRedactPii: true,
 		pricingModel: "opus",
@@ -158,6 +161,7 @@ const MAPPING: Record<string, keyof Config> = {
 	"indexer.chunk_budget_chars": "indexerChunkBudgetChars",
 	"indexer.ignore": "indexerIgnore",
 	"indexer.context_headers": "indexerContextHeaders",
+	"indexer.auto_index": "indexerAutoIndex",
 	"indexer.redact_secrets": "indexerRedactSecrets",
 	"memory.redact_pii": "memoryRedactPii",
 	"storage.path": "storagePath",

package/src/engine.ts

@@ -166,6 +166,7 @@ export class ContextEngine {
 	private searchesSinceDecision = 0;
 	private hasRecordedDecision = false;
 	private vecBackfillStarted = false;
+	private autoIndexDone = false;
 
 	constructor(options: EngineOptions) {
 		this.projectDir = fs.realpathSync(path.resolve(options.projectDir));
@@ -336,6 +337,34 @@ export class ContextEngine {
 		};
 	}
 
+	/**
+	 * One-shot background incremental index, kicked at session start by the
+	 * extension (config `indexer.auto_index`, default true). Returns null when
+	 * disabled or when this engine already ran (or skipped) its auto-index, so
+	 * callers never stack a second background run. A REJECTED run (backend
+	 * init, model download, embedding crash) un-consumes the one-shot so the
+	 * next kick retries — the index is still empty, after all. The pipeline
+	 * lock in runIndexing serializes this with a manual /context-index or a
+	 * lazy search-triggered run, so races are safe — just wasteful, which the
+	 * one-shot guard prevents.
+	 */
+	async autoIndexOnce(): Promise<{
+		indexed: number;
+		chunks: number;
+		deleted: number;
+		skipped: number;
+		errors: string[];
+	} | null> {
+		if (!this.config.indexerAutoIndex || this.autoIndexDone) return null;
+		this.autoIndexDone = true;
+		try {
+			return await this.reindex({});
+		} catch (err) {
+			this.autoIndexDone = false;
+			throw err;
+		}
+	}
+
 	// ── Memory writes ─────────────────────────────────────────────────────
 
 	async recordDecision(
@@ -1252,7 +1281,9 @@ export class ContextEngine {
 		lines.push(
 			chunks > 0
 				? `${chunks.toLocaleString()} chunks indexed and searchable via context_search.`
-				: "Index is empty — run /context-index (or call reindex) to build it.",
+				: this.config.indexerAutoIndex
+					? "Index is empty — a background index starts automatically at session start (indexer.auto_index); retry context_search shortly."
+					: "Index is empty — run /context-index (or call reindex) to build it.",
 		);
 		return lines.join("\n");
 	}

package/src/extension.ts

@@ -16,12 +16,61 @@ import { applyContextSearchGuidance } from "./guidance.js";
 /** Tools whose payloads are worth capturing for session timelines. */
 const CAPTURED_TOOLS = new Set(["bash", "read", "edit", "write", "grep", "find", "ls"]);
 
+/**
+ * Delay before the session-start auto-index kicks in: long enough that pi's
+ * initial render, skill loading and the user's first keystrokes win the event
+ * loop, short enough that the index is warm before the agent needs it.
+ */
+const AUTO_INDEX_DELAY_MS = 2000;
+
 export default function contextEngineExtension(pi: ExtensionAPI) {
 	let engine: ContextEngine | null = null;
 	let currentCwd = process.cwd();
 	let promptNumber = 0;
 	let lazyIndexStarted = false;
 	let bootstrapSent = false;
+	let autoIndexTimer: NodeJS.Timeout | null = null;
+	let indexInFlight = false;
+
+	/**
+	 * Kick the one-shot background index (session-start auto-index and the
+	 * first-search lazy path share this). Never stacks: `lazyIndexStarted`
+	 * guards re-entry, `indexInFlight` lets context_search say "still
+	 * building" instead of re-kicking, and `engine.autoIndexOnce()` is
+	 * itself one-shot per engine instance.
+	 */
+	const kickBackgroundIndex = (eng: ContextEngine, notify?: (msg: string, level: "info" | "warning") => void) => {
+		lazyIndexStarted = true;
+		indexInFlight = true;
+		void eng
+			.autoIndexOnce()
+			.then((result) => {
+				if (!result || !notify || engine !== eng) return; // session moved on
+				if (result.errors.length > 0) {
+					notify(`Context index: background index hit ${result.errors.length} error(s)`, "warning");
+				} else if (result.indexed > 0) {
+					notify(`Context index updated: ${result.indexed} file(s), ${result.chunks} chunk(s)`, "info");
+				}
+			})
+			.catch((err: unknown) => {
+				// Backend init, model download and embedding failures REJECT the
+				// reindex (result.errors only covers per-file errors), so without
+				// this branch a failed startup index would die silently and leave
+				// the lazy path burned on an empty index. Free the lazy path so a
+				// later search retries (autoIndexOnce un-consumes its one-shot),
+				// and surface the failure when we have a notify channel.
+				if (engine === eng) {
+					lazyIndexStarted = false;
+					if (notify) {
+						const detail = (err instanceof Error ? err.message : String(err)).slice(0, 120);
+						notify(`Context index: background indexing failed — run /context-index for details (${detail})`, "warning");
+					}
+				}
+			})
+			.finally(() => {
+				indexInFlight = false;
+			});
+	};
 
 	const getEngine = (cwd?: string): ContextEngine => {
 		const dir = cwd ?? currentCwd;
@@ -68,15 +117,29 @@ export default function contextEngineExtension(pi: ExtensionAPI) {
 			if (!query) return text("Query cannot be empty.");
 			if (query.length > 10_000) return text("Query too long (max 10000 characters).");
 
-			// Lazy index on first search: kick off in the background and tell the
+			// Empty index: a background build may already be running (session-start
+			// auto-index or an earlier search). Never stack a second run; tell the
 			// caller to retry instead of blocking for a multi-minute first index.
-			if (eng.chunkCount() === 0 && !lazyIndexStarted) {
-				lazyIndexStarted = true;
-				void eng.reindex({}).catch(() => undefined);
-				return text(
-					`Index for ${eng.projectName} is empty — indexing started in the background. ` +
-						"Retry this search in a few seconds, or run /context-index for a synchronous index.",
-				);
+			// `indexer.auto_index: false` disables this lazy kick too — one switch
+			// for all automatic indexing.
+			if (eng.chunkCount() === 0) {
+				if (indexInFlight) {
+					return text(
+						`Index for ${eng.projectName} is still building in the background — retry this search in a few seconds.`,
+					);
+				}
+				if (!eng.config.indexerAutoIndex) {
+					return text(
+						`Index for ${eng.projectName} is empty and indexer.auto_index is off — run /context-index (or call reindex) to build it.`,
+					);
+				}
+				if (!lazyIndexStarted) {
+					kickBackgroundIndex(eng);
+					return text(
+						`Index for ${eng.projectName} is empty — indexing started in the background. ` +
+							"Retry this search in a few seconds, or run /context-index for a synchronous index.",
+					);
+				}
 			}
 
 			const result = await eng.search(query, { topK: params.top_k, maxTokens: params.max_tokens });
@@ -309,6 +372,39 @@ export default function contextEngineExtension(pi: ExtensionAPI) {
 		engine.startBackgroundWork();
 		lazyIndexStarted = false;
 		bootstrapSent = false;
+		if (autoIndexTimer) {
+			clearTimeout(autoIndexTimer);
+			autoIndexTimer = null;
+		}
+		// Auto-index (config `indexer.auto_index`): kick an incremental reindex
+		// shortly after the session is ready, so the agent never discovers an
+		// empty or stale index mid-task. Incremental runs are cheap when nothing
+		// changed (walk + hash only); the pipeline lock serializes this with a
+		// manual /context-index or a lazy search-triggered run. Errors surface as
+		// a notification, never as a session failure.
+		if (engine.config.indexerAutoIndex) {
+			const eng = engine;
+			autoIndexTimer = setTimeout(() => {
+				autoIndexTimer = null;
+				if (lazyIndexStarted) return; // a search already kicked a background index
+				if (engine !== eng) {
+					// getEngine switched projects mid-delay (a tool ran with a
+					// different cwd): index the LIVE engine instead of the retired
+					// one, so the new project still gets its warmup instead of
+					// waiting for an unlucky lazy kick.
+					if (engine) {
+						kickBackgroundIndex(engine, (msg, level) => {
+							if (ctx.hasUI) ctx.ui.notify(msg, level);
+						});
+					}
+					return;
+				}
+				kickBackgroundIndex(eng, (msg, level) => {
+					if (ctx.hasUI) ctx.ui.notify(msg, level);
+				});
+			}, AUTO_INDEX_DELAY_MS);
+			autoIndexTimer.unref?.();
+		}
 	});
 
 	pi.on("before_agent_start", async (event, ctx) => {
@@ -355,6 +451,10 @@ export default function contextEngineExtension(pi: ExtensionAPI) {
 	});
 
 	pi.on("session_shutdown", async () => {
+		if (autoIndexTimer) {
+			clearTimeout(autoIndexTimer);
+			autoIndexTimer = null;
+		}
 		await engine?.shutdown().catch(() => undefined);
 		engine = null;
 	});

package/src/indexer/embedder.ts

@@ -7,7 +7,7 @@
 import path from "node:path";
 import fs from "node:fs";
 import os from "node:os";
-import { sha256Hex } from "../utils.js";
+import { sha256Hex, yieldToEventLoop } from "../utils.js";
 
 export const DEFAULT_LOCAL_MODEL = "Xenova/all-MiniLM-L6-v2";
 /** R11: code-trained local embedder, measured in README (not a default — see eval table). */
@@ -38,6 +38,14 @@ export function transformersCacheDir(): string {
 	return resolveModelCacheDir();
 }
 
+/**
+ * Texts per embed() call in LocalBackend. ONNX inference and tokenization run
+ * on the main thread, so this bounds one synchronous block: 64 texts measured
+ * at ~0.5-0.9 s per call (a visible TUI freeze, see
+ * scripts/bench-event-loop.mts); 16 lands around ~100 ms per block.
+ */
+const LOCAL_SUB_BATCH = 16;
+
 interface TransformersModule {
 	pipeline: (
 		task: string,
@@ -135,14 +143,18 @@ export class LocalBackend implements EmbeddingBackend {
 		this.dimension = probe[0]!.length;
 	}
 
-	async embedTexts(texts: string[], _batchSize = 64): Promise<number[][]> {
+	async embedTexts(texts: string[], _batchSize?: number): Promise<number[][]> {
 		const embed = await this.getPipeline();
 		if (texts.length === 0) return [];
 		if (texts.length === 1) return embed(texts);
 		const out: number[][] = [];
-		for (let i = 0; i < texts.length; i += 64) {
-			const batch = await embed(texts.slice(i, i + 64));
+		for (let i = 0; i < texts.length; i += LOCAL_SUB_BATCH) {
+			const batch = await embed(texts.slice(i, i + LOCAL_SUB_BATCH));
 			out.push(...batch);
+			// Cooperative yield: one embed() call blocks the main thread for the
+			// whole sub-batch (tokenization + native ONNX run), so without this
+			// the UI starves for the whole indexing run.
+			if (i + LOCAL_SUB_BATCH < texts.length) await yieldToEventLoop();
 		}
 		return out;
 	}

package/src/indexer/manifest.ts

@@ -17,29 +17,47 @@ interface ManifestData {
 export class Manifest {
 	private entries: Record<string, string> = {};
 	private schemaVersionValue: number = CURRENT_SCHEMA_VERSION;
+	/** Set when an existing manifest was unreadable or not in the recognized
+	 * format — the pipeline wipes the index stores so no stale rows survive. */
+	private storeResetNeeded = false;
 	private lastGitShaValue: string | null = null;
 	private embeddingDimValue: number | null = null;
 
 	constructor(private readonly manifestPath: string) {
-		if (fs.existsSync(manifestPath)) {
-			try {
-				const loaded = JSON.parse(fs.readFileSync(manifestPath, "utf-8")) as ManifestData;
-				if (loaded && typeof loaded === "object") {
-					if ("__schema_version" in loaded) {
-						this.schemaVersionValue = loaded.__schema_version ?? 1;
-						this.entries = loaded.files ?? {};
-						this.lastGitShaValue = loaded.last_git_sha ?? null;
-						this.embeddingDimValue = loaded.embedding_dim ?? null;
-					} else {
-						// Old plain-dict format (pre-v0.2) — treat as version 1.
-						this.schemaVersionValue = 1;
-						this.entries = loaded as unknown as Record<string, string>;
-					}
-				}
-			} catch {
-				this.entries = {};
+		if (!fs.existsSync(manifestPath)) return;
+		// An existing manifest must prove it is in the recognized versioned
+		// format; until it does, it is treated as foreign (version 0) and the
+		// pipeline wipes the stores before rebuilding. No compat parsing of
+		// pre-versioning shapes.
+		let recognized = false;
+		try {
+			const loaded = JSON.parse(fs.readFileSync(manifestPath, "utf-8")) as ManifestData;
+			const version = loaded && typeof loaded === "object" ? loaded.__schema_version : undefined;
+			const files = loaded && typeof loaded === "object" ? loaded.files : undefined;
+			if (
+				typeof version === "number" &&
+				version > 0 &&
+				files !== null &&
+				typeof files === "object" &&
+				!Array.isArray(files)
+			) {
+				recognized = true;
+				this.schemaVersionValue = version;
+				this.entries = files ?? {};
+				this.lastGitShaValue = loaded.last_git_sha ?? null;
+				this.embeddingDimValue = loaded.embedding_dim ?? null;
 			}
+		} catch {
+			// corrupt JSON — falls through to the reset below
 		}
+		if (!recognized) this.schemaVersionValue = 0;
+		this.storeResetNeeded = !recognized || this.schemaVersionValue !== CURRENT_SCHEMA_VERSION;
+	}
+
+	/** True when the on-disk manifest was unusable/foreign and the index stores
+	 * must be wiped before this run repopulates them. */
+	get needsStoreReset(): boolean {
+		return this.storeResetNeeded;
 	}
 
 	get lastGitSha(): string | null {