diff --git a/src/agents/pi-embedded-runner/thinking.test.ts b/src/agents/pi-embedded-runner/thinking.test.ts index 6a2481748a1..e3d0a8291b6 100644 --- a/src/agents/pi-embedded-runner/thinking.test.ts +++ b/src/agents/pi-embedded-runner/thinking.test.ts @@ -3,6 +3,22 @@ import { describe, expect, it } from "vitest"; import { castAgentMessage } from "../test-helpers/agent-message-fixtures.js"; import { dropThinkingBlocks, isAssistantMessageWithContent } from "./thinking.js"; +function dropSingleAssistantContent(content: Array>) { + const messages: AgentMessage[] = [ + castAgentMessage({ + role: "assistant", + content, + }), + ]; + + const result = dropThinkingBlocks(messages); + return { + assistant: result[0] as Extract, + messages, + result, + }; +} + describe("isAssistantMessageWithContent", () => { it("accepts assistant messages with array content and rejects others", () => { const assistant = castAgentMessage({ @@ -30,32 +46,18 @@ describe("dropThinkingBlocks", () => { }); it("drops thinking blocks while preserving non-thinking assistant content", () => { - const messages: AgentMessage[] = [ - castAgentMessage({ - role: "assistant", - content: [ - { type: "thinking", thinking: "internal" }, - { type: "text", text: "final" }, - ], - }), - ]; - - const result = dropThinkingBlocks(messages); - const assistant = result[0] as Extract; + const { assistant, messages, result } = dropSingleAssistantContent([ + { type: "thinking", thinking: "internal" }, + { type: "text", text: "final" }, + ]); expect(result).not.toBe(messages); expect(assistant.content).toEqual([{ type: "text", text: "final" }]); }); it("keeps assistant turn structure when all content blocks were thinking", () => { - const messages: AgentMessage[] = [ - castAgentMessage({ - role: "assistant", - content: [{ type: "thinking", thinking: "internal-only" }], - }), - ]; - - const result = dropThinkingBlocks(messages); - const assistant = result[0] as Extract; + const { assistant } = dropSingleAssistantContent([ + { type: "thinking", thinking: "internal-only" }, + ]); expect(assistant.content).toEqual([{ type: "text", text: "" }]); }); });