fix(coding-agent): keep prompt and tool state across context handlers (#9846)

Since 0.86 the prompt and tool declarations live in transcript system
messages, so a context handler that filters or slices messages could drop
them. After extension-driven compaction this sent requests without built-in
tools and made Codex emit raw tool-call text.

context handlers now see the conversation only. An unchanged list keeps the
transcript as is; a changed list gets the replayed prompt sections and tool
declarations as one leading system message. Handler-added system messages
are kept after it.

Add context_with_system, which runs after every context handler on the full
transcript and sends its result verbatim, for extensions that need to edit
system messages per request. Dropping the leading system message is reported
as an extension error.

closes #9789, closes #9822

packages/coding-agent/CHANGELOG.md

@@ -19,6 +19,7 @@
 - Added append-only model-context edits. For example, `sessionManager.appendContextEdit(entryId, null)` omits one message from future provider context without changing raw history, usage, or UI history.
 - Added actionable `turn_end` and `agent_before_settle` extension boundaries. Return `{ entries: [...event.entries, draft], continue: true }` to persist structural entries in order and ensure one next provider request without changing steering or follow-up scheduling.
 - Added retain-none compaction input: `sessionManager.appendCompaction(summary, null, tokensBefore)` stores the compaction's own ID as its kept boundary.
+- Added the `context_with_system` extension event, which runs after `context` handlers on the full transcript including system messages and sends its result verbatim. See [`context_with_system`](docs/extensions.md#context_with_system).
 - Added per-model image resize profiles through `inputLimits.images.resize` in `models.json`, applied to file attachments, image reads, and tool-result images ([#9631](https://github.com/earendil-works/pi/issues/9631)).
 
 ### Fixed
@@ -27,6 +28,7 @@
 - Fixed context-invisible boundary metadata and replacement edits causing newly appended or replaced input to be summarized before its first provider request.
 - Fixed edited-context accounting both discarding valid assistant usage captured after the latest context edit and reusing that usage after a later compaction made it stale.
 - Fixed selected error retries and final length/overflow recovery retaining abandoned model attempts in future provider context; post-run recovery omissions are now persisted without hiding raw transcript history or changing queue scheduling.
+- Fixed `context` handlers that filter or slice messages dropping the prompt and tool declarations, which after extension-driven compaction left requests without built-in tools or made Codex emit raw tool-call text. Handlers no longer see system messages; Pi restores the prompt and tool state after they run. See [`context`](docs/extensions.md#context) ([#9789](https://github.com/earendil-works/pi/issues/9789), [#9822](https://github.com/earendil-works/pi/issues/9822)).
 - Fixed `/bug` allowing uploads in offline mode while preserving local zip exports.
 - Fixed idle prompt-cache warming rebuilding expired caches when its timer or an extension decision is delayed.
 - Improved crash diagnostics with hints identifying loaded extensions that appear in the stack trace.

packages/coding-agent/docs/extensions.md

@@ -294,7 +294,8 @@ user sends prompt ────────────────────
   │   ┌─── turn (repeats while LLM calls tools) ───┐       │
   │   │                                            │       │
   │   ├─► turn_start                               │       │
-  │   ├─► context (can modify messages)            │       │
+  │   ├─► context (can modify conversation messages)   │
+  │   ├─► context_with_system (can modify the full transcript)
   │   ├─► before_provider_headers (can mutate headers)     |
   │   ├─► before_provider_request (can inspect or replace payload)
   │   ├─► after_provider_response (status + headers, before stream consume)
@@ -725,12 +726,31 @@ Fired before each LLM call. Modify messages non-destructively. See [Session Form
 
 ```typescript
 pi.on("context", async (event, ctx) => {
-  // event.messages - deep copy, safe to modify
+  // event.messages - deep copy without system messages, safe to modify
   const filtered = event.messages.filter(m => !shouldPrune(m));
   return { messages: filtered };
 });
 ```
 
+`event.messages` holds the conversation without system messages. The prompt and tool declarations belong to Pi and are not part of this hook: when the handler returns a changed list, Pi replays the current prompt sections and tool declarations into one leading system message ahead of the returned messages. Filtering, windowing, or slicing from a compaction summary therefore cannot drop the prompt or the tools. An unchanged list keeps mid-conversation system messages in place, so models that accept them retain their cached prefix. System messages a handler adds are kept after Pi's head. To change the prompt or the tool set durably, use [`before_agent_start`](#before_agent_start) or `pi.setActiveTools()`; to edit system messages for one request, use [`context_with_system`](#context_with_system).
+
+#### context_with_system
+
+Fired before each LLM call, after every `context` handler has run and Pi has restored the prompt and tool state. `event.messages` is the full transcript, including the leading system message and any mid-conversation prompt or tool patches (see [Session Format](session-format.md#sessionmessageentry)). The returned messages are sent as they are: this hook owns the prompt and tool declarations for the request.
+
+```typescript
+import { getCurrentSystemMessage } from "@earendil-works/pi-ai";
+
+pi.on("context_with_system", async (event, ctx) => {
+  const cut = findCutIndex(event.messages);
+  // Fold the dropped prefix so its prompt and tool state survives as the new head.
+  const head = getCurrentSystemMessage(event.messages.slice(0, cut));
+  return { messages: head ? [head, ...event.messages.slice(cut)] : event.messages.slice(cut) };
+});
+```
+
+Rules: keep a system message at index 0 (providers read the prompt and initial tool declarations there; Pi reports an error if a handler drops it). Removing a system message removes the tool declarations and section patches it carries. Check your output with `getCurrentSystemPrompt()` and `getCurrentTools()` from `@earendil-works/pi-ai`. Handlers run in extension load order; a `systemPrompt` forced from `before_agent_start` is still projected onto the request afterwards.
+
 #### before_provider_headers
 
 Fired after the outgoing HTTP headers are assembled. Use it to add, override, or remove request headers.

packages/coding-agent/src/core/extensions/index.ts

@@ -57,6 +57,7 @@ export type {
 	// Event Results
 	ContextEventResult,
 	ContextUsage,
+	ContextWithSystemEvent,
 	CustomEntryDraft,
 	CustomMessageEntryDraft,
 	CustomToolCallEvent,

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

@@ -3,7 +3,13 @@
  */
 
 import type { AgentMessage } from "@earendil-works/pi-agent-core";
-import type { ImageContent, Model, Provider, ProviderHeaders } from "@earendil-works/pi-ai";
+import {
+	getCurrentSystemMessage,
+	type ImageContent,
+	type Model,
+	type Provider,
+	type ProviderHeaders,
+} from "@earendil-works/pi-ai";
 import type { KeyId } from "@earendil-works/pi-tui";
 import { type Theme, theme } from "../../modes/interactive/theme/theme.ts";
 import type { CacheWarmingAction } from "../cache-warmer.ts";
@@ -32,6 +38,7 @@ import type {
 	ContextEvent,
 	ContextEventResult,
 	ContextUsage,
+	ContextWithSystemEvent,
 	EntryRenderer,
 	Extension,
 	ExtensionActions,
@@ -169,6 +176,7 @@ type RunnerEmitEvent = Exclude<
 	| ToolResultEvent
 	| UserBashEvent
 	| ContextEvent
+	| ContextWithSystemEvent
 	| CacheWarmingDecisionEvent
 	| BeforeProviderRequestEvent
 	| BeforeProviderHeadersEvent
@@ -258,6 +266,28 @@ function snapshotEventHandlers(extensions: Extension[], event: ExtensionEvent["t
 	return extensions.map((ext) => ({ ext, handlers: ext.handlers.get(event)?.slice() ?? [] }));
 }
 
+function sameMessages(left: AgentMessage[], right: AgentMessage[]): boolean {
+	return left.length === right.length && left.every((message, index) => message === right[index]);
+}
+
+/**
+ * Re-attach the prompt and tool state after a `context` handler. Handlers only see the
+ * conversation; the system messages belong to Pi. An unchanged conversation keeps every
+ * system message in place, so models with mid-conversation support keep their cached
+ * prefix. A changed one gets the replayed prompt sections and tool declarations as one
+ * leading system message, so pruning, windowing, or slicing from a compaction summary
+ * cannot drop them.
+ */
+function restoreSystemMessages(
+	current: AgentMessage[],
+	visible: AgentMessage[],
+	returned: AgentMessage[],
+): AgentMessage[] {
+	if (sameMessages(returned, visible)) return current;
+	const head = getCurrentSystemMessage(current);
+	return head ? [head, ...returned] : returned;
+}
+
 export async function emitProjectTrustEvent(
 	extensionsResult: LoadExtensionsResult,
 	event: ProjectTrustEvent,
@@ -1152,6 +1182,11 @@ export class ExtensionRunner {
 		return undefined;
 	}
 
+	/**
+	 * Run the request-time transforms in two phases. `context` handlers see the conversation
+	 * only and Pi restores the prompt and tool state after each; `context_with_system`
+	 * handlers then see the full transcript and their output is used as returned.
+	 */
 	async emitContext(messages: AgentMessage[]): Promise<AgentMessage[]> {
 		const ctx = this.createContext();
 		let currentMessages = structuredClone(messages);
@@ -1159,12 +1194,17 @@ export class ExtensionRunner {
 		for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "context")) {
 			for (const handler of handlers) {
 				try {
-					const event: ContextEvent = { type: "context", messages: currentMessages };
-					const handlerResult = await handler(event, ctx);
+					const visibleMessages = currentMessages.filter((message) => message.role !== "system");
+					const visibleSnapshot = visibleMessages.slice();
+					const event: ContextEvent = { type: "context", messages: visibleMessages };
+					const handlerResult = (await handler(event, ctx)) as ContextEventResult | undefined;
 
-					if (handlerResult && (handlerResult as ContextEventResult).messages) {
-						currentMessages = (handlerResult as ContextEventResult).messages!;
-					}
+					// Handlers may return a new list or edit event.messages in place.
+					const returned =
+						handlerResult?.messages ??
+						(sameMessages(visibleMessages, visibleSnapshot) ? undefined : visibleMessages);
+					if (!returned) continue;
+					currentMessages = restoreSystemMessages(currentMessages, visibleSnapshot, returned);
 				} catch (err) {
 					const message = err instanceof Error ? err.message : String(err);
 					const stack = err instanceof Error ? err.stack : undefined;
@@ -1178,6 +1218,35 @@ export class ExtensionRunner {
 			}
 		}
 
+		for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "context_with_system")) {
+			for (const handler of handlers) {
+				try {
+					const hadLeadingSystemMessage = currentMessages[0]?.role === "system";
+					const event: ContextWithSystemEvent = { type: "context_with_system", messages: currentMessages };
+					const handlerResult = (await handler(event, ctx)) as ContextEventResult | undefined;
+					currentMessages = handlerResult?.messages ?? currentMessages;
+					// Providers read the prompt and initial tools from the leading system message.
+					// Losing it is never intended; report it but honor the handler's output.
+					if (hadLeadingSystemMessage && currentMessages[0]?.role !== "system") {
+						this.emitError({
+							extensionPath: ext.path,
+							event: "context_with_system",
+							error: "Handler removed the leading system message; the request has no prompt or initial tool declarations. Keep it at index 0 or replace a dropped prefix with getCurrentSystemMessage().",
+						});
+					}
+				} catch (err) {
+					const message = err instanceof Error ? err.message : String(err);
+					const stack = err instanceof Error ? err.stack : undefined;
+					this.emitError({
+						extensionPath: ext.path,
+						event: "context_with_system",
+						error: message,
+						stack,
+					});
+				}
+			}
+		}
+
 		return currentMessages;
 	}
 

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

@@ -688,12 +688,28 @@ export type SessionEvent =
 // Agent Events
 // ============================================================================
 
-/** Fired before each LLM call. Can modify messages. */
+/**
+ * Fired before each LLM call. Can modify messages.
+ *
+ * `messages` holds the conversation without system messages. The prompt and tool state
+ * belong to Pi: it restores them after the handler returns, so a handler cannot drop
+ * them and does not need to preserve them.
+ */
 export interface ContextEvent {
 	type: "context";
 	messages: AgentMessage[];
 }
 
+/**
+ * Fired before each LLM call, after every `context` handler has run and Pi has restored
+ * the prompt and tool state. `messages` is the full transcript including system messages,
+ * and the result is sent as returned: the handler owns the prompt and tool declarations.
+ */
+export interface ContextWithSystemEvent {
+	type: "context_with_system";
+	messages: AgentMessage[];
+}
+
 /** Fired before a provider request is sent. Can replace the payload. */
 export interface BeforeProviderRequestEvent {
 	type: "before_provider_request";
@@ -1156,6 +1172,7 @@ export type ExtensionEvent =
 	| ResourcesDiscoverEvent
 	| SessionEvent
 	| ContextEvent
+	| ContextWithSystemEvent
 	| CacheWarmingDecisionEvent
 	| BeforeProviderRequestEvent
 	| BeforeProviderHeadersEvent
@@ -1362,6 +1379,7 @@ export interface ExtensionAPI {
 	): () => void;
 	on(event: "session_tree", handler: ExtensionHandler<SessionTreeEvent>): () => void;
 	on(event: "context", handler: ExtensionHandler<ContextEvent, ContextEventResult>): () => void;
+	on(event: "context_with_system", handler: ExtensionHandler<ContextWithSystemEvent, ContextEventResult>): () => void;
 	on(
 		event: "cache_warming_decision",
 		handler: ExtensionHandler<CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult>,

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

@@ -51,6 +51,7 @@ export {
 	type CompactionEntryDraft,
 	type ContextEditEntryDraft,
 	type ContextEvent,
+	type ContextWithSystemEvent,
 	type CustomEntryDraft,
 	type CustomMessageEntryDraft,
 	defineTool,

packages/coding-agent/src/index.ts

@@ -81,6 +81,7 @@ export type {
 	ContextEvent,
 	ContextEventResult,
 	ContextUsage,
+	ContextWithSystemEvent,
 	CustomEntryDraft,
 	CustomMessageEntryDraft,
 	CustomToolCallEvent,

packages/coding-agent/test/suite/regressions/9789-context-handler-system-messages.test.ts

@@ -0,0 +1,226 @@
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
+import {
+	fauxAssistantMessage,
+	getCurrentSystemPrompt,
+	getCurrentTools,
+	type TranscriptContext,
+} from "@earendil-works/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ExtensionFactory } from "../../../src/index.ts";
+import { createHarness, type Harness } from "../harness.ts";
+
+/** Compaction supplied by an extension hook, as in the reported sessions. */
+const compactViaHook: ExtensionFactory = (pi) => {
+	pi.on("session_before_compact", async (event) => ({
+		compaction: {
+			summary: "extension summary",
+			firstKeptEntryId: event.preparation.firstKeptEntryId,
+			tokensBefore: event.preparation.tokensBefore,
+			details: { source: "test" },
+		},
+	}));
+};
+
+function captureRequest(harness: Harness, text: string): () => TranscriptContext {
+	let request: TranscriptContext | undefined;
+	harness.setResponses([
+		(context) => {
+			request = context;
+			return fauxAssistantMessage(text);
+		},
+	]);
+	return () => {
+		if (!request) throw new Error("expected a provider request");
+		return request;
+	};
+}
+
+function toolNames(context: TranscriptContext): string[] {
+	return getCurrentTools(context.messages).map((tool) => tool.name);
+}
+
+async function compactSession(harness: Harness): Promise<void> {
+	harness.settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } });
+	harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]);
+	await harness.session.prompt("first");
+	await harness.session.prompt("second");
+	await harness.session.compact();
+	expect(harness.session.messages.map((message) => message.role).slice(0, 2)).toEqual(["system", "compactionSummary"]);
+}
+
+describe("context handlers and system messages", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) harnesses.pop()?.cleanup();
+	});
+
+	// Regression #9789, #9822: pruning from the compaction summary dropped the prompt and tool checkpoint.
+	it("keeps the prompt and tools when a handler slices from the compaction summary", async () => {
+		const seen: AgentMessage[][] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				compactViaHook,
+				(pi) => {
+					pi.on("context", async (event) => {
+						seen.push(event.messages);
+						const summary = event.messages.findIndex((message) => message.role === "compactionSummary");
+						return { messages: event.messages.slice(summary) };
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		await compactSession(harness);
+		const getRequest = captureRequest(harness, "after compaction");
+
+		await harness.session.prompt("third");
+
+		const request = getRequest();
+		expect(seen.at(-1)?.some((message) => message.role === "system")).toBe(false);
+		expect(request.messages[0]?.role).toBe("system");
+		expect(toolNames(request)).toEqual(harness.session.getActiveToolNames());
+		expect(getCurrentSystemPrompt(request.messages)).toBe(harness.session.systemPrompt);
+		expect(request.messages.filter((message) => message.role === "system")).toHaveLength(1);
+	});
+
+	it("keeps mid-conversation system messages in place when a handler leaves the conversation unchanged", async () => {
+		let turn = 0;
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("before_agent_start", (event) => {
+						if (++turn === 2) event.systemPromptOptions.sections.plan_mode = "Plan only.";
+					});
+					pi.on("context", async (event) => ({ messages: event.messages }));
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("one")]);
+		await harness.session.prompt("first");
+		const getRequest = captureRequest(harness, "two");
+
+		await harness.session.prompt("second");
+
+		const systemMessages = getRequest().messages.filter((message) => message.role === "system");
+		expect(systemMessages).toHaveLength(2);
+		expect(systemMessages[1]?.sections).toEqual({ plan_mode: "<plan_mode>\nPlan only.\n</plan_mode>" });
+	});
+
+	it("applies in-place edits to event.messages without a return value", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("context", async (event) => {
+						event.messages.splice(0, 0, {
+							role: "user",
+							content: [{ type: "text", text: "injected" }],
+							timestamp: 0,
+						});
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		const getRequest = captureRequest(harness, "done");
+
+		await harness.session.prompt("hello");
+
+		const request = getRequest();
+		expect(request.messages.map((message) => message.role)).toEqual(["system", "user", "user"]);
+		expect(toolNames(request)).toEqual(harness.session.getActiveToolNames());
+	});
+
+	it("keeps system messages a handler adds after the replayed head", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("context", async (event) => ({
+						messages: [{ role: "system", content: "ephemeral reminder", timestamp: 0 }, ...event.messages],
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		const getRequest = captureRequest(harness, "done");
+
+		await harness.session.prompt("hello");
+
+		const request = getRequest();
+		expect(request.messages.map((message) => message.role)).toEqual(["system", "system", "user"]);
+		expect(toolNames(request)).toEqual(harness.session.getActiveToolNames());
+		expect(getCurrentSystemPrompt(request.messages)).toContain(harness.session.systemPrompt);
+		expect(getCurrentSystemPrompt(request.messages)).toContain("ephemeral reminder");
+	});
+});
+
+describe("context_with_system handlers", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) harnesses.pop()?.cleanup();
+	});
+
+	it("runs after context handlers on the restored transcript and sends its output verbatim", async () => {
+		const seen: AgentMessage[][] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				compactViaHook,
+				(pi) => {
+					pi.on("context_with_system", async (event) => {
+						seen.push(event.messages);
+						return {
+							messages: event.messages.map((message) =>
+								message.role === "system" && message.toolsAdded
+									? { ...message, toolsAdded: message.toolsAdded.filter((tool) => tool.name !== "bash") }
+									: message,
+							),
+						};
+					});
+					// Registered after, but runs first: context handlers precede context_with_system.
+					pi.on("context", async (event) => {
+						const summary = event.messages.findIndex((message) => message.role === "compactionSummary");
+						return { messages: event.messages.slice(summary) };
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		await compactSession(harness);
+		const getRequest = captureRequest(harness, "after compaction");
+
+		await harness.session.prompt("third");
+
+		const input = seen.at(-1);
+		expect(input?.[0]?.role).toBe("system");
+		expect(input?.[1]?.role).toBe("compactionSummary");
+		expect(harness.session.getActiveToolNames()).toContain("bash");
+		expect(toolNames(getRequest())).toEqual(harness.session.getActiveToolNames().filter((name) => name !== "bash"));
+	});
+
+	it("reports a handler that drops the leading system message but honors its output", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("context_with_system", async (event) => ({
+						messages: event.messages.filter((message) => message.role !== "system"),
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		const errors: string[] = [];
+		harness.session.extensionRunner.onError((error) => {
+			errors.push(`${error.event}: ${error.error}`);
+		});
+		const getRequest = captureRequest(harness, "done");
+
+		await harness.session.prompt("hello");
+
+		expect(getRequest().messages.map((message) => message.role)).toEqual(["user"]);
+		expect(errors).toEqual([
+			expect.stringMatching(/^context_with_system: Handler removed the leading system message/),
+		]);
+	});
+});