fix: carry reply tags across streamed chunks

src/agents/pi-embedded-subscribe.handlers.messages.ts

@@ -226,24 +226,27 @@ export function handleMessageEnd(
         );
       } else {
         ctx.state.lastBlockReplyText = text;
-        const {
-          text: cleanedText,
-          mediaUrls,
-          audioAsVoice,
-          replyToId,
-          replyToTag,
-          replyToCurrent,
-        } = parseReplyDirectives(text);
-        // Emit if there's content OR audioAsVoice flag (to propagate the flag).
-        if (cleanedText || (mediaUrls && mediaUrls.length > 0) || audioAsVoice) {
-          void onBlockReply({
+        const splitResult = ctx.consumeReplyDirectives(text, { final: true });
+        if (splitResult) {
+          const {
             text: cleanedText,
-            mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+            mediaUrls,
             audioAsVoice,
             replyToId,
             replyToTag,
             replyToCurrent,
-          });
+          } = splitResult;
+          // Emit if there's content OR audioAsVoice flag (to propagate the flag).
+          if (cleanedText || (mediaUrls && mediaUrls.length > 0) || audioAsVoice) {
+            void onBlockReply({
+              text: cleanedText,
+              mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+              audioAsVoice,
+              replyToId,
+              replyToTag,
+              replyToCurrent,
+            });
+          }
         }
       }
     }
@@ -254,6 +257,30 @@ export function handleMessageEnd(
     ctx.emitReasoningStream(rawThinking);
   }
 
+  if (ctx.state.blockReplyBreak === "text_end" && onBlockReply) {
+    const tailResult = ctx.consumeReplyDirectives("", { final: true });
+    if (tailResult) {
+      const {
+        text: cleanedText,
+        mediaUrls,
+        audioAsVoice,
+        replyToId,
+        replyToTag,
+        replyToCurrent,
+      } = tailResult;
+      if (cleanedText || (mediaUrls && mediaUrls.length > 0) || audioAsVoice) {
+        void onBlockReply({
+          text: cleanedText,
+          mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+          audioAsVoice,
+          replyToId,
+          replyToTag,
+          replyToCurrent,
+        });
+      }
+    }
+  }
+
   ctx.state.deltaBuffer = "";
   ctx.state.blockBuffer = "";
   ctx.blockChunker?.reset();

src/agents/pi-embedded-subscribe.handlers.types.ts

@@ -1,6 +1,7 @@
 import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";
 
 import type { ReasoningLevel } from "../auto-reply/thinking.js";
+import type { ReplyDirectiveParseResult } from "../auto-reply/reply/reply-directives.js";
 import type { InlineCodeState } from "../markdown/code-spans.js";
 import type { EmbeddedBlockChunker } from "./pi-embedded-block-chunker.js";
 import type { MessagingToolSend } from "./pi-embedded-messaging.js";
@@ -77,6 +78,10 @@ export type EmbeddedPiSubscribeContext = {
   emitBlockChunk: (text: string) => void;
   flushBlockReplyBuffer: () => void;
   emitReasoningStream: (text: string) => void;
+  consumeReplyDirectives: (
+    text: string,
+    options?: { final?: boolean },
+  ) => ReplyDirectiveParseResult | null;
   resetAssistantMessageState: (nextAssistantTextBaseline: number) => void;
   resetForCompactionRetry: () => void;
   finalizeAssistantTexts: (args: {

src/agents/pi-embedded-subscribe.reply-tags.test.ts

@@ -0,0 +1,106 @@
+import type { AssistantMessage } from "@mariozechner/pi-ai";
+import { describe, expect, it, vi } from "vitest";
+import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
+
+type StubSession = {
+  subscribe: (fn: (evt: unknown) => void) => () => void;
+};
+
+describe("subscribeEmbeddedPiSession reply tags", () => {
+  it("carries reply_to_current across tag-only block chunks", () => {
+    let handler: ((evt: unknown) => void) | undefined;
+    const session: StubSession = {
+      subscribe: (fn) => {
+        handler = fn;
+        return () => {};
+      },
+    };
+
+    const onBlockReply = vi.fn();
+
+    subscribeEmbeddedPiSession({
+      session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+      runId: "run",
+      onBlockReply,
+      blockReplyBreak: "text_end",
+      blockReplyChunking: {
+        minChars: 1,
+        maxChars: 50,
+        breakPreference: "newline",
+      },
+    });
+
+    handler?.({ type: "message_start", message: { role: "assistant" } });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: {
+        type: "text_delta",
+        delta: "[[reply_to_current]]\nHello",
+      },
+    });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_end" },
+    });
+
+    const assistantMessage = {
+      role: "assistant",
+      content: [{ type: "text", text: "[[reply_to_current]]\nHello" }],
+    } as AssistantMessage;
+    handler?.({ type: "message_end", message: assistantMessage });
+
+    expect(onBlockReply).toHaveBeenCalledTimes(1);
+    const payload = onBlockReply.mock.calls[0]?.[0];
+    expect(payload?.text).toBe("Hello");
+    expect(payload?.replyToCurrent).toBe(true);
+    expect(payload?.replyToTag).toBe(true);
+  });
+
+  it("flushes trailing directive tails on stream end", () => {
+    let handler: ((evt: unknown) => void) | undefined;
+    const session: StubSession = {
+      subscribe: (fn) => {
+        handler = fn;
+        return () => {};
+      },
+    };
+
+    const onBlockReply = vi.fn();
+
+    subscribeEmbeddedPiSession({
+      session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+      runId: "run",
+      onBlockReply,
+      blockReplyBreak: "text_end",
+      blockReplyChunking: {
+        minChars: 1,
+        maxChars: 50,
+        breakPreference: "newline",
+      },
+    });
+
+    handler?.({ type: "message_start", message: { role: "assistant" } });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_delta", delta: "Hello [[" },
+    });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_end" },
+    });
+
+    const assistantMessage = {
+      role: "assistant",
+      content: [{ type: "text", text: "Hello [[" }],
+    } as AssistantMessage;
+    handler?.({ type: "message_end", message: assistantMessage });
+
+    expect(onBlockReply).toHaveBeenCalledTimes(2);
+    expect(onBlockReply.mock.calls[0]?.[0]?.text).toBe("Hello");
+    expect(onBlockReply.mock.calls[1]?.[0]?.text).toBe("[[");
+  });
+});

src/agents/pi-embedded-subscribe.ts

@@ -1,4 +1,5 @@
 import { parseReplyDirectives } from "../auto-reply/reply/reply-directives.js";
+import { createStreamingDirectiveAccumulator } from "../auto-reply/reply/streaming-directives.js";
 import { formatToolAggregate } from "../auto-reply/tool-meta.js";
 import { createSubsystemLogger } from "../logging/subsystem.js";
 import type { InlineCodeState } from "../markdown/code-spans.js";
@@ -75,11 +76,13 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
   const messagingToolSentTargets = state.messagingToolSentTargets;
   const pendingMessagingTexts = state.pendingMessagingTexts;
   const pendingMessagingTargets = state.pendingMessagingTargets;
+  const replyDirectiveAccumulator = createStreamingDirectiveAccumulator();
 
   const resetAssistantMessageState = (nextAssistantTextBaseline: number) => {
     state.deltaBuffer = "";
     state.blockBuffer = "";
     blockChunker?.reset();
+    replyDirectiveAccumulator.reset();
     state.blockState.thinking = false;
     state.blockState.final = false;
     state.blockState.inlineCode = createInlineCodeState();
@@ -374,7 +377,8 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     assistantTexts.push(chunk);
     rememberAssistantText(chunk);
     if (!params.onBlockReply) return;
-    const splitResult = parseReplyDirectives(chunk);
+    const splitResult = replyDirectiveAccumulator.consume(chunk);
+    if (!splitResult) return;
     const {
       text: cleanedText,
       mediaUrls,
@@ -395,6 +399,9 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     });
   };
 
+  const consumeReplyDirectives = (text: string, options?: { final?: boolean }) =>
+    replyDirectiveAccumulator.consume(text, options);
+
   const flushBlockReplyBuffer = () => {
     if (!params.onBlockReply) return;
     if (blockChunker?.hasBuffered()) {
@@ -447,6 +454,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     emitBlockChunk,
     flushBlockReplyBuffer,
     emitReasoningStream,
+    consumeReplyDirectives,
     resetAssistantMessageState,
     resetForCompactionRetry,
     finalizeAssistantTexts,