From 52812adf5f114e10198cfc27cbf689f1034f08e7 Mon Sep 17 00:00:00 2001
From: Jay Zeng <jayzeng@users.noreply.github.com>
Date: Mon, 13 Apr 2026 23:52:30 -0700
Subject: [PATCH] fix: harden exit summary shutdown

Signed-off-by: Jay Zeng <jayzeng@users.noreply.github.com>
---
 AGENTS.md         |  5 +++++
 index.ts          | 39 ++++++++++++++++++++++++++++++---
 test/unit.test.ts | 55 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 95 insertions(+), 4 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index cbb9156..3c1991e 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -33,6 +33,11 @@ Runtime data lives outside the repo under `~/.pi/agent/memory/` (`MEMORY.md`, `S
 
 ## Testing Guidelines
 
+- Enforce TDD for every behavior change: follow `red -> green -> refactor`.
+- Start by establishing a verifiable baseline: run the relevant existing tests before edits, and record the exact command + outcome in the PR/commit notes.
+- Add or update a failing test first that reproduces the bug or captures the new requirement; implement code only after the test fails for the expected reason.
+- Keep tests green after implementation and after any refactor; do not merge with skipped failing tests.
+- Every bug fix must include a regression test that fails before the fix and passes after it.
 - Tests touch `~/.pi/agent/memory/`; ensure backups/restores remain intact and new tests don’t leak user data.
 - Prefer behavior-focused assertions (tool availability, file contents, cross-session recall). Keep timeouts generous for model latency.
 
diff --git a/index.ts b/index.ts
index 961258e..2f5c09c 100644
--- a/index.ts
+++ b/index.ts
@@ -341,8 +341,41 @@ function formatExitSummaryEntry(
 	return [`<!-- ${timestamp} [${sessionId}] -->`, header, "", summary.trim()].join("\n");
 }
 
+function getSessionBranch(ctx: ExtensionContext): SessionEntry[] | null {
+	const sessionManager = ctx.sessionManager as ExtensionContext["sessionManager"] & {
+		getBranch?: () => SessionEntry[];
+	};
+	if (typeof sessionManager?.getBranch !== "function") {
+		return null;
+	}
+	return sessionManager.getBranch();
+}
+
+async function resolveExitSummaryApiKey(ctx: ExtensionContext): Promise<string | undefined> {
+	if (!ctx.model) return undefined;
+
+	const modelRegistry = ctx.modelRegistry as ExtensionContext["modelRegistry"] & {
+		getApiKey?: (model: NonNullable<ExtensionContext["model"]>) => Promise<string | undefined>;
+		getApiKeyForProvider?: (provider: string) => Promise<string | undefined>;
+	};
+
+	if (typeof modelRegistry?.getApiKey === "function") {
+		return modelRegistry.getApiKey(ctx.model);
+	}
+
+	if (typeof modelRegistry?.getApiKeyForProvider === "function") {
+		return modelRegistry.getApiKeyForProvider(ctx.model.provider);
+	}
+
+	return undefined;
+}
+
 async function generateExitSummary(ctx: ExtensionContext): Promise<ExitSummaryResult> {
-	const branch = ctx.sessionManager.getBranch();
+	const branch = getSessionBranch(ctx);
+	if (!branch) {
+		return { summary: null, error: "Session branch unavailable", hasMessages: false };
+	}
+
 	const messages = branch
 		.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
 		.map((entry) => entry.message);
@@ -355,11 +388,11 @@ async function generateExitSummary(ctx: ExtensionContext): Promise<ExitSummaryRe
 		return { summary: null, error: "No active model", hasMessages: true };
 	}
 
-	const apiKey = await ctx.modelRegistry.getApiKey(ctx.model);
+	const apiKey = await resolveExitSummaryApiKey(ctx);
 	if (!apiKey) {
 		return {
 			summary: null,
-			error: `No API key for ${ctx.model.provider}/${ctx.model.id}`,
+			error: `API key resolution unavailable for ${ctx.model.provider}/${ctx.model.id}`,
 			hasMessages: true,
 		};
 	}
diff --git a/test/unit.test.ts b/test/unit.test.ts
index d87140e..49ad202 100644
--- a/test/unit.test.ts
+++ b/test/unit.test.ts
@@ -83,6 +83,27 @@ function createMockCtx(sessionId = "abcdef1234567890") {
 	};
 }
 
+function createShutdownCtx(options?: {
+	sessionId?: string;
+	branch?: any[];
+	model?: { provider: string; id: string };
+	modelRegistry?: Record<string, unknown>;
+}) {
+	const sessionId = options?.sessionId ?? "abcdef1234567890";
+	return {
+		sessionManager: {
+			getSessionId: () => sessionId,
+			getBranch: () => options?.branch ?? [],
+		},
+		model: options?.model,
+		modelRegistry: options?.modelRegistry ?? {},
+		hasUI: false,
+		ui: {
+			notify: mock(() => {}),
+		},
+	};
+}
+
 // We need to import the default export to register tools
 import registerExtension from "../index.js";
 
@@ -983,7 +1004,7 @@ describe("lifecycle hooks", () => {
 		_setQmdAvailable(true);
 		scheduleQmdUpdate();
 		expect(_getUpdateTimer()).not.toBeNull();
-		await hooks.session_shutdown({}, {});
+		await hooks.session_shutdown({}, createShutdownCtx());
 		expect(_getUpdateTimer()).toBeNull();
 	});
 
@@ -993,6 +1014,38 @@ describe("lifecycle hooks", () => {
 		await hooks.session_shutdown({}, {});
 	});
 
+	test("session_shutdown writes fallback summary when getApiKey is unavailable", async () => {
+		const ctx = createShutdownCtx({
+			branch: [
+				{
+					type: "message",
+					message: {
+						role: "user",
+						content: [{ type: "text", text: "Please remember we chose SQLite." }],
+						timestamp: Date.now(),
+					},
+				},
+				{
+					type: "message",
+					message: {
+						role: "assistant",
+						content: [{ type: "text", text: "Noted." }],
+						timestamp: Date.now(),
+					},
+				},
+			],
+			model: { provider: "openai", id: "gpt-4o-mini" },
+			modelRegistry: {},
+		});
+
+		await hooks.session_shutdown({}, ctx);
+
+		const content = fs.readFileSync(dailyPath(todayStr()), "utf-8");
+		expect(content).toContain("## Session Summary (auto, exit: session-end)");
+		expect(content).toContain("Auto-summary unavailable");
+		expect(content).toContain("API key resolution unavailable");
+	});
+
 	// -- session_before_compact --
 
 	test("session_before_compact appends handoff when scratchpad has open items", async () => {