feat(coding-agent): allow message_end replacements

closes #3982

packages/coding-agent/CHANGELOG.md

@@ -4,9 +4,14 @@
 
 ### Added
 
+- Added `message_end` extension result support for replacing finalized messages, enabling extensions to override assistant usage cost ([#3982](https://github.com/badlogic/pi-mono/issues/3982)).
 - Added top-level `name` support to `pi.registerProvider()` so extension-registered providers can show a friendly name in `/login` ([#3956](https://github.com/badlogic/pi-mono/issues/3956)).
 - Added `ctx.ui.getEditorComponent()` so extensions can wrap the currently configured custom editor factory ([#3935](https://github.com/badlogic/pi-mono/issues/3935)).
 
+### Removed
+
+- Removed Google Gemini CLI and Google Antigravity built-in login, default model, documentation, and example extension support.
+
 ### Fixed
 
 - Updated `@mariozechner/clipboard` to an attested release so package managers with trust policies do not reject installs ([#3946](https://github.com/badlogic/pi-mono/issues/3946)).

packages/coding-agent/docs/extensions.md

@@ -527,6 +527,7 @@ Fired for message lifecycle updates.
 
 - `message_start` and `message_end` fire for user, assistant, and toolResult messages.
 - `message_update` fires for assistant streaming updates.
+- `message_end` handlers can return `{ message }` to replace the finalized message. The replacement must keep the same `role`.
 
 ```typescript
 pi.on("message_start", async (event, ctx) => {
@@ -539,7 +540,20 @@ pi.on("message_update", async (event, ctx) => {
 });
 
 pi.on("message_end", async (event, ctx) => {
-  // event.message
+  if (event.message.role !== "assistant") return;
+
+  return {
+    message: {
+      ...event.message,
+      usage: {
+        ...event.message.usage,
+        cost: {
+          ...event.message.usage.cost,
+          total: 0.123,
+        },
+      },
+    },
+  };
 });
 ```
 
@@ -2557,7 +2571,6 @@ All examples in [examples/extensions/](../examples/extensions/).
 | `session-name.ts` | Name sessions for selector | `setSessionName`, `getSessionName` |
 | `bookmark.ts` | Bookmark entries for /tree | `setLabel` |
 | **Misc** |||
-| `antigravity-image-gen.ts` | Image generation tool | `registerTool`, Google Antigravity |
 | `inline-bash.ts` | Inline bash in tool calls | `on("tool_call")` |
 | `bash-spawn-hook.ts` | Adjust bash command, cwd, and env before execution | `createBashTool`, `spawnHook` |
 | `with-deps/` | Extension with npm dependencies | Package structure with `package.json` |

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

@@ -610,6 +610,22 @@ export class AgentSession {
 		return undefined;
 	}
 
+	private _replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void {
+		// Agent-core stores the finalized message object in its state before emitting message_end.
+		// SessionManager persistence happens later in _processAgentEvent() with event.message.
+		// Mutating this object in place keeps agent state, later turn/agent events, listeners,
+		// and the eventual SessionManager.appendMessage(event.message) persistence in sync.
+		if (target === replacement) {
+			return;
+		}
+
+		const targetRecord = target as unknown as Record<string, unknown>;
+		for (const key of Object.keys(targetRecord)) {
+			delete targetRecord[key];
+		}
+		Object.assign(targetRecord, replacement);
+	}
+
 	/** Emit extension events based on agent events */
 	private async _emitExtensionEvent(event: AgentEvent): Promise<void> {
 		if (event.type === "agent_start") {
@@ -651,7 +667,10 @@ export class AgentSession {
 				type: "message_end",
 				message: event.message,
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			const replacement = await this._extensionRunner.emitMessageEnd(extensionEvent);
+			if (replacement) {
+				this._replaceMessageInPlace(event.message, replacement);
+			}
 		} else if (event.type === "tool_execution_start") {
 			const extensionEvent: ToolExecutionStartEvent = {
 				type: "tool_execution_start",

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

@@ -34,6 +34,8 @@ import type {
 	InputEvent,
 	InputEventResult,
 	InputSource,
+	MessageEndEvent,
+	MessageEndEventResult,
 	MessageRenderer,
 	ProviderConfig,
 	RegisteredCommand,
@@ -118,6 +120,7 @@ type RunnerEmitEvent = Exclude<
 	| ContextEvent
 	| BeforeProviderRequestEvent
 	| BeforeAgentStartEvent
+	| MessageEndEvent
 	| ResourcesDiscoverEvent
 	| InputEvent
 >;
@@ -708,6 +711,48 @@ export class ExtensionRunner {
 		return result as RunnerEmitResult<TEvent>;
 	}
 
+	async emitMessageEnd(event: MessageEndEvent): Promise<AgentMessage | undefined> {
+		const ctx = this.createContext();
+		let currentMessage = event.message;
+		let modified = false;
+
+		for (const ext of this.extensions) {
+			const handlers = ext.handlers.get("message_end");
+			if (!handlers || handlers.length === 0) continue;
+
+			for (const handler of handlers) {
+				try {
+					const currentEvent: MessageEndEvent = { ...event, message: currentMessage };
+					const handlerResult = (await handler(currentEvent, ctx)) as MessageEndEventResult | undefined;
+					if (!handlerResult?.message) continue;
+
+					if (handlerResult.message.role !== currentMessage.role) {
+						this.emitError({
+							extensionPath: ext.path,
+							event: "message_end",
+							error: "message_end handlers must return a message with the same role",
+						});
+						continue;
+					}
+
+					currentMessage = handlerResult.message;
+					modified = true;
+				} 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: "message_end",
+						error: message,
+						stack,
+					});
+				}
+			}
+		}
+
+		return modified ? currentMessage : undefined;
+	}
+
 	async emitToolResult(event: ToolResultEvent): Promise<ToolResultEventResult | undefined> {
 		const ctx = this.createContext();
 		const currentEvent: ToolResultEvent = { ...event };

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

@@ -993,6 +993,11 @@ export interface ToolResultEventResult {
 	isError?: boolean;
 }
 
+export interface MessageEndEventResult {
+	/** Replace the finalized message. The replacement must keep the original message role. */
+	message?: AgentMessage;
+}
+
 export interface BeforeAgentStartEventResult {
 	message?: Pick<CustomMessage, "customType" | "content" | "display" | "details">;
 	/** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */
@@ -1101,7 +1106,7 @@ export interface ExtensionAPI {
 	on(event: "turn_end", handler: ExtensionHandler<TurnEndEvent>): void;
 	on(event: "message_start", handler: ExtensionHandler<MessageStartEvent>): void;
 	on(event: "message_update", handler: ExtensionHandler<MessageUpdateEvent>): void;
-	on(event: "message_end", handler: ExtensionHandler<MessageEndEvent>): void;
+	on(event: "message_end", handler: ExtensionHandler<MessageEndEvent, MessageEndEventResult>): void;
 	on(event: "tool_execution_start", handler: ExtensionHandler<ToolExecutionStartEvent>): void;
 	on(event: "tool_execution_update", handler: ExtensionHandler<ToolExecutionUpdateEvent>): void;
 	on(event: "tool_execution_end", handler: ExtensionHandler<ToolExecutionEndEvent>): void;

packages/coding-agent/test/suite/agent-session-runtime.test.ts

@@ -120,6 +120,47 @@ describe("AgentSessionRuntime characterization", () => {
 		return { runtime, faux, tempDir };
 	}
 
+	it("persists message_end assistant replacements to the session manager", async () => {
+		const { runtime } = await createRuntimeForTest((pi: ExtensionAPI) => {
+			pi.on("message_end", (event) => {
+				if (event.message.role !== "assistant") return;
+
+				return {
+					message: {
+						...event.message,
+						usage: {
+							...event.message.usage,
+							cost: {
+								...event.message.usage.cost,
+								total: 0.123,
+							},
+						},
+					},
+				};
+			});
+		});
+
+		await runtime.session.prompt("hello");
+
+		const sessionAssistant = runtime.session.messages.find((message) => message.role === "assistant");
+		expect(sessionAssistant?.role).toBe("assistant");
+		if (sessionAssistant?.role !== "assistant") {
+			throw new Error("missing assistant message");
+		}
+		expect(sessionAssistant.usage.cost.total).toBe(0.123);
+
+		const persistedAssistant = runtime.session.sessionManager
+			.getEntries()
+			.filter((entry) => entry.type === "message")
+			.map((entry) => entry.message)
+			.find((message) => message.role === "assistant");
+		expect(persistedAssistant?.role).toBe("assistant");
+		if (persistedAssistant?.role !== "assistant") {
+			throw new Error("missing persisted assistant message");
+		}
+		expect(persistedAssistant.usage.cost.total).toBe(0.123);
+	});
+
 	it("emits session_before_switch and session_start for new and resume flows", async () => {
 		const events: RecordedSessionEvent[] = [];
 		const { runtime } = await createRuntimeForTest((pi: ExtensionAPI) => {

packages/coding-agent/test/suite/regressions/3982-message-end-cost-override.test.ts

@@ -0,0 +1,56 @@
+import { fauxAssistantMessage } from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import { createHarness, type Harness } from "../harness.js";
+
+describe("regression #3982: message_end cost override", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("allows extensions to replace finalized assistant usage cost", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("message_end", (event) => {
+						if (event.message.role !== "assistant") return;
+
+						return {
+							message: {
+								...event.message,
+								usage: {
+									...event.message.usage,
+									cost: {
+										...event.message.usage.cost,
+										total: 0.123,
+									},
+								},
+							},
+						};
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("hello")]);
+
+		await harness.session.prompt("hi");
+
+		const assistantMessage = harness.session.messages.find((message) => message.role === "assistant");
+		expect(assistantMessage?.role).toBe("assistant");
+		if (assistantMessage?.role !== "assistant") {
+			throw new Error("missing assistant message");
+		}
+		expect(assistantMessage.usage.cost.total).toBe(0.123);
+
+		const messageEnd = harness.eventsOfType("message_end").find((event) => event.message.role === "assistant");
+		expect(messageEnd?.message.role).toBe("assistant");
+		if (messageEnd?.message.role !== "assistant") {
+			throw new Error("missing assistant message_end event");
+		}
+		expect(messageEnd.message.usage.cost.total).toBe(0.123);
+	});
+});