diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts index c08ac557bde..d9cd44e7a7f 100644 --- a/src/agents/pi-embedded-subscribe.handlers.messages.ts +++ b/src/agents/pi-embedded-subscribe.handlers.messages.ts @@ -16,6 +16,7 @@ import { extractThinkingFromTaggedText, formatReasoningMessage, promoteThinkingTagsToBlocks, + stripTrailingPartialThinkingTagFragment, } from "./pi-embedded-utils.js"; const stripTrailingDirective = (text: string): string => { @@ -30,9 +31,6 @@ const stripTrailingDirective = (text: string): string => { return text.slice(0, openIndex); }; -const stripTrailingPartialThinkingTag = (text: string): string => - text.replace(/<\s*\/?\s*(?:think(?:ing)?|thought|antthinking)\s*$/i, "").trimEnd(); - function emitReasoningEnd(ctx: EmbeddedPiSubscribeContext) { if (!ctx.state.reasoningStreamOpen) { return; @@ -192,7 +190,7 @@ export function handleMessageUpdate( } const parsedDelta = visibleDelta ? ctx.consumePartialReplyDirectives(visibleDelta) : null; const parsedFull = parseReplyDirectives(stripTrailingDirective(next)); - const cleanedText = stripTrailingPartialThinkingTag(parsedFull.text); + const cleanedText = stripTrailingPartialThinkingTagFragment(parsedFull.text); const mediaUrls = parsedDelta?.mediaUrls; const hasMedia = Boolean(mediaUrls && mediaUrls.length > 0); const hasAudio = Boolean(parsedDelta?.audioAsVoice); diff --git a/src/agents/pi-embedded-utils.e2e.test.ts b/src/agents/pi-embedded-utils.e2e.test.ts index ecb8dace5a1..ee13e527c2a 100644 --- a/src/agents/pi-embedded-utils.e2e.test.ts +++ b/src/agents/pi-embedded-utils.e2e.test.ts @@ -2,7 +2,9 @@ import type { AssistantMessage } from "@mariozechner/pi-ai"; import { describe, expect, it } from "vitest"; import { extractAssistantText, + extractThinkingFromTaggedStream, formatReasoningMessage, + stripTrailingPartialThinkingTagFragment, stripDowngradedToolCallText, } from "./pi-embedded-utils.js"; @@ -603,6 +605,19 @@ describe("formatReasoningMessage", () => { }); }); +describe("thinking tag fragment handling", () => { + it("strips dangling closing tag fragments from tagged stream extraction", () => { + expect(extractThinkingFromTaggedStream("step one and twostep one and two { + expect(stripTrailingPartialThinkingTagFragment("Reasoning line { it("strips [Historical context: ...] blocks", () => { const text = `[Historical context: a different model called tool "exec" with arguments {"command":"git status"}]`; diff --git a/src/agents/pi-embedded-utils.ts b/src/agents/pi-embedded-utils.ts index 2f38cecb697..2bae7aa5ac7 100644 --- a/src/agents/pi-embedded-utils.ts +++ b/src/agents/pi-embedded-utils.ts @@ -390,15 +390,31 @@ export function extractThinkingFromTaggedText(text: string): string { return result.trim(); } +export function stripTrailingPartialThinkingTagFragment(text: string): string { + if (!text) { + return text; + } + const match = text.match(/<\s*\/?\s*([a-z]*)\s*$/i); + if (!match || typeof match.index !== "number") { + return text; + } + const prefix = (match[1] ?? "").toLowerCase(); + const targets = ["think", "thinking", "thought", "antthinking"]; + const isThinkingPrefix = + prefix.length === 0 || targets.some((target) => target.startsWith(prefix)); + if (!isThinkingPrefix) { + return text; + } + return text.slice(0, match.index).trimEnd(); +} + export function extractThinkingFromTaggedStream(text: string): string { if (!text) { return ""; } - const stripTrailingPartialTag = (value: string) => - value.replace(/<\s*\/?\s*(?:think(?:ing)?|thought|antthinking)\s*$/i, "").trimEnd(); const closed = extractThinkingFromTaggedText(text); if (closed) { - return stripTrailingPartialTag(closed); + return stripTrailingPartialThinkingTagFragment(closed); } const openRe = /<\s*(?:think(?:ing)?|thought|antthinking)\s*>/gi; @@ -411,10 +427,10 @@ export function extractThinkingFromTaggedStream(text: string): string { const lastOpen = openMatches[openMatches.length - 1]; const lastClose = closeMatches[closeMatches.length - 1]; if (lastClose && (lastClose.index ?? -1) > (lastOpen.index ?? -1)) { - return stripTrailingPartialTag(closed); + return stripTrailingPartialThinkingTagFragment(closed); } const start = (lastOpen.index ?? 0) + lastOpen[0].length; - return stripTrailingPartialTag(text.slice(start).trim()); + return stripTrailingPartialThinkingTagFragment(text.slice(start).trim()); } export function inferToolMetaFromArgs(toolName: string, args: unknown): string | undefined { diff --git a/src/auto-reply/reply/agent-runner-execution.ts b/src/auto-reply/reply/agent-runner-execution.ts index 4becf72c780..3fdafd821bd 100644 --- a/src/auto-reply/reply/agent-runner-execution.ts +++ b/src/auto-reply/reply/agent-runner-execution.ts @@ -321,11 +321,11 @@ export async function runAgentTurnWithFallback(params: { onReasoningStream: params.typingSignals.shouldStartOnReasoning || params.opts?.onReasoningStream ? async (payload) => { - await params.typingSignals.signalReasoningDelta(); await params.opts?.onReasoningStream?.({ text: payload.text, mediaUrls: payload.mediaUrls, }); + await params.typingSignals.signalReasoningDelta(); } : undefined, onReasoningEnd: params.opts?.onReasoningEnd,