From 5bc6afd15a9f05b77d957796e7998ceded6bcd1b Mon Sep 17 00:00:00 2001
From: anan <>
Date: Sat, 1 Aug 2026 21:03:25 +0800
Subject: [PATCH 1/2] fix: close MCP client on session_shutdown so pi -p can
 exit

The extension spawns the codebase-memory-mcp child process at load time
but never closes it. The child's stdio pipes keep the parent event loop
alive, so pi in non-interactive print mode (-p / --mode json) hangs
forever after answering instead of exiting.

Register an idempotent session_shutdown handler (the documented hook for
session-scoped resources) that closes the MCP client, letting the event
loop drain and the process exit normally.

Repro: pi -p "ping" with this extension loaded never exits.
After: exits cleanly; interactive sessions are unaffected because
session_shutdown only fires at session end/switch, not per turn.
---
 extensions/codebase-memory.ts | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/extensions/codebase-memory.ts b/extensions/codebase-memory.ts
index 8c188b0..61c9403 100644
--- a/extensions/codebase-memory.ts
+++ b/extensions/codebase-memory.ts
@@ -97,5 +97,12 @@ export default function codebaseMemoryExtension(pi: ExtensionAPI) {
 		});
 	}
 
+	// Close the MCP child process so the parent event loop can drain
+	// (without this, pi -p hangs after answering because stdio pipes stay open).
+	pi.on("session_shutdown", async () => {
+		await client?.close().catch(() => undefined);
+		client = undefined;
+	});
+
 	void register().catch(() => undefined); // silent if binary not installed yet
 }

From 83e389ad6cf2d6e65731f0c242b1e2bb4f376efe Mon Sep 17 00:00:00 2001
From: Richard Levitte <richard@levitte.org>
Date: Sat, 8 Aug 2026 10:05:38 +0200
Subject: [PATCH 2/2] Close in-flight connections on session_shutdown

If shutdown lands while connect() is still pending, client is
undefined and the handler no-ops; the spawn then completes into
a leak that holds the event loop open just the same.

Assisted-by: Pi:moonshot/kimi-k3
---
 extensions/codebase-memory.ts | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/extensions/codebase-memory.ts b/extensions/codebase-memory.ts
index 61c9403..a51e2fd 100644
--- a/extensions/codebase-memory.ts
+++ b/extensions/codebase-memory.ts
@@ -28,6 +28,7 @@ const HINT =
 
 export default function codebaseMemoryExtension(pi: ExtensionAPI) {
 	let client: Client | undefined;
+	let shuttingDown = false;
 	const seen = new Set<string>();
 
 	async function getClient(reconnect = false): Promise<Client> {
@@ -38,6 +39,13 @@ export default function codebaseMemoryExtension(pi: ExtensionAPI) {
 		if (client) return client;
 		const c = new Client({ name: "pi-codebase-memory", version: "0.1.2" });
 		await c.connect(new StdioClientTransport({ command: BIN, args: [], env: process.env as Record<string, string> }));
+		if (shuttingDown) {
+			// Shutdown landed while connect() was in flight; don't adopt
+			// the freshly spawned server, or it leaks and holds the event
+			// loop open just the same.
+			await c.close().catch(() => {});
+			throw new Error("session is shutting down");
+		}
 		client = c;
 		return c;
 	}
@@ -100,6 +108,7 @@ export default function codebaseMemoryExtension(pi: ExtensionAPI) {
 	// Close the MCP child process so the parent event loop can drain
 	// (without this, pi -p hangs after answering because stdio pipes stay open).
 	pi.on("session_shutdown", async () => {
+		shuttingDown = true;
 		await client?.close().catch(() => undefined);
 		client = undefined;
 	});