diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.does-not-append-text-end-content-is.e2e.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.does-not-append-text-end-content-is.e2e.test.ts index 964ff5b3a..690a1d7ab 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.does-not-append-text-end-content-is.e2e.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.does-not-append-text-end-content-is.e2e.test.ts @@ -13,7 +13,7 @@ describe("subscribeEmbeddedPiSession", () => { { tag: "antthinking", open: "", close: "" }, ] as const; - it("does not append when text_end content is a prefix of deltas", () => { + function setupTextEndSubscription() { let handler: ((evt: unknown) => void) | undefined; const session: StubSession = { subscribe: (fn) => { @@ -31,103 +31,59 @@ describe("subscribeEmbeddedPiSession", () => { blockReplyBreak: "text_end", }); - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_delta", - delta: "Hello world", - }, - }); + const emit = (evt: unknown) => handler?.(evt); - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_end", - content: "Hello", - }, - }); - - expect(onBlockReply).toHaveBeenCalledTimes(1); - expect(subscription.assistantTexts).toEqual(["Hello world"]); - }); - it("does not append when text_end content is already contained", () => { - let handler: ((evt: unknown) => void) | undefined; - const session: StubSession = { - subscribe: (fn) => { - handler = fn; - return () => {}; - }, + const emitDelta = (delta: string) => { + emit({ + type: "message_update", + message: { role: "assistant" }, + assistantMessageEvent: { + type: "text_delta", + delta, + }, + }); }; - const onBlockReply = vi.fn(); - - const subscription = subscribeEmbeddedPiSession({ - session: session as unknown as Parameters[0]["session"], - runId: "run", - onBlockReply, - blockReplyBreak: "text_end", - }); - - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_delta", - delta: "Hello world", - }, - }); - - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_end", - content: "world", - }, - }); - - expect(onBlockReply).toHaveBeenCalledTimes(1); - expect(subscription.assistantTexts).toEqual(["Hello world"]); - }); - it("appends suffix when text_end content extends deltas", () => { - let handler: ((evt: unknown) => void) | undefined; - const session: StubSession = { - subscribe: (fn) => { - handler = fn; - return () => {}; - }, + const emitTextEnd = (content: string) => { + emit({ + type: "message_update", + message: { role: "assistant" }, + assistantMessageEvent: { + type: "text_end", + content, + }, + }); }; - const onBlockReply = vi.fn(); + return { onBlockReply, subscription, emitDelta, emitTextEnd }; + } - const subscription = subscribeEmbeddedPiSession({ - session: session as unknown as Parameters[0]["session"], - runId: "run", - onBlockReply, - blockReplyBreak: "text_end", - }); + it.each([ + { + name: "does not append when text_end content is a prefix of deltas", + delta: "Hello world", + content: "Hello", + expected: "Hello world", + }, + { + name: "does not append when text_end content is already contained", + delta: "Hello world", + content: "world", + expected: "Hello world", + }, + { + name: "appends suffix when text_end content extends deltas", + delta: "Hello", + content: "Hello world", + expected: "Hello world", + }, + ])("$name", ({ delta, content, expected }) => { + const { onBlockReply, subscription, emitDelta, emitTextEnd } = setupTextEndSubscription(); - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_delta", - delta: "Hello", - }, - }); - - handler?.({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_end", - content: "Hello world", - }, - }); + emitDelta(delta); + emitTextEnd(content); expect(onBlockReply).toHaveBeenCalledTimes(1); - expect(subscription.assistantTexts).toEqual(["Hello world"]); + expect(subscription.assistantTexts).toEqual([expected]); }); });