fix: keep partial think-tag fragments out of streams

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);

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("<think>step one and two</")).toBe("step one and two");
+    expect(extractThinkingFromTaggedStream("<think>step one and two</th")).toBe("step one and two");
+  });
+
+  it("strips dangling opening/closing think fragments while preserving plain text", () => {
+    expect(stripTrailingPartialThinkingTagFragment("Reasoning line <th")).toBe("Reasoning line");
+    expect(stripTrailingPartialThinkingTagFragment("Reasoning line </")).toBe("Reasoning line");
+    expect(stripTrailingPartialThinkingTagFragment("2 < 3")).toBe("2 < 3");
+  });
+});
+
 describe("stripDowngradedToolCallText", () => {
   it("strips [Historical context: ...] blocks", () => {
     const text = `[Historical context: a different model called tool "exec" with arguments {"command":"git status"}]`;

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 {

src/auto-reply/reply/agent-runner-execution.ts

@@ -301,11 +301,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,