From 98d3281cad05f2836f6d5f9365faef0409c105b0 Mon Sep 17 00:00:00 2001
From: dundunHa <lixinpengde@gmail.com>
Date: Sat, 8 Aug 2026 03:26:08 +0800
Subject: [PATCH] fix(pi): stream assistant deltas when pi omits the cumulative
 message (#2978)

* fix(pi): stream assistant deltas when pi omits the cumulative message

pi 0.84.0 removed the cumulative `message` snapshot from JSON/RPC
`message_update` events, so `handleMessageUpdate` dereferenced
`undefined` and took down the whole daemon on the first assistant delta.
Every prompt to a pi agent crash-looped the worker.

`message_update` carries an `assistantMessageEvent`, which pi only emits
for assistant messages, so a missing snapshot is treated as an assistant
update. `message_start` still carries the message, so the response id
keeps flowing through `activeAssistantMessageId`.

Fixes #2940
Fixes #2946
Fixes #2950

* fix(pi): give the message_update COMPAT tag a version and removal date

Follows docs/protocol-compatibility.md: every shim tag needs a name, a
version, and a removal condition or date so `rg "COMPAT\("` stays a
schedulable cleanup backlog.

---------

Co-authored-by: dundunHa <dundunhade@gmail.com>
---
 .../server/agent/providers/pi/agent.test.ts   | 57 +++++++++++++++++++
 .../src/server/agent/providers/pi/agent.ts    |  4 +-
 .../server/agent/providers/pi/rpc-types.ts    |  4 +-
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts
index d6b0f9377..c367b57ab 100644
--- a/packages/server/src/server/agent/providers/pi/agent.test.ts
+++ b/packages/server/src/server/agent/providers/pi/agent.test.ts
@@ -662,6 +662,63 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("streams assistant text and reasoning when Pi omits the cumulative message", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    await session.startTurn("hello");
+    fakeSession.emit({
+      type: "message_start",
+      message: { role: "assistant", content: [], responseId: "response-1" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "hel" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "lo" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_delta", delta: "thinking" },
+    });
+
+    expect(events.timelineItems()).toEqual([
+      { type: "assistant_message", text: "hel", messageId: "response-1" },
+      { type: "assistant_message", text: "lo", messageId: "response-1" },
+      { type: "reasoning", text: "thinking" },
+    ]);
+  });
+
+  test("generates one message id when Pi omits both message start and the cumulative message", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    await session.startTurn("hello");
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "hel" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "lo" },
+    });
+
+    const [firstChunk, secondChunk] = events.timelineItems();
+    expect(firstChunk).toMatchObject({
+      type: "assistant_message",
+      text: "hel",
+      messageId: expect.any(String),
+    });
+    const firstMessageId = (firstChunk as { messageId: string }).messageId;
+    expect(secondChunk).toEqual({
+      type: "assistant_message",
+      text: "lo",
+      messageId: firstMessageId,
+    });
+  });
+
   test("emits live user messages with submitted Pi tree entry ids", async () => {
     const { pi, session, events } = await createSession();
     const fakeSession = pi.latestSession();
diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts
index f08c68e28..04adc29b9 100644
--- a/packages/server/src/server/agent/providers/pi/agent.ts
+++ b/packages/server/src/server/agent/providers/pi/agent.ts
@@ -2191,12 +2191,12 @@ export class PiRpcAgentSession implements AgentSession {
     event: Extract<PiAgentSessionEvent, { type: "message_update" }>,
     turnId: string | undefined,
   ): void {
-    if (event.message.role !== "assistant") {
+    if (event.message && event.message.role !== "assistant") {
       return;
     }
     if (event.assistantMessageEvent.type === "text_delta") {
       // Pi-compatible runtimes may emit updates without a preceding message_start.
-      this.activeAssistantMessageId ??= event.message.responseId || randomUUID();
+      this.activeAssistantMessageId ??= event.message?.responseId || randomUUID();
       this.emit({
         type: "timeline",
         provider: this.provider,
diff --git a/packages/server/src/server/agent/providers/pi/rpc-types.ts b/packages/server/src/server/agent/providers/pi/rpc-types.ts
index 949ff9303..a84b42876 100644
--- a/packages/server/src/server/agent/providers/pi/rpc-types.ts
+++ b/packages/server/src/server/agent/providers/pi/rpc-types.ts
@@ -160,7 +160,9 @@ export type PiAgentSessionEvent =
   | { type: "message_end"; message: PiAgentMessage }
   | {
       type: "message_update";
-      message: PiAgentMessage;
+      // COMPAT(piCumulativeMessageUpdate): added in v0.3.0, remove after 2027-02-07 once the pi
+      // floor is >=0.84. pi <=0.83 repeats the cumulative assistant message on every update.
+      message?: PiAgentMessage;
       assistantMessageEvent: PiAssistantMessageEvent;
     }
   | {