fix(pi): preserve content block order while streaming

Pi identifies streaming blocks by contentIndex, but the provider discarded that identity and emitted deltas in arrival order. Buffer later blocks until earlier indexed blocks end so late reasoning stays ahead of assistant text.

packages/server/src/server/agent/providers/pi/agent.test.ts

@@ -604,6 +604,98 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("streams Pi content blocks by content index when an earlier thinking delta arrives late", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    await session.startTurn("hello");
+    fakeSession.emit({
+      type: "message_start",
+      message: {
+        role: "assistant",
+        responseId: "response-1",
+        content: [
+          { type: "thinking", thinking: "" },
+          { type: "text", text: "" },
+        ],
+      },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_start", contentIndex: 0 },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_delta", contentIndex: 0, delta: "verify" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_start", contentIndex: 1 },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", contentIndex: 1, delta: "Answer" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_delta", contentIndex: 0, delta: "." },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_end", contentIndex: 0 },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", contentIndex: 1, delta: " complete." },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_end", contentIndex: 1 },
+    });
+
+    expect(events.timelineItems()).toEqual([
+      { type: "reasoning", text: "verify" },
+      { type: "reasoning", text: "." },
+      { type: "assistant_message", text: "Answer", messageId: "response-1" },
+      { type: "assistant_message", text: " complete.", messageId: "response-1" },
+    ]);
+  });
+
+  test("flushes indexed Pi blocks under the previous message id when its end event is missing", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    await session.startTurn("hello");
+    fakeSession.emit({
+      type: "message_start",
+      message: {
+        role: "assistant",
+        responseId: "response-1",
+        content: [
+          { type: "thinking", thinking: "" },
+          { type: "text", text: "" },
+        ],
+      },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_delta", contentIndex: 0, delta: "think" },
+    });
+    fakeSession.emit({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", contentIndex: 1, delta: "answer" },
+    });
+    fakeSession.emit({
+      type: "message_start",
+      message: { role: "assistant", responseId: "response-2", content: [] },
+    });
+
+    expect(events.timelineItems()).toEqual([
+      { type: "reasoning", text: "think" },
+      { type: "assistant_message", text: "answer", messageId: "response-1" },
+    ]);
+  });
+
   test("streams Pi task calls as sub-agent cards with lifecycle status", async () => {
     const { pi, session, events } = await createSession();
     const fakeSession = pi.latestSession();

packages/server/src/server/agent/providers/pi/agent.ts

@@ -63,6 +63,7 @@ import {
 } from "./history-mapper.js";
 import { materializeProviderImage } from "../provider-image-output.js";
 import { PiCliRuntime } from "./cli-runtime.js";
+import { PiMessageContentStream, type PiMessageContentChunk } from "./message-content-stream.js";
 import { revertPiConversation } from "./rewind.js";
 import { listPiImportableSessions, readPiImportSessionConfig } from "./session-descriptor.js";
 import type { PiRuntime, PiRuntimeSession, PiStartSessionInput } from "./runtime.js";
@@ -1205,6 +1206,7 @@ export class PiRpcAgentSession implements AgentSession {
   private activeTurnId: string | null = null;
   private activeClientMessageId: string | null = null;
   private activeAssistantMessageId: string | null = null;
+  private readonly activeAssistantContent = new PiMessageContentStream();
   private activeTurnStarted = false;
   private activeTurnStartedEmitted = false;
   private pendingSettledMessages: PiAgentMessage[] | null = null;
@@ -1300,6 +1302,7 @@ export class PiRpcAgentSession implements AgentSession {
     this.lastInterruptedTurnId = null;
     this.activeClientMessageId = options?.clientMessageId ?? null;
     this.activeAssistantMessageId = null;
+    this.activeAssistantContent.reset();
     this.activeTurnStarted = false;
     this.activeTurnStartedEmitted = false;
     this.pendingSettledMessages = null;
@@ -2101,7 +2104,7 @@ export class PiRpcAgentSession implements AgentSession {
         });
         return;
       case "message_start":
-        this.handleMessageStart(event);
+        this.handleMessageStart(event, turnId);
         return;
       case "message_end":
         this.handleMessageEnd(event, turnId);
@@ -2246,37 +2249,55 @@ export class PiRpcAgentSession implements AgentSession {
     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.emit({
-        type: "timeline",
-        provider: this.provider,
-        turnId,
-        item: {
-          type: "assistant_message",
-          text: event.assistantMessageEvent.delta ?? "",
-          messageId: this.activeAssistantMessageId,
-        },
-      });
-      return;
-    }
-    if (event.assistantMessageEvent.type === "thinking_delta") {
+    const message = event.message?.role === "assistant" ? event.message : undefined;
+    this.emitMessageContentChunks(
+      this.activeAssistantContent.update(event.assistantMessageEvent, message),
+      turnId,
+      message?.responseId,
+    );
+  }
+
+  private emitMessageContentChunks(
+    chunks: PiMessageContentChunk[],
+    turnId: string | undefined,
+    responseId?: string,
+  ): void {
+    for (const chunk of chunks) {
+      if (chunk.type === "text") {
+        // Pi-compatible runtimes may emit updates without a preceding message_start.
+        this.activeAssistantMessageId ??= responseId || randomUUID();
+        this.emit({
+          type: "timeline",
+          provider: this.provider,
+          turnId,
+          item: {
+            type: "assistant_message",
+            text: chunk.text,
+            messageId: this.activeAssistantMessageId,
+          },
+        });
+        continue;
+      }
       this.emit({
         type: "timeline",
         provider: this.provider,
         turnId,
         item: {
           type: "reasoning",
-          text: event.assistantMessageEvent.delta ?? "",
+          text: chunk.text,
         },
       });
     }
   }
 
-  private handleMessageStart(event: Extract<PiAgentSessionEvent, { type: "message_start" }>): void {
+  private handleMessageStart(
+    event: Extract<PiAgentSessionEvent, { type: "message_start" }>,
+    turnId: string | undefined,
+  ): void {
     if (event.message.role === "assistant") {
+      this.emitMessageContentChunks(this.activeAssistantContent.finish(), turnId);
       this.activeAssistantMessageId = event.message.responseId || null;
+      this.activeAssistantContent.start(event.message.content);
     }
   }
 
@@ -2285,6 +2306,7 @@ export class PiRpcAgentSession implements AgentSession {
     turnId: string | undefined,
   ): void {
     if (event.message.role === "assistant") {
+      this.emitMessageContentChunks(this.activeAssistantContent.finish(), turnId);
       this.activeAssistantMessageId = null;
       return;
     }
@@ -2355,6 +2377,7 @@ export class PiRpcAgentSession implements AgentSession {
       this.lastInterruptedTurnId = null;
       return;
     }
+    this.emitMessageContentChunks(this.activeAssistantContent.finish(), turnId);
     this.activeTurnId = null;
     this.activeClientMessageId = null;
     this.activeAssistantMessageId = null;

packages/server/src/server/agent/providers/pi/message-content-stream.ts

@@ -0,0 +1,123 @@
+import type { PiAssistantContent, PiAssistantMessageEvent, PiAgentMessage } from "./rpc-types.js";
+
+export interface PiMessageContentChunk {
+  type: "text" | "thinking";
+  text: string;
+}
+
+interface PendingContentBlock {
+  chunks: PiMessageContentChunk[];
+  complete: boolean;
+}
+
+// Pi may revisit an earlier contentIndex while a later block is already streaming. Paseo's
+// timeline is append-only, so later blocks wait here until every earlier block has ended.
+export class PiMessageContentStream {
+  private readonly blocks = new Map<number, PendingContentBlock>();
+  private readonly completedIndexes = new Set<number>();
+
+  start(content: PiAssistantContent[]): void {
+    this.reset();
+    this.seedBlocks(content);
+  }
+
+  update(
+    event: PiAssistantMessageEvent,
+    message?: Extract<PiAgentMessage, { role: "assistant" }>,
+  ): PiMessageContentChunk[] {
+    this.seedBlocks(message?.content ?? []);
+    const contentIndex = readContentIndex(event);
+
+    if (event.type === "text_delta" || event.type === "thinking_delta") {
+      const chunk = {
+        type: event.type === "text_delta" ? ("text" as const) : ("thinking" as const),
+        text: event.delta ?? "",
+      };
+      if (contentIndex === null) {
+        return [chunk];
+      }
+      this.ensureBlock(contentIndex).chunks.push(chunk);
+      return this.flushReadyBlocks();
+    }
+
+    if (contentIndex === null) {
+      return [];
+    }
+    if (isContentStart(event)) {
+      this.ensureBlock(contentIndex);
+      return this.flushReadyBlocks();
+    }
+    if (isContentEnd(event)) {
+      this.ensureBlock(contentIndex).complete = true;
+      return this.flushReadyBlocks();
+    }
+    return [];
+  }
+
+  finish(): PiMessageContentChunk[] {
+    for (const block of this.blocks.values()) {
+      block.complete = true;
+    }
+    const chunks = this.flushReadyBlocks();
+    this.reset();
+    return chunks;
+  }
+
+  reset(): void {
+    this.blocks.clear();
+    this.completedIndexes.clear();
+  }
+
+  private seedBlocks(content: PiAssistantContent[]): void {
+    for (let index = 0; index < content.length; index += 1) {
+      this.ensureBlock(index);
+    }
+  }
+
+  private ensureBlock(contentIndex: number): PendingContentBlock {
+    const existing = this.blocks.get(contentIndex);
+    if (existing) {
+      return existing;
+    }
+    const block = { chunks: [], complete: false };
+    if (!this.completedIndexes.has(contentIndex)) {
+      this.blocks.set(contentIndex, block);
+    }
+    return block;
+  }
+
+  private flushReadyBlocks(): PiMessageContentChunk[] {
+    const chunks: PiMessageContentChunk[] = [];
+    while (this.blocks.size > 0) {
+      const contentIndex = Math.min(...this.blocks.keys());
+      const block = this.blocks.get(contentIndex);
+      if (!block) break;
+      chunks.push(...block.chunks.splice(0));
+      if (!block.complete) break;
+      this.blocks.delete(contentIndex);
+      this.completedIndexes.add(contentIndex);
+    }
+    return chunks;
+  }
+}
+
+function readContentIndex(event: PiAssistantMessageEvent): number | null {
+  if (!("contentIndex" in event)) return null;
+  return Number.isInteger(event.contentIndex) && (event.contentIndex ?? -1) >= 0
+    ? (event.contentIndex ?? null)
+    : null;
+}
+
+function isContentStart(event: PiAssistantMessageEvent): boolean {
+  return (
+    event.type === "text_start" ||
+    event.type === "thinking_start" ||
+    event.type === "toolcall_start"
+  );
+}
+
+function isContentEnd(event: PiAssistantMessageEvent): boolean {
+  return (
+    event.type === "text_end" || event.type === "thinking_end" || event.type === "toolcall_end"
+  );
+}

packages/server/src/server/agent/providers/pi/rpc-types.ts

@@ -149,9 +149,11 @@ export interface PiRpcResponse {
 }
 
 export type PiAssistantMessageEvent =
-  | { type: "text_delta"; delta?: string }
-  | { type: "thinking_delta"; delta?: string }
-  | { type: "start" | "text_start" | "text_end" | "thinking_start" | "thinking_end" | "done" };
+  | { type: "text_delta"; contentIndex?: number; delta?: string }
+  | { type: "thinking_delta"; contentIndex?: number; delta?: string }
+  | { type: "text_start" | "text_end" | "thinking_start" | "thinking_end"; contentIndex?: number }
+  | { type: "toolcall_start" | "toolcall_delta" | "toolcall_end"; contentIndex?: number }
+  | { type: "start" | "done" };
 
 export type PiAgentSessionEvent =
   | { type: "agent_start" }