diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts index 5246c43a741..9aa445a1ab6 100644 --- a/src/agents/pi-embedded-subscribe.handlers.messages.ts +++ b/src/agents/pi-embedded-subscribe.handlers.messages.ts @@ -165,10 +165,9 @@ export function handleMessageUpdate( } } - const streamedReasoning = extractThinkingFromTaggedStream(ctx.state.deltaBuffer); if (ctx.state.streamReasoning) { // Handle partial tags: stream whatever reasoning is visible so far. - ctx.emitReasoningStream(streamedReasoning); + ctx.emitReasoningStream(extractThinkingFromTaggedStream(ctx.state.deltaBuffer)); } const next = ctx @@ -277,9 +276,7 @@ export function handleMessageEnd( }); const rawThinking = ctx.state.includeReasoning || ctx.state.streamReasoning - ? extractAssistantThinking(assistantMessage) || - extractThinkingFromTaggedStream(ctx.state.deltaBuffer) || - extractThinkingFromTaggedText(rawText) + ? extractAssistantThinking(assistantMessage) || extractThinkingFromTaggedText(rawText) : ""; const formattedReasoning = rawThinking ? formatReasoningMessage(rawThinking) : ""; const trimmedText = text.trim(); diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.emits-reasoning-as-separate-message-enabled.e2e.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.emits-reasoning-as-separate-message-enabled.e2e.test.ts index 767c344aab8..98b4ce09237 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.emits-reasoning-as-separate-message-enabled.e2e.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.emits-reasoning-as-separate-message-enabled.e2e.test.ts @@ -68,36 +68,4 @@ describe("subscribeEmbeddedPiSession", () => { ]); }, ); - - it.each(THINKING_TAG_CASES)( - "falls back to streamed <%s> reasoning when message_end has no thinking block", - ({ open, close }) => { - const { emit, onBlockReply } = createReasoningBlockReplyHarness(); - - emit({ type: "message_start", message: { role: "assistant" } }); - emit({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { type: "text_delta", delta: `${open}Because` }, - }); - emit({ - type: "message_update", - message: { role: "assistant" }, - assistantMessageEvent: { - type: "text_delta", - delta: ` it helps${close}\n\nFinal answer`, - }, - }); - - emit({ - type: "message_end", - message: { - role: "assistant", - content: [{ type: "text", text: "Final answer" }], - } as AssistantMessage, - }); - - expectReasoningAndAnswerCalls(onBlockReply); - }, - ); });