From 5bc6afd15a9f05b77d957796e7998ceded6bcd1b Mon Sep 17 00:00:00 2001
From: anan <>
Date: Sat, 1 Aug 2026 21:03:25 +0800
Subject: [PATCH] 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
 }