From c3fd36e2e9d2fd276922a91c88eb0a75ecd0acbe Mon Sep 17 00:00:00 2001 From: "paseo-bot[bot]" <266920839+paseo-bot[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:42:39 +0000 Subject: [PATCH] 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. --- .../server/agent/providers/pi/agent.test.ts | 92 +++++++++++++ .../src/server/agent/providers/pi/agent.ts | 61 ++++++--- .../providers/pi/message-content-stream.ts | 123 ++++++++++++++++++ .../server/agent/providers/pi/rpc-types.ts | 8 +- 4 files changed, 262 insertions(+), 22 deletions(-) create mode 100644 packages/server/src/server/agent/providers/pi/message-content-stream.ts 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 dafe769bef..cf5a7d10ca 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/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(); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 9251c682d0..7822681982 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/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): void { + private handleMessageStart( + event: Extract, + 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; diff --git a/packages/server/src/server/agent/providers/pi/message-content-stream.ts b/packages/server/src/server/agent/providers/pi/message-content-stream.ts new file mode 100644 index 0000000000..10e240f0d3 --- /dev/null +++ b/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(); + private readonly completedIndexes = new Set(); + + start(content: PiAssistantContent[]): void { + this.reset(); + this.seedBlocks(content); + } + + update( + event: PiAssistantMessageEvent, + message?: Extract, + ): 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" + ); +} 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 ce3241a19c..1ccb5b5960 100644 --- a/packages/server/src/server/agent/providers/pi/rpc-types.ts +++ b/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" }