@krosskinetic/pi-zg 0.1.0 → 0.1.1

package/README.md

@@ -1,22 +1,24 @@
 **WORK IN PROGRESS / ALPHA**
 
-# pi-zg
+# @krosskinetic/pi-zg
 
 A [Pi](https://github.com/earendil-works/pi-mono) package that natively integrates
 [zvec-grep](https://github.com/) (`zg`) semantic code search into Pi: it manages
 `zg`'s shared server, offers to build a missing index interactively, and gives
-the agent three tools (`zg_search`, `zg_rg`, `zg_index`) plus three commands
-(`/zg-status`, `/zg-index`, `/zg-server`). It calls an existing local `zg` CLI;
+the agent three tools (`zg_search`, `zg_rg`, `zg_index`) plus four commands
+(`/zg-settings`, `/zg-status`, `/zg-index`, `/zg-server`). It calls an existing local `zg` CLI;
 it does not install or bundle zvec-grep itself.
 
 ## Prerequisites
 
 Install `@zvec/zvec-grep` separately and ensure `zg` is on `PATH`.
+This extension invokes that local CLI; it does not install, bundle, or
+configure zvec-grep for you.
 
 ## Install into Pi
 
 ```bash
-pi install npm:pi-zg
+pi install npm:@krosskinetic/pi-zg
 ```
 
 For local development, from a checkout of this repo:
@@ -25,6 +27,21 @@ For local development, from a checkout of this repo:
 pi -e .
 ```
 
+## Quick start
+
+1. Install and configure `zg`, then verify that `zg version` works from your
+   project directory.
+2. Install this package and start Pi in the project you want to search.
+3. At session start, pi-zg checks whether `zg` is available, whether its
+   shared server is ready, and whether the project has an index. Unless
+   disabled, it starts a stopped shared server.
+4. Ask Pi a natural-language code-search question. Its `zg_search` tool
+   searches the index. If the project has not been indexed, pi-zg offers to
+   build one interactively.
+
+For example, ask: “Where is authentication token refresh implemented?” For an
+exact identifier, literal, or regex search, Pi can use `zg_rg` instead.
+
 ## What it does
 
 ### Shared server (`zg server`)
@@ -35,8 +52,10 @@ configured via `zg install`) may also use. This extension:
 - **Auto-starts** it at session start if it isn't already running
   (`--no-zg-autostart` disables this). `zg server on` is idempotent, so this
   is safe to run every session.
-- **Never auto-stops it** — since it's shared, stopping it could break other
-  tools relying on it. Use `/zg-server off` to stop it explicitly.
+- **Stops it at Pi session shutdown only when Pi started it.** Pi records the
+  daemon PID and rechecks it before stopping the server, so a server that was
+  already running—or has since been replaced—is left alone. Use
+  `/zg-server off` to stop a shared server explicitly.
 
 With the server running, `zg query` refreshes the index in the background
 automatically after file changes, so `zg_index` is rarely needed once a
@@ -55,14 +74,45 @@ project has an initial index.
 - **`zg_index`** — build, rebuild, or drop the persistent index. Gated by
   prompt guidelines so the agent only uses it when the user explicitly asks.
 
+`zg_search` accepts a natural-language `query` and an optional `limit` of
+1–100 results (zg defaults to 7). `zg_rg` accepts a regex `pattern`, optional
+paths, `fixedString` for literal matching, and one `glob` filter. Its output
+and the output of the other tools are capped at 2,000 lines or 50 KB.
+
 ### Commands (human-invoked)
 
+- **`/zg-settings`** — interactive configuration for Pi's search and refresh
+  mode, default embedding model, embedding device, and provider API key.
 - **`/zg-status`** — zg version, server state, and index status/coverage.
 - **`/zg-index [--rebuild|--drop]`** — build, rebuild, or drop the index
   directly, without going through the LLM. Confirms before `--drop`.
 - **`/zg-server <on|off|status>`** — explicit manual control of the shared
   daemon. Confirms before `off`, since other tools may depend on it.
 
+### Settings (`/zg-settings`)
+
+`/zg-settings` exposes the configuration that most directly affects Pi search:
+
+- **Search and refresh mode** — choose between a shared server with background
+  refresh, direct queries that refresh before every semantic search, or direct
+  queries that use the current index and can be stale. The choice applies to
+  the current Pi session. Switching to either direct mode stops a running
+  shared server only after confirmation, because the daemon owns index writes
+  and may be used by other tools.
+- **Default embedding model** — sets zg's persistent default for newly built
+  indexes, for example `local/potion-code-16m-v2` or
+  `qwen/text-embedding-v4`. Existing indexes keep their recorded embedding
+  schema.
+- **Embedding device** — configures `auto`, `cpu`, `metal`, `vulkan`, or
+  `cuda` for a specified local model. The command can restart a running shared
+  server after confirmation so the changed runtime takes effect.
+- **Provider API key** — saves credentials for a named embedding provider via
+  `zg config provider set`. The value is passed directly to zg and is not
+  shown in Pi notifications.
+
+For advanced index-selection rules, remote endpoints, authentication, and
+other less-common settings, use the underlying `zg` CLI directly.
+
 ### Status
 
 The footer shows `server ●/○` and `index ✓/✗` (colored via the active
@@ -70,11 +120,19 @@ theme), refreshed at session start and at the start of every turn.
 
 ### Flags
 
-- `--no-zg-autostart` — disable automatically starting the shared server.
+- `--no-zg-autostart` — disable automatically starting the shared server at
+  Pi session start.
 - `--no-zg-onboard` — disable the interactive "build an index?" offer;
   `zg_search` fails with a manual-fix message instead (useful for
   non-interactive/scripted `pi -p` runs).
 
+Pass these when launching Pi, for example:
+
+```bash
+pi --no-zg-autostart
+pi --no-zg-onboard
+```
+
 ## Non-goals
 
 - Does not register zg's MCP server as an actual MCP tool source inside Pi
@@ -83,4 +141,4 @@ theme), refreshed at session start and at the start of every turn.
   via subprocess spawns before every call.
 - Does not override Pi's built-in `grep` tool. `zg_search`/`zg_rg` are
   purely additive.
-- Does not auto-stop the shared `zg` server.
+- Does not stop a shared `zg` server that Pi did not start.

package/extensions/pi-zg.ts

@@ -10,8 +10,8 @@ import { Type } from "typebox";
  *
  * Beyond wrapping `zg query`/`zg index`, this extension:
  * - Auto-starts zg's shared server at session start so queries get
- *   background index auto-refresh for free (never auto-stops it, since
- *   it's a daemon other agents/tools may also depend on).
+ *   background index auto-refresh for free, then shuts it down at Pi session
+ *   end only when Pi owns that exact daemon PID.
  * - Caches zg's availability/server/index state per session (refreshed at
  *   session_start and each turn_start) instead of re-deriving it with
  *   extra `zg` subprocess spawns before every tool call.
@@ -25,6 +25,8 @@ import { Type } from "typebox";
 
 const STATUS_KEY = "pi-zg";
 const DEFAULT_LOCAL_MODEL = "local/potion-code-16m-v2";
+type ZgQueryMode = "server" | "direct";
+type ZgRefreshPolicy = "background" | "wait" | "off";
 
 interface ZgState {
 	/** Whether refreshZgState has run at least once this session. */
@@ -34,10 +36,18 @@ interface ZgState {
 	version?: string;
 	/** Whether the shared zg server daemon is up and ready. */
 	serverRunning: boolean;
+	/** PID reported by the ready shared server, when available. */
+	serverPid?: string;
+	/** PID of the server Pi started and therefore owns for this session. */
+	ownedServerPid?: string;
 	/** Whether the current project has a ready index. */
 	indexed: boolean;
 	/** Whether the user already declined the "build an index?" offer this session. */
 	declinedIndexOffer: boolean;
+	/** Pi's active indexed-search transport for this session. */
+	queryMode: ZgQueryMode;
+	/** Pi's active freshness policy for this session. */
+	refreshPolicy: ZgRefreshPolicy;
 }
 
 function createZgState(): ZgState {
@@ -47,6 +57,8 @@ function createZgState(): ZgState {
 		serverRunning: false,
 		indexed: false,
 		declinedIndexOffer: false,
+		queryMode: "server",
+		refreshPolicy: "background",
 	};
 }
 
@@ -70,6 +82,15 @@ function renderStatus(ctx: ExtensionContext, state: ZgState) {
 	ctx.ui.setStatus(STATUS_KEY, `${server} ${index}`);
 }
 
+function parseServerPid(output: string): string | undefined {
+	return output.match(/^PID:\s*(\d+)$/m)?.[1];
+}
+
+function markServerOwnedByPi(state: ZgState) {
+	// Do not risk stopping a daemon if zg did not report a stable PID for it.
+	if (state.serverRunning && state.serverPid) state.ownedServerPid = state.serverPid;
+}
+
 /** Refresh cached zg availability/server/index state and update the footer. */
 async function refreshZgState(
 	pi: ExtensionAPI,
@@ -84,6 +105,7 @@ async function refreshZgState(
 
 	if (!state.available) {
 		state.serverRunning = false;
+		state.serverPid = undefined;
 		state.indexed = false;
 		renderStatus(ctx, state);
 		return state;
@@ -94,6 +116,7 @@ async function refreshZgState(
 		execZg(pi, ["status", "--check-ready"], ctx, { signal, timeout: 5_000 }),
 	]);
 	state.serverRunning = server.code === 0;
+	state.serverPid = state.serverRunning ? parseServerPid(server.stdout) : undefined;
 	state.indexed = status.code === 0;
 
 	renderStatus(ctx, state);
@@ -114,6 +137,180 @@ function parseHitCount(output: string): number | undefined {
 	return match ? Number(match[1]) : undefined;
 }
 
+function searchModeLabel(state: ZgState): string {
+	if (state.queryMode === "server") return "server / background refresh";
+	return state.refreshPolicy === "wait" ? "direct / refresh before search" : "direct / current index";
+}
+
+async function ensureServerRunning(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<boolean> {
+	if (state.serverRunning) return true;
+	ctx.ui.setWorkingMessage("Starting zg server...");
+	const result = await execZg(pi, ["server", "on"], ctx, { timeout: 20_000 });
+	ctx.ui.setWorkingMessage();
+	await refreshZgState(pi, ctx, state);
+	if (result.code === 0) markServerOwnedByPi(state);
+	if (result.code !== 0 || !state.serverRunning) {
+		ctx.ui.notify("zg server could not be started. Choose a direct search mode in /zg-settings or inspect /zg-server status.", "error");
+		return false;
+	}
+	return true;
+}
+
+/**
+ * A daemon has an exclusive lease on index writes, so direct refresh must stop
+ * it first. Ask because that daemon can be shared with other agent clients.
+ */
+async function stopServerForDirectMode(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<boolean> {
+	if (!state.serverRunning) return true;
+	const confirmed = await ctx.ui.confirm(
+		"Switch to direct searches?",
+		"Direct refresh writes to the index itself and requires stopping the shared zg server. This may affect other tools using it.",
+	);
+	if (!confirmed) return false;
+
+	ctx.ui.setWorkingMessage("Stopping zg server...");
+	const result = await execZg(pi, ["server", "off"], ctx, { timeout: 20_000 });
+	ctx.ui.setWorkingMessage();
+	await refreshZgState(pi, ctx, state);
+	if (result.code !== 0 || state.serverRunning) {
+		ctx.ui.notify("zg server could not be stopped; direct refresh was not enabled.", "error");
+		return false;
+	}
+	state.ownedServerPid = undefined;
+	return true;
+}
+
+async function offerServerRestart(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<void> {
+	if (!state.serverRunning) return;
+	const restart = await ctx.ui.confirm(
+		"Restart zg server?",
+		"Restarting applies the changed embedding runtime to the shared server. Other tools using it will briefly lose access.",
+	);
+	if (!restart) return;
+	ctx.ui.setWorkingMessage("Restarting zg server...");
+	const stopped = await execZg(pi, ["server", "off"], ctx, { timeout: 20_000 });
+	if (stopped.code === 0) state.ownedServerPid = undefined;
+	const started = stopped.code === 0 ? await execZg(pi, ["server", "on"], ctx, { timeout: 20_000 }) : stopped;
+	ctx.ui.setWorkingMessage();
+	await refreshZgState(pi, ctx, state);
+	if (started.code === 0) markServerOwnedByPi(state);
+	ctx.ui.notify(
+		started.code === 0 && state.serverRunning ? "zg server restarted." : "zg server restart failed; inspect /zg-server status.",
+		started.code === 0 && state.serverRunning ? "info" : "error",
+	);
+}
+
+async function configureSearchMode(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<void> {
+	const choice = await ctx.ui.select("Search and refresh mode", [
+		"Server — background refresh (fast repeated searches)",
+		"Direct — refresh before each search (no daemon)",
+		"Direct — use current index (fast, may be stale)",
+		"Cancel",
+	]);
+	if (!choice || choice === "Cancel") return;
+
+	if (choice.startsWith("Server")) {
+		state.queryMode = "server";
+		state.refreshPolicy = "background";
+		if (await ensureServerRunning(pi, ctx, state)) ctx.ui.notify("zg searches use the shared server with background refresh.", "info");
+		return;
+	}
+
+	if (!(await stopServerForDirectMode(pi, ctx, state))) return;
+	state.queryMode = "direct";
+	state.refreshPolicy = choice.startsWith("Direct — refresh") ? "wait" : "off";
+	ctx.ui.notify(
+		state.refreshPolicy === "wait"
+			? "zg searches refresh the index before each semantic query."
+			: "zg searches use the current index and may be stale after edits.",
+		"info",
+	);
+}
+
+async function configureDefaultModel(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<void> {
+	const model = await ctx.ui.input("Default embedding model", "e.g. local/potion-code-16m-v2 or qwen/text-embedding-v4");
+	if (!model?.trim()) return;
+	const result = await execZg(pi, ["config", "model", "set", model.trim(), "--default"], ctx);
+	if (result.code !== 0) {
+		ctx.ui.notify("Could not set the embedding model. Inspect zg configuration in a terminal for details.", "error");
+		return;
+	}
+	ctx.ui.notify(`Default embedding model set to ${model.trim()}. Existing indexes keep their current schema.`, "info");
+	await offerServerRestart(pi, ctx, state);
+}
+
+async function configureEmbeddingDevice(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<void> {
+	const model = await ctx.ui.input("Model to configure", "e.g. local/potion-code-16m-v2");
+	if (!model?.trim()) return;
+	const device = await ctx.ui.select("Embedding device", ["auto", "cpu", "metal", "vulkan", "cuda", "Cancel"]);
+	if (!device || device === "Cancel") return;
+	const result = await execZg(pi, ["config", "model", "set", model.trim(), "--device", device], ctx);
+	if (result.code !== 0) {
+		ctx.ui.notify("Could not set the embedding device. Inspect zg configuration in a terminal for details.", "error");
+		return;
+	}
+	ctx.ui.notify(`Embedding device for ${model.trim()} set to ${device}.`, "info");
+	await offerServerRestart(pi, ctx, state);
+}
+
+async function configureProviderKey(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
+	const provider = await ctx.ui.input("Embedding provider", "e.g. qwen");
+	if (!provider?.trim()) return;
+	const apiKey = await ctx.ui.input("API key", "Paste the provider API key");
+	if (!apiKey?.trim()) return;
+	const result = await execZg(pi, ["config", "provider", "set", provider.trim(), "--api-key", apiKey.trim()], ctx);
+	if (result.code !== 0) {
+		ctx.ui.notify("Could not save the provider API key. Inspect zg configuration in a terminal for details.", "error");
+		return;
+	}
+	ctx.ui.notify(`Credentials saved for ${provider.trim()}.`, "info");
+}
+
+async function openZgSettings(pi: ExtensionAPI, ctx: ExtensionContext, state: ZgState): Promise<void> {
+	if (!state.checked) await refreshZgState(pi, ctx, state);
+
+	while (true) {
+		const choice = await ctx.ui.select(`zg settings — ${searchModeLabel(state)}`, [
+			"Search and refresh mode",
+			"Default embedding model",
+			"Embedding device",
+			"Provider API key",
+			"View active settings",
+			"Done",
+		]);
+		if (!choice || choice === "Done") return;
+
+		switch (choice) {
+			case "Search and refresh mode":
+				await configureSearchMode(pi, ctx, state);
+				break;
+			case "Default embedding model":
+				await configureDefaultModel(pi, ctx, state);
+				break;
+			case "Embedding device":
+				await configureEmbeddingDevice(pi, ctx, state);
+				break;
+			case "Provider API key":
+				await configureProviderKey(pi, ctx);
+				break;
+			case "View active settings": {
+				await refreshZgState(pi, ctx, state);
+				const detail = await execZg(pi, ["status"], ctx);
+				ctx.ui.notify(
+					[
+						`Pi search mode: ${searchModeLabel(state)} (session setting)`,
+						`zg server: ${state.serverRunning ? "running" : "stopped"}`,
+						"",
+						(detail.stdout || detail.stderr).trim(),
+					].join("\n"),
+					state.indexed ? "info" : "warning",
+				);
+				break;
+			}
+		}
+	}
+}
+
 /**
  * Lazily offer to build a missing index right where the LLM discovered it's
  * missing, instead of dead-ending with a "go run this command yourself" error.
@@ -140,7 +337,7 @@ async function offerToBuildIndex(
 	}
 
 	ctx.ui.setWorkingMessage("Building zg index...");
-	let result = await execZg(pi, ["index"], ctx, { signal, timeout: 300_000 });
+	let result = await execZg(pi, ["index", "--mode", state.queryMode], ctx, { signal, timeout: 300_000 });
 
 	if (result.code !== 0 && /embedding/i.test(result.stderr)) {
 		// No default embedding model configured yet -- offer to set one and retry.
@@ -157,7 +354,10 @@ async function offerToBuildIndex(
 		}
 		if (model) {
 			await execZg(pi, ["config", "model", "set", model, "--default"], ctx, { signal });
-			result = await execZg(pi, ["index", "--embedding", model], ctx, { signal, timeout: 300_000 });
+			result = await execZg(pi, ["index", "--mode", state.queryMode, "--embedding", model], ctx, {
+				signal,
+				timeout: 300_000,
+			});
 		}
 	}
 
@@ -190,16 +390,33 @@ export default function (pi: ExtensionAPI) {
 	pi.on("session_start", async (_event, ctx) => {
 		await refreshZgState(pi, ctx, state);
 
-		if (!pi.getFlag("no-zg-autostart") && state.available && !state.serverRunning) {
+		if (!pi.getFlag("no-zg-autostart") && state.queryMode === "server" && state.available && !state.serverRunning) {
 			// Fire-and-forget: don't block startup on daemon warmup. `zg server on`
 			// is idempotent, so this is safe even if something else started it
 			// in the meantime.
 			execZg(pi, ["server", "on"], ctx, { timeout: 20_000 })
-				.then(() => refreshZgState(pi, ctx, state))
+				.then(async (result) => {
+					await refreshZgState(pi, ctx, state);
+					if (result.code === 0) markServerOwnedByPi(state);
+				})
 				.catch(() => {});
 		}
 	});
 
+	pi.on("session_shutdown", async (_event, ctx) => {
+		if (!state.ownedServerPid) return;
+
+		// Recheck the PID before stopping anything. If another process replaced the
+		// daemon, it is no longer the server Pi started and must be left alone.
+		const server = await execZg(pi, ["server", "status", "--check-ready"], ctx, { timeout: 5_000 });
+		if (server.code !== 0 || parseServerPid(server.stdout) !== state.ownedServerPid) return;
+
+		await execZg(pi, ["server", "off"], ctx, { timeout: 20_000 });
+		state.serverRunning = false;
+		state.serverPid = undefined;
+		state.ownedServerPid = undefined;
+	});
+
 	pi.on("turn_start", async (_event, ctx) => {
 		if (state.available) {
 			await refreshZgState(pi, ctx, state);
@@ -231,13 +448,16 @@ export default function (pi: ExtensionAPI) {
 				);
 			}
 
+			if (state.queryMode === "server" && !(await ensureServerRunning(pi, ctx, state))) {
+				throw new Error("zg server is unavailable. Use /zg-settings to choose a direct search mode or /zg-server to inspect it.");
+			}
 			if (!state.indexed && !(await offerToBuildIndex(pi, ctx, state, signal))) {
 				throw new Error(
 					`This project does not appear to be indexed. Run \`zg index\` manually in ${ctx.cwd}, then retry.`,
 				);
 			}
 
-			const args = ["query"];
+			const args = ["query", "--mode", state.queryMode, "--refresh", state.refreshPolicy];
 			if (params.limit !== undefined) args.push("--limit", String(params.limit));
 			args.push(params.query);
 
@@ -323,7 +543,7 @@ export default function (pi: ExtensionAPI) {
 				throw new Error("`zg` is not available on PATH. Install @zvec/zvec-grep separately before indexing.");
 			}
 
-			const args = ["index"];
+			const args = ["index", "--mode", state.queryMode];
 			if (params.drop) args.push("--drop", "--yes");
 			else if (params.rebuild) args.push("--rebuild");
 
@@ -341,6 +561,13 @@ export default function (pi: ExtensionAPI) {
 		},
 	});
 
+	pi.registerCommand("zg-settings", {
+		description: "Configure Pi search mode plus zg embedding model, device, and provider credentials",
+		handler: async (_args, ctx) => {
+			await openZgSettings(pi, ctx, state);
+		},
+	});
+
 	pi.registerCommand("zg-status", {
 		description: "Report zg version, server, and index status for this project",
 		handler: async (_args, ctx) => {
@@ -354,6 +581,7 @@ export default function (pi: ExtensionAPI) {
 			const lines = [
 				`zg: ${state.version || "available"}`,
 				`server: ${state.serverRunning ? "running" : "stopped"}`,
+				`Pi search mode: ${searchModeLabel(state)} (session setting)`,
 				"",
 				(detail.stdout || detail.stderr).trim(),
 			];
@@ -378,7 +606,7 @@ export default function (pi: ExtensionAPI) {
 				if (!confirmed) return;
 			}
 
-			const cmdArgs = ["index"];
+			const cmdArgs = ["index", "--mode", state.queryMode];
 			if (drop) cmdArgs.push("--drop", "--yes");
 			else if (rebuild) cmdArgs.push("--rebuild");
 
@@ -417,8 +645,11 @@ export default function (pi: ExtensionAPI) {
 				if (!confirmed) return;
 			}
 
+			const wasServerRunning = state.serverRunning;
 			const result = await execZg(pi, ["server", action], ctx, { timeout: 20_000 });
 			await refreshZgState(pi, ctx, state);
+			if (action === "on" && !wasServerRunning && result.code === 0) markServerOwnedByPi(state);
+			if (action === "off" && result.code === 0) state.ownedServerPid = undefined;
 			ctx.ui.notify(
 				(result.stdout || result.stderr).trim() || `zg server ${action} done.`,
 				result.code === 0 ? "info" : "error",

package/package.json

@@ -1,9 +1,16 @@
 {
 	"name": "@krosskinetic/pi-zg",
-	"version": "0.1.0",
+	"version": "0.1.1",
 	"description": "Pi extension for zvec-grep (zg): semantic code search, managed ripgrep, and shared-server-aware index management",
 	"type": "module",
-	"keywords": ["pi-package", "pi", "pi-coding-agent", "zvec-grep", "semantic-search", "code-search"],
+	"keywords": [
+		"pi-package",
+		"pi",
+		"pi-coding-agent",
+		"zvec-grep",
+		"semantic-search",
+		"code-search"
+	],
 	"license": "MIT",
 	"author": "KrossKinetic",
 	"repository": {
@@ -13,9 +20,15 @@
 	"homepage": "https://github.com/KrossKinetic/pi-zg#readme",
 	"bugs": "https://github.com/KrossKinetic/pi-zg/issues",
 	"pi": {
-		"extensions": ["./extensions"]
+		"extensions": [
+			"./extensions"
+		]
 	},
-	"files": ["extensions", "README.md", "LICENSE"],
+	"files": [
+		"extensions",
+		"README.md",
+		"LICENSE"
+	],
 	"peerDependencies": {
 		"@earendil-works/pi-coding-agent": "*",
 		"typebox": "*"