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" }