diff --git a/src/agents/pi-embedded-subscribe.e2e-harness.ts b/src/agents/pi-embedded-subscribe.e2e-harness.ts index 6cba5899b2..918685a184 100644 --- a/src/agents/pi-embedded-subscribe.e2e-harness.ts +++ b/src/agents/pi-embedded-subscribe.e2e-harness.ts @@ -6,6 +6,7 @@ type SubscribeEmbeddedPiSession = typeof subscribeEmbeddedPiSession; type SubscribeEmbeddedPiSessionParams = Parameters[0]; type PiSession = Parameters[0]["session"]; type OnBlockReply = NonNullable; +type BlockReplyChunking = NonNullable; export const THINKING_TAG_CASES = [ { tag: "think", open: "", close: "" }, @@ -70,6 +71,25 @@ export function createParagraphChunkedBlockReplyHarness(params: { return { emit, onBlockReply, subscription }; } +export function createTextEndBlockReplyHarness(params?: { + onBlockReply?: OnBlockReply; + runId?: string; + blockReplyChunking?: BlockReplyChunking; +}): { + emit: (evt: unknown) => void; + onBlockReply: OnBlockReply; + subscription: ReturnType; +} { + const onBlockReply: OnBlockReply = params?.onBlockReply ?? (() => {}); + const { emit, subscription } = createSubscribedSessionHarness({ + runId: params?.runId ?? "run", + onBlockReply, + blockReplyBreak: "text_end", + blockReplyChunking: params?.blockReplyChunking, + }); + return { emit, onBlockReply, subscription }; +} + export function extractAgentEventPayloads(calls: Array): Array> { return calls .map((call) => { @@ -120,6 +140,31 @@ export function emitAssistantTextDeltaAndEnd(params: { params.emit({ type: "message_end", message: assistantMessage }); } +export function emitAssistantTextDelta(params: { + emit: (evt: unknown) => void; + delta: string; +}): void { + params.emit({ + type: "message_update", + message: { role: "assistant" }, + assistantMessageEvent: { type: "text_delta", delta: params.delta }, + }); +} + +export function emitAssistantTextEnd(params: { + emit: (evt: unknown) => void; + content?: string; +}): void { + params.emit({ + type: "message_update", + message: { role: "assistant" }, + assistantMessageEvent: + typeof params.content === "string" + ? { type: "text_end", content: params.content } + : { type: "text_end" }, + }); +} + export function expectFencedChunks(calls: Array, expectedPrefix: string): void { expect(calls.length).toBeGreaterThan(1); for (const call of calls) {