From 477975713839452a525502c92e5952cc5dafe531 Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Sat, 20 Jun 2026 22:47:49 +0700
Subject: [PATCH] Forward live usage stream events

The wire schema already allowed completed-turn usage, but rejected active usage updates before clients could see them. Accept usage_updated payloads and preserve turn ids on streamed turn events.
---
 .../src/messages.stream-parsing.test.ts       | 39 +++++++++++++++++++
 packages/protocol/src/messages.ts             | 13 +++++++
 2 files changed, 52 insertions(+)

diff --git a/packages/protocol/src/messages.stream-parsing.test.ts b/packages/protocol/src/messages.stream-parsing.test.ts
index d4ea77d20..61f10ab4c 100644
--- a/packages/protocol/src/messages.stream-parsing.test.ts
+++ b/packages/protocol/src/messages.stream-parsing.test.ts
@@ -160,6 +160,45 @@ describe("shared messages stream parsing", () => {
     }
   });
 
+  it("parses agent_stream usage_updated events", () => {
+    const fixture = {
+      type: "agent_stream",
+      payload: {
+        agentId: "agent_live",
+        timestamp: "2026-02-08T20:10:00.000Z",
+        event: {
+          type: "usage_updated",
+          provider: "claude",
+          turnId: "foreground-turn-1",
+          usage: {
+            contextWindowUsedTokens: 24_074,
+            contextWindowMaxTokens: 200_000,
+          },
+        },
+      },
+    };
+
+    const parsed = AgentStreamMessageSchema.parse(fixture);
+    expect(parsed.payload.event).toEqual({
+      type: "usage_updated",
+      provider: "claude",
+      turnId: "foreground-turn-1",
+      usage: {
+        contextWindowUsedTokens: 24_074,
+        contextWindowMaxTokens: 200_000,
+      },
+    });
+
+    const wrapped = WSOutboundMessageSchema.parse({
+      type: "session",
+      message: fixture,
+    });
+    expect(wrapped.message.type).toBe("agent_stream");
+    if (wrapped.message.type === "agent_stream") {
+      expect(wrapped.message.payload.event.type).toBe("usage_updated");
+    }
+  });
+
   it("parses representative sub_agent tool_call event", () => {
     const parsed = AgentStreamMessageSchema.parse({
       type: "agent_stream",
diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts
index 2757f4f58..772ad18a2 100644
--- a/packages/protocol/src/messages.ts
+++ b/packages/protocol/src/messages.ts
@@ -594,11 +594,19 @@ export const AgentStreamEventPayloadSchema = z.discriminatedUnion("type", [
   z.object({
     type: z.literal("turn_started"),
     provider: AgentProviderSchema,
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("turn_completed"),
     provider: AgentProviderSchema,
     usage: AgentUsageSchema.optional(),
+    turnId: z.string().optional(),
+  }),
+  z.object({
+    type: z.literal("usage_updated"),
+    provider: AgentProviderSchema,
+    usage: AgentUsageSchema,
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("turn_failed"),
@@ -606,27 +614,32 @@ export const AgentStreamEventPayloadSchema = z.discriminatedUnion("type", [
     error: z.string(),
     code: z.string().optional(),
     diagnostic: z.string().optional(),
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("turn_canceled"),
     provider: AgentProviderSchema,
     reason: z.string(),
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("timeline"),
     provider: AgentProviderSchema,
     item: AgentTimelineItemPayloadSchema,
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("permission_requested"),
     provider: AgentProviderSchema,
     request: AgentPermissionRequestPayloadSchema,
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("permission_resolved"),
     provider: AgentProviderSchema,
     requestId: z.string(),
     resolution: AgentPermissionResponseSchema,
+    turnId: z.string().optional(),
   }),
   z.object({
     type: z.literal("attention_required"),