fix: restore telegram draft streaming partials

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

@@ -1,6 +1,5 @@
 import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";
 
-import { parseReplyDirectives } from "../auto-reply/reply/reply-directives.js";
 import { emitAgentEvent } from "../infra/agent-events.js";
 import {
   isMessagingToolDuplicateNormalized,
@@ -111,35 +110,39 @@ export function handleMessageUpdate(
     })
     .trim();
   if (next && next !== ctx.state.lastStreamedAssistant) {
-    const previousText = ctx.state.lastStreamedAssistant ?? "";
-    const { text: cleanedText, mediaUrls } = parseReplyDirectives(next);
-    const { text: previousCleanedText } = parseReplyDirectives(previousText);
-    if (cleanedText.startsWith(previousCleanedText)) {
-      const deltaText = cleanedText.slice(previousCleanedText.length);
+    const parsedDelta = chunk ? ctx.consumePartialReplyDirectives(chunk) : null;
+    const deltaText = parsedDelta?.text ?? "";
+    const mediaUrls = parsedDelta?.mediaUrls;
+    if (!deltaText && (!mediaUrls || mediaUrls.length === 0) && !parsedDelta?.audioAsVoice) {
       ctx.state.lastStreamedAssistant = next;
-      emitAgentEvent({
-        runId: ctx.params.runId,
-        stream: "assistant",
-        data: {
-          text: cleanedText,
-          delta: deltaText,
-          mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
-        },
+      return;
+    }
+    const previousCleaned = ctx.state.lastStreamedAssistantCleaned ?? "";
+    const cleanedText = `${previousCleaned}${deltaText}`;
+    ctx.state.lastStreamedAssistant = next;
+    ctx.state.lastStreamedAssistantCleaned = cleanedText;
+    emitAgentEvent({
+      runId: ctx.params.runId,
+      stream: "assistant",
+      data: {
+        text: cleanedText,
+        delta: deltaText,
+        mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+      },
+    });
+    void ctx.params.onAgentEvent?.({
+      stream: "assistant",
+      data: {
+        text: cleanedText,
+        delta: deltaText,
+        mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+      },
+    });
+    if (ctx.params.onPartialReply && ctx.state.shouldEmitPartialReplies) {
+      void ctx.params.onPartialReply({
+        text: cleanedText,
+        mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
       });
-      void ctx.params.onAgentEvent?.({
-        stream: "assistant",
-        data: {
-          text: cleanedText,
-          delta: deltaText,
-          mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
-        },
-      });
-      if (ctx.params.onPartialReply && ctx.state.shouldEmitPartialReplies) {
-        void ctx.params.onPartialReply({
-          text: cleanedText,
-          mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
-        });
-      }
     }
   }
 

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

@@ -38,6 +38,7 @@ export type EmbeddedPiSubscribeState = {
   blockBuffer: string;
   blockState: { thinking: boolean; final: boolean; inlineCode: InlineCodeState };
   lastStreamedAssistant?: string;
+  lastStreamedAssistantCleaned?: string;
   lastStreamedReasoning?: string;
   lastBlockReplyText?: string;
   assistantMessageIndex: number;
@@ -82,6 +83,10 @@ export type EmbeddedPiSubscribeContext = {
     text: string,
     options?: { final?: boolean },
   ) => ReplyDirectiveParseResult | null;
+  consumePartialReplyDirectives: (
+    text: string,
+    options?: { final?: boolean },
+  ) => ReplyDirectiveParseResult | null;
   resetAssistantMessageState: (nextAssistantTextBaseline: number) => void;
   resetForCompactionRetry: () => void;
   finalizeAssistantTexts: (args: {

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

@@ -103,4 +103,50 @@ describe("subscribeEmbeddedPiSession reply tags", () => {
     expect(onBlockReply.mock.calls[0]?.[0]?.text).toBe("Hello");
     expect(onBlockReply.mock.calls[1]?.[0]?.text).toBe("[[");
   });
+
+  it("streams partial replies past reply_to tags split across chunks", () => {
+    let handler: ((evt: unknown) => void) | undefined;
+    const session: StubSession = {
+      subscribe: (fn) => {
+        handler = fn;
+        return () => {};
+      },
+    };
+
+    const onPartialReply = vi.fn();
+
+    subscribeEmbeddedPiSession({
+      session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+      runId: "run",
+      onPartialReply,
+    });
+
+    handler?.({ type: "message_start", message: { role: "assistant" } });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_delta", delta: "[[reply_to:1897" },
+    });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_delta", delta: "]] Hello" },
+    });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_delta", delta: " world" },
+    });
+    handler?.({
+      type: "message_update",
+      message: { role: "assistant" },
+      assistantMessageEvent: { type: "text_end" },
+    });
+
+    const lastPayload = onPartialReply.mock.calls.at(-1)?.[0];
+    expect(lastPayload?.text).toBe("Hello world");
+    for (const call of onPartialReply.mock.calls) {
+      expect(call[0]?.text?.includes("[[reply_to")).toBe(false);
+    }
+  });
 });

src/agents/pi-embedded-subscribe.ts

@@ -47,6 +47,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     // Track if a streamed chunk opened a <think> block (stateful across chunks).
     blockState: { thinking: false, final: false, inlineCode: createInlineCodeState() },
     lastStreamedAssistant: undefined,
+    lastStreamedAssistantCleaned: undefined,
     lastStreamedReasoning: undefined,
     lastBlockReplyText: undefined,
     assistantMessageIndex: 0,
@@ -77,16 +78,19 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
   const pendingMessagingTexts = state.pendingMessagingTexts;
   const pendingMessagingTargets = state.pendingMessagingTargets;
   const replyDirectiveAccumulator = createStreamingDirectiveAccumulator();
+  const partialReplyDirectiveAccumulator = createStreamingDirectiveAccumulator();
 
   const resetAssistantMessageState = (nextAssistantTextBaseline: number) => {
     state.deltaBuffer = "";
     state.blockBuffer = "";
     blockChunker?.reset();
     replyDirectiveAccumulator.reset();
+    partialReplyDirectiveAccumulator.reset();
     state.blockState.thinking = false;
     state.blockState.final = false;
     state.blockState.inlineCode = createInlineCodeState();
     state.lastStreamedAssistant = undefined;
+    state.lastStreamedAssistantCleaned = undefined;
     state.lastBlockReplyText = undefined;
     state.lastStreamedReasoning = undefined;
     state.lastReasoningSent = undefined;
@@ -447,6 +451,8 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
 
   const consumeReplyDirectives = (text: string, options?: { final?: boolean }) =>
     replyDirectiveAccumulator.consume(text, options);
+  const consumePartialReplyDirectives = (text: string, options?: { final?: boolean }) =>
+    partialReplyDirectiveAccumulator.consume(text, options);
 
   const flushBlockReplyBuffer = () => {
     if (!params.onBlockReply) {
@@ -509,6 +515,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     flushBlockReplyBuffer,
     emitReasoningStream,
     consumeReplyDirectives,
+    consumePartialReplyDirectives,
     resetAssistantMessageState,
     resetForCompactionRetry,
     finalizeAssistantTexts,