fix(coding-agent): preserve usage in streaming events (#7982)

Fixes #7911

packages/coding-agent/docs/json.md

@@ -19,6 +19,7 @@ type JsonAgentSessionEvent =
   | Exclude<AgentSessionEvent, { type: "message_update" }>
   | {
       type: "message_update";
+      usage: Usage;
       assistantMessageEvent: WithoutPartial<AssistantMessageEvent>;
     };
 ```
@@ -73,16 +74,17 @@ Followed by events as they occur:
 {"type":"agent_start"}
 {"type":"turn_start"}
 {"type":"message_start","message":{"role":"assistant","content":[],...}}
-{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}}
+{"type":"message_update","usage":{...},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}}
 {"type":"message_end","message":{...}}
 {"type":"turn_end","message":{...},"toolResults":[]}
 {"type":"agent_end","messages":[...]}
 ```
 
 `message_update` records are delta-only. They omit both the cumulative `message` field and
-`assistantMessageEvent.partial` to keep stream size linear. Use `contentIndex` and `delta`
-to assemble live text, thinking, or tool-call arguments if needed. `message_end` contains
-the final authoritative message.
+`assistantMessageEvent.partial` to keep stream size linear. The top-level `usage` field contains
+the latest cumulative provider-reported usage and may remain zero when a provider only reports
+usage at completion. Use `contentIndex` and `delta` to assemble live text, thinking, or tool-call
+arguments if needed. `message_end` contains the final authoritative message.
 
 ## Example
 

packages/coding-agent/docs/rpc.md

@@ -919,6 +919,14 @@ Emitted during streaming of assistant messages. Contains a delta event without a
 ```json
 {
   "type": "message_update",
+  "usage": {
+    "input": 100,
+    "output": 1,
+    "cacheRead": 0,
+    "cacheWrite": 0,
+    "totalTokens": 101,
+    "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}
+  },
   "assistantMessageEvent": {
     "type": "text_delta",
     "contentIndex": 0,
@@ -943,12 +951,15 @@ The `assistantMessageEvent` field contains one of these delta types:
 
 Example streaming a text response:
 ```json
-{"type":"message_update","assistantMessageEvent":{"type":"text_start","contentIndex":0}}
-{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}}
-{"type":"message_update","assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":" world"}}
-{"type":"message_update","assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"Hello world"}}
+{"type":"message_update","usage":{...},"assistantMessageEvent":{"type":"text_start","contentIndex":0}}
+{"type":"message_update","usage":{...},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"Hello"}}
+{"type":"message_update","usage":{...},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":" world"}}
+{"type":"message_update","usage":{...},"assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"Hello world"}}
 ```
 
+The top-level `usage` field contains the latest cumulative provider-reported usage. It may remain
+zero until completion when a provider does not report usage during streaming.
+
 `message_update` intentionally omits the former cumulative `message` field and
 `assistantMessageEvent.partial`. Clients that need a live partial message must assemble it
 from `message_start` and subsequent events using `contentIndex`. Treat `message_end.message`

packages/coding-agent/src/modes/json-event.ts

@@ -1,3 +1,4 @@
+import type { Usage } from "@earendil-works/pi-ai";
 import type { AgentSessionEvent } from "../core/agent-session.ts";
 
 type WithoutPartial<T> = T extends { partial: unknown } ? Omit<T, "partial"> : T;
@@ -8,6 +9,7 @@ type ToJsonEvent<T> = T extends {
 }
 	? {
 			type: "message_update";
+			usage: Usage;
 			assistantMessageEvent: WithoutPartial<TAssistantMessageEvent>;
 		}
 	: T;
@@ -21,7 +23,8 @@ type JsonMessageUpdateEvent = Extract<JsonAgentSessionEvent, { type: "message_up
 /**
  * Remove cumulative assistant snapshots from streaming wire events.
  * `message_start` provides the initial message, deltas build it, and
- * `message_end` provides the final authoritative message.
+ * `message_end` provides the final authoritative message. Cumulative usage
+ * remains available because its size is constant.
  */
 export function toJsonEvent(event: MessageUpdateEvent): JsonMessageUpdateEvent;
 export function toJsonEvent(event: AgentSessionEvent): JsonAgentSessionEvent;
@@ -29,12 +32,15 @@ export function toJsonEvent(event: AgentSessionEvent): JsonAgentSessionEvent {
 	if (event.type !== "message_update") {
 		return event;
 	}
+	if (event.message.role !== "assistant") {
+		throw new Error("message_update message is not an assistant message");
+	}
 
 	const assistantMessageEvent = event.assistantMessageEvent;
 	if (!("partial" in assistantMessageEvent)) {
-		return { type: "message_update", assistantMessageEvent };
+		return { type: "message_update", usage: event.message.usage, assistantMessageEvent };
 	}
 
 	const { partial: _partial, ...deltaEvent } = assistantMessageEvent;
-	return { type: "message_update", assistantMessageEvent: deltaEvent };
+	return { type: "message_update", usage: event.message.usage, assistantMessageEvent: deltaEvent };
 }

packages/coding-agent/test/suite/regressions/7911-json-stream-usage.test.ts

@@ -0,0 +1,32 @@
+import { fauxAssistantMessage } from "@earendil-works/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import { toJsonEvent } from "../../../src/modes/json-event.ts";
+import { createHarness, type Harness } from "../harness.ts";
+
+describe("regression #7911: JSON message updates retain usage", () => {
+	let harness: Harness | undefined;
+
+	afterEach(() => {
+		harness?.cleanup();
+	});
+
+	it("includes cumulative usage without cumulative message snapshots", async () => {
+		harness = await createHarness();
+		harness.setResponses([fauxAssistantMessage("hello")]);
+
+		await harness.session.prompt("respond");
+
+		// #7290's delta-only wire projection dropped this fixed-size metadata with the snapshots.
+		const update = harness
+			.eventsOfType("message_update")
+			.find((event) => event.message.role === "assistant" && event.message.usage.totalTokens > 0);
+		if (!update || update.message.role !== "assistant") {
+			throw new Error("Expected an assistant update with populated usage");
+		}
+
+		const wireUpdate = toJsonEvent(update);
+		expect(wireUpdate.usage).toEqual(update.message.usage);
+		expect(wireUpdate).not.toHaveProperty("message");
+		expect(wireUpdate.assistantMessageEvent).not.toHaveProperty("partial");
+	});
+});