Improve compaction hooks: add signal, no timeout, SessionManager cleanup, docs

packages/coding-agent/src/core/agent-session.ts

@@ -24,7 +24,7 @@ import { exportSessionToHtml } from "./export-html.js";
 import type { HookRunner, SessionEventResult, TurnEndEvent, TurnStartEvent } from "./hooks/index.js";
 import type { BashExecutionMessage } from "./messages.js";
 import { getApiKeyForModel, getAvailableModels } from "./model-config.js";
-import { type CompactionEntry, loadSessionFromEntries, type SessionManager } from "./session-manager.js";
+import type { CompactionEntry, SessionManager } from "./session-manager.js";
 import type { SettingsManager, SkillsSettings } from "./settings-manager.js";
 import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
 
@@ -510,7 +510,7 @@ export class AgentSession {
 	 */
 	async reset(): Promise<boolean> {
 		const previousSessionFile = this.sessionFile;
-		const entries = this.sessionManager.loadEntries();
+		const entries = this.sessionManager.getEntries();
 
 		// Emit before_clear event (can be cancelled)
 		if (this._hookRunner?.hasHandlers("session")) {
@@ -748,7 +748,7 @@ export class AgentSession {
 				throw new Error(`No API key for ${this.model.provider}`);
 			}
 
-			const entries = this.sessionManager.loadEntries();
+			const entries = this.sessionManager.getEntries();
 			const settings = this.settingsManager.getCompactionSettings();
 
 			const preparation = prepareCompaction(entries, settings);
@@ -783,6 +783,7 @@ export class AgentSession {
 					customInstructions,
 					model: this.model,
 					resolveApiKey: this._resolveApiKey,
+					signal: this._compactionAbortController.signal,
 				})) as SessionEventResult | undefined;
 
 				if (result?.cancel) {
@@ -811,9 +812,9 @@ export class AgentSession {
 			}
 
 			this.sessionManager.saveCompaction(compactionEntry);
-			const newEntries = this.sessionManager.loadEntries();
-			const loaded = loadSessionFromEntries(newEntries);
-			this.agent.replaceMessages(loaded.messages);
+			const newEntries = this.sessionManager.getEntries();
+			const sessionContext = this.sessionManager.buildSessionContext();
+			this.agent.replaceMessages(sessionContext.messages);
 
 			if (this._hookRunner) {
 				await this._hookRunner.emit({
@@ -909,7 +910,7 @@ export class AgentSession {
 				return;
 			}
 
-			const entries = this.sessionManager.loadEntries();
+			const entries = this.sessionManager.getEntries();
 
 			const preparation = prepareCompaction(entries, settings);
 			if (!preparation) {
@@ -944,6 +945,7 @@ export class AgentSession {
 					customInstructions: undefined,
 					model: this.model,
 					resolveApiKey: this._resolveApiKey,
+					signal: this._autoCompactionAbortController.signal,
 				})) as SessionEventResult | undefined;
 
 				if (hookResult?.cancel) {
@@ -973,9 +975,9 @@ export class AgentSession {
 			}
 
 			this.sessionManager.saveCompaction(compactionEntry);
-			const newEntries = this.sessionManager.loadEntries();
-			const loaded = loadSessionFromEntries(newEntries);
-			this.agent.replaceMessages(loaded.messages);
+			const newEntries = this.sessionManager.getEntries();
+			const sessionContext = this.sessionManager.buildSessionContext();
+			this.agent.replaceMessages(sessionContext.messages);
 
 			if (this._hookRunner) {
 				await this._hookRunner.emit({
@@ -1281,7 +1283,7 @@ export class AgentSession {
 	 */
 	async switchSession(sessionPath: string): Promise<boolean> {
 		const previousSessionFile = this.sessionFile;
-		const oldEntries = this.sessionManager.loadEntries();
+		const oldEntries = this.sessionManager.getEntries();
 
 		// Emit before_switch event (can be cancelled)
 		if (this._hookRunner?.hasHandlers("session")) {
@@ -1306,8 +1308,8 @@ export class AgentSession {
 		this.sessionManager.setSessionFile(sessionPath);
 
 		// Reload messages
-		const entries = this.sessionManager.loadEntries();
-		const loaded = loadSessionFromEntries(entries);
+		const entries = this.sessionManager.getEntries();
+		const sessionContext = this.sessionManager.buildSessionContext();
 
 		// Emit session event to hooks
 		if (this._hookRunner) {
@@ -1324,22 +1326,22 @@ export class AgentSession {
 		// Emit session event to custom tools
 		await this._emitToolSessionEvent("switch", previousSessionFile);
 
-		this.agent.replaceMessages(loaded.messages);
+		this.agent.replaceMessages(sessionContext.messages);
 
 		// Restore model if saved
-		const savedModel = this.sessionManager.loadModel();
-		if (savedModel) {
+		if (sessionContext.model) {
 			const availableModels = (await getAvailableModels()).models;
-			const match = availableModels.find((m) => m.provider === savedModel.provider && m.id === savedModel.modelId);
+			const match = availableModels.find(
+				(m) => m.provider === sessionContext.model!.provider && m.id === sessionContext.model!.modelId,
+			);
 			if (match) {
 				this.agent.setModel(match);
 			}
 		}
 
 		// Restore thinking level if saved (setThinkingLevel clamps to model capabilities)
-		const savedThinking = this.sessionManager.loadThinkingLevel();
-		if (savedThinking) {
-			this.setThinkingLevel(savedThinking as ThinkingLevel);
+		if (sessionContext.thinkingLevel) {
+			this.setThinkingLevel(sessionContext.thinkingLevel as ThinkingLevel);
 		}
 
 		this._reconnectToAgent();
@@ -1357,7 +1359,7 @@ export class AgentSession {
 	 */
 	async branch(entryIndex: number): Promise<{ selectedText: string; cancelled: boolean }> {
 		const previousSessionFile = this.sessionFile;
-		const entries = this.sessionManager.loadEntries();
+		const entries = this.sessionManager.getEntries();
 		const selectedEntry = entries[entryIndex];
 
 		if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") {
@@ -1394,8 +1396,8 @@ export class AgentSession {
 		}
 
 		// Reload messages from entries (works for both file and in-memory mode)
-		const newEntries = this.sessionManager.loadEntries();
-		const loaded = loadSessionFromEntries(newEntries);
+		const newEntries = this.sessionManager.getEntries();
+		const sessionContext = this.sessionManager.buildSessionContext();
 
 		// Emit branch event to hooks (after branch completes)
 		if (this._hookRunner) {
@@ -1414,7 +1416,7 @@ export class AgentSession {
 		await this._emitToolSessionEvent("branch", previousSessionFile);
 
 		if (!skipConversationRestore) {
-			this.agent.replaceMessages(loaded.messages);
+			this.agent.replaceMessages(sessionContext.messages);
 		}
 
 		return { selectedText, cancelled: false };
@@ -1424,7 +1426,7 @@ export class AgentSession {
 	 * Get all user messages from session for branch selector.
 	 */
 	getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> {
-		const entries = this.sessionManager.loadEntries();
+		const entries = this.sessionManager.getEntries();
 		const result: Array<{ entryIndex: number; text: string }> = [];
 
 		for (let i = 0; i < entries.length; i++) {
@@ -1570,7 +1572,7 @@ export class AgentSession {
 		previousSessionFile: string | null,
 	): Promise<void> {
 		const event: ToolSessionEvent = {
-			entries: this.sessionManager.loadEntries(),
+			entries: this.sessionManager.getEntries(),
 			sessionFile: this.sessionFile,
 			previousSessionFile,
 			reason,

packages/coding-agent/src/core/hooks/runner.ts

@@ -11,6 +11,7 @@ import type {
 	HookEvent,
 	HookEventContext,
 	HookUIContext,
+	SessionEvent,
 	SessionEventResult,
 	ToolCallEvent,
 	ToolCallEventResult,
@@ -229,9 +230,17 @@ export class HookRunner {
 
 			for (const handler of handlers) {
 				try {
-					const timeout = createTimeout(this.timeout);
-					const handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
-					timeout.clear();
+					// No timeout for before_compact events (like tool_call, they may take a while)
+					const isBeforeCompact = event.type === "session" && (event as SessionEvent).reason === "before_compact";
+					let handlerResult: unknown;
+
+					if (isBeforeCompact) {
+						handlerResult = await handler(event, ctx);
+					} else {
+						const timeout = createTimeout(this.timeout);
+						handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
+						timeout.clear();
+					}
 
 					// For session events, capture the result (for before_* cancellation)
 					if (event.type === "session" && handlerResult) {

packages/coding-agent/src/core/hooks/types.ts

@@ -141,6 +141,8 @@ export type SessionEvent =
 			model: Model<any>;
 			/** Resolve API key for any model (checks settings, OAuth, env vars) */
 			resolveApiKey: (model: Model<any>) => Promise<string | undefined>;
+			/** Abort signal - hooks should pass this to LLM calls and check it periodically */
+			signal: AbortSignal;
 	  })
 	| (SessionEventBase & {
 			reason: "compact";

packages/coding-agent/src/core/model-config.ts

@@ -380,8 +380,13 @@ export async function getAvailableModels(
 }
 
 /**
- * Find a specific model by provider and ID
- * Returns { model, error } - either model or error message
+ * Find a specific model by provider and ID.
+ *
+ * Searches models from:
+ * 1. Built-in models from @mariozechner/pi-ai
+ * 2. Custom models defined in ~/.pi/agent/models.json
+ *
+ * Returns { model, error } - either the model or an error message.
  */
 export function findModel(
 	provider: string,

packages/coding-agent/src/core/sdk.ts

@@ -496,7 +496,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 	};
 
 	// Check if session has existing data to restore
-	const existingSession = sessionManager.loadSession();
+	const existingSession = sessionManager.buildSessionContext();
 	time("loadSession");
 	const hasExistingSession = existingSession.messages.length > 0;
 

packages/coding-agent/src/core/session-manager.ts

@@ -54,7 +54,7 @@ export type SessionEntry =
 	| ModelChangeEntry
 	| CompactionEntry;
 
-export interface LoadedSession {
+export interface SessionContext {
 	messages: AppMessage[];
 	thinkingLevel: string;
 	model: { provider: string; modelId: string } | null;
@@ -78,6 +78,7 @@ export const SUMMARY_PREFIX = `The conversation history before this point was co
 export const SUMMARY_SUFFIX = `
 </summary>`;
 
+/** Exported for compaction.test.ts */
 export function createSummaryMessage(summary: string): AppMessage {
 	return {
 		role: "user",
@@ -86,6 +87,7 @@ export function createSummaryMessage(summary: string): AppMessage {
 	};
 }
 
+/** Exported for compaction.test.ts */
 export function parseSessionEntries(content: string): SessionEntry[] {
 	const entries: SessionEntry[] = [];
 	const lines = content.trim().split("\n");
@@ -112,7 +114,15 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
 	return null;
 }
 
-export function loadSessionFromEntries(entries: SessionEntry[]): LoadedSession {
+/**
+ * Build the session context from entries. This is what gets sent to the LLM.
+ *
+ * If there's a compaction entry, returns the summary message plus messages
+ * from `firstKeptEntryIndex` onwards. Otherwise returns all messages.
+ *
+ * Also extracts the current thinking level and model from the entries.
+ */
+export function buildSessionContext(entries: SessionEntry[]): SessionContext {
 	let thinkingLevel = "off";
 	let model: { provider: string; modelId: string } | null = null;
 
@@ -299,7 +309,7 @@ export class SessionManager {
 		}
 	}
 
-	saveMessage(message: any): void {
+	saveMessage(message: AppMessage): void {
 		const entry: SessionMessageEntry = {
 			type: "message",
 			timestamp: new Date().toISOString(),
@@ -335,29 +345,21 @@ export class SessionManager {
 		this._persist(entry);
 	}
 
-	loadSession(): LoadedSession {
-		const entries = this.loadEntries();
-		return loadSessionFromEntries(entries);
+	/**
+	 * Build the session context (what gets sent to the LLM).
+	 * If compacted, returns summary + kept messages. Otherwise all messages.
+	 * Includes thinking level and model.
+	 */
+	buildSessionContext(): SessionContext {
+		return buildSessionContext(this.getEntries());
 	}
 
-	loadMessages(): AppMessage[] {
-		return this.loadSession().messages;
-	}
-
-	loadThinkingLevel(): string {
-		return this.loadSession().thinkingLevel;
-	}
-
-	loadModel(): { provider: string; modelId: string } | null {
-		return this.loadSession().model;
-	}
-
-	loadEntries(): SessionEntry[] {
-		if (this.inMemoryEntries.length > 0) {
-			return [...this.inMemoryEntries];
-		} else {
-			return loadEntriesFromFile(this.sessionFile);
-		}
+	/**
+	 * Get all session entries. Returns a defensive copy.
+	 * Use buildSessionContext() if you need the messages for the LLM.
+	 */
+	getEntries(): SessionEntry[] {
+		return [...this.inMemoryEntries];
 	}
 
 	createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {

packages/coding-agent/src/index.ts

@@ -119,13 +119,13 @@ export {
 	readOnlyTools,
 } from "./core/sdk.js";
 export {
+	buildSessionContext,
 	type CompactionEntry,
 	createSummaryMessage,
 	getLatestCompactionEntry,
-	type LoadedSession,
-	loadSessionFromEntries,
 	type ModelChangeEntry,
 	parseSessionEntries,
+	type SessionContext as LoadedSession,
 	type SessionEntry,
 	type SessionHeader,
 	type SessionInfo,

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -351,7 +351,7 @@ export class InteractiveMode {
 		}
 
 		// Load session entries if any
-		const entries = this.session.sessionManager.loadEntries();
+		const entries = this.session.sessionManager.getEntries();
 
 		// Set TUI-based UI context for custom tools
 		const uiContext = this.createHookUIContext();
@@ -1067,7 +1067,7 @@ export class InteractiveMode {
 			this.updateEditorBorderColor();
 		}
 
-		const compactionEntry = getLatestCompactionEntry(this.sessionManager.loadEntries());
+		const compactionEntry = getLatestCompactionEntry(this.sessionManager.getEntries());
 
 		for (const message of messages) {
 			if (isBashExecutionMessage(message)) {
@@ -1137,7 +1137,7 @@ export class InteractiveMode {
 		this.renderMessages(state.messages, { updateFooter: true, populateHistory: true });
 
 		// Show compaction info if session was compacted
-		const entries = this.sessionManager.loadEntries();
+		const entries = this.sessionManager.getEntries();
 		const compactionCount = entries.filter((e) => e.type === "compaction").length;
 		if (compactionCount > 0) {
 			const times = compactionCount === 1 ? "1 time" : `${compactionCount} times`;
@@ -1185,7 +1185,7 @@ export class InteractiveMode {
 		// Emit shutdown event to hooks
 		const hookRunner = this.session.hookRunner;
 		if (hookRunner?.hasHandlers("session")) {
-			const entries = this.sessionManager.loadEntries();
+			const entries = this.sessionManager.getEntries();
 			await hookRunner.emit({
 				type: "session",
 				entries,
@@ -1924,7 +1924,7 @@ export class InteractiveMode {
 	}
 
 	private async handleCompactCommand(customInstructions?: string): Promise<void> {
-		const entries = this.sessionManager.loadEntries();
+		const entries = this.sessionManager.getEntries();
 		const messageCount = entries.filter((e) => e.type === "message").length;
 
 		if (messageCount < 2) {

packages/coding-agent/src/modes/print-mode.ts

@@ -28,7 +28,7 @@ export async function runPrintMode(
 	initialAttachments?: Attachment[],
 ): Promise<void> {
 	// Load entries once for session start events
-	const entries = session.sessionManager.loadEntries();
+	const entries = session.sessionManager.getEntries();
 
 	// Hook runner already has no-op UI context by default (set in main.ts)
 	// Set up hooks for print mode (no UI)

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -121,7 +121,7 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 	});
 
 	// Load entries once for session start events
-	const entries = session.sessionManager.loadEntries();
+	const entries = session.sessionManager.getEntries();
 
 	// Set up hooks with RPC-based UI context
 	const hookRunner = session.hookRunner;