From 40c6eabb8f9b34ca0ffb652d6e3a7929b5e2eee5 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 30 Apr 2026 21:19:09 +0200 Subject: [PATCH] feat(coding-agent): allow message_end replacements closes #3982 --- packages/coding-agent/CHANGELOG.md | 5 ++ packages/coding-agent/docs/extensions.md | 17 +++++- .../coding-agent/src/core/agent-session.ts | 21 ++++++- .../src/core/extensions/runner.ts | 45 +++++++++++++++ .../coding-agent/src/core/extensions/types.ts | 7 ++- .../test/suite/agent-session-runtime.test.ts | 41 ++++++++++++++ .../3982-message-end-cost-override.test.ts | 56 +++++++++++++++++++ 7 files changed, 188 insertions(+), 4 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/3982-message-end-cost-override.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d5f1ebb5..e29be22d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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)). diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 5864a01c..e20924c4 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/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` | diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 7a26deb4..bf87179f 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/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; + 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 { 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", diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 03c98077..9b7f2ddf 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/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; } + async emitMessageEnd(event: MessageEndEvent): Promise { + 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 { const ctx = this.createContext(); const currentEvent: ToolResultEvent = { ...event }; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index a91d1f2f..0df7649f 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/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; /** 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): void; on(event: "message_start", handler: ExtensionHandler): void; on(event: "message_update", handler: ExtensionHandler): void; - on(event: "message_end", handler: ExtensionHandler): void; + on(event: "message_end", handler: ExtensionHandler): void; on(event: "tool_execution_start", handler: ExtensionHandler): void; on(event: "tool_execution_update", handler: ExtensionHandler): void; on(event: "tool_execution_end", handler: ExtensionHandler): void; diff --git a/packages/coding-agent/test/suite/agent-session-runtime.test.ts b/packages/coding-agent/test/suite/agent-session-runtime.test.ts index 220abfef..2b99d3ad 100644 --- a/packages/coding-agent/test/suite/agent-session-runtime.test.ts +++ b/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) => { diff --git a/packages/coding-agent/test/suite/regressions/3982-message-end-cost-override.test.ts b/packages/coding-agent/test/suite/regressions/3982-message-end-cost-override.test.ts new file mode 100644 index 00000000..e2f4bd4f --- /dev/null +++ b/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); + }); +});