feat(coding-agent): allow message_end replacements

closes #3982

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);
+	});
+});