From bc1e95fb6231634cee57f65d04e8279c85558a07 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Fri, 31 Jul 2026 13:45:07 -0700
Subject: [PATCH] fix(streaming): preserve incomplete bracket tails in final
 replies (#116983)

Co-authored-by: Peter Steinberger <steipete@macos.shared>
---
 ...-agent-subscribe.handlers.messages.test.ts | 35 +++++++++++++++++++
 ...edded-agent-subscribe.handlers.messages.ts | 17 +++++----
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts
index e1fee6ba3961..0d1ca8ffe3eb 100644
--- a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts
+++ b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts
@@ -470,6 +470,41 @@ describe("handleMessageUpdate text signatures", () => {
     expect(context.state.lastStreamedAssistantCleaned).toBe("Hello");
   });
 
+  it.each([
+    {
+      name: "the directive accumulator has no parsed result",
+      text: "answer part A msg [[E1008]timeout] answer part B",
+      hasParsedDirectives: false,
+    },
+    {
+      name: "the directive accumulator flushes a buffered tail",
+      text: "answer part A msg [[E1008]timeout] answer part B",
+      hasParsedDirectives: true,
+    },
+    {
+      name: "the final text ends with one bracket",
+      text: "answer part A [",
+      hasParsedDirectives: true,
+    },
+  ])("keeps literal final text when $name", ({ text, hasParsedDirectives }) => {
+    const onAgentEvent = vi.fn();
+    const context = createMessageUpdateContext({
+      onAgentEvent,
+      ...(hasParsedDirectives ? {} : { consumePartialReplyDirectives: vi.fn(() => null) }),
+    });
+
+    updateMessage(context, {
+      message: { role: "assistant", content: [] },
+      assistantMessageEvent: { type: "text_end", content: text },
+    });
+
+    expect(context.state.lastStreamedAssistantCleaned).toBe(text);
+    expect(firstMockArg(onAgentEvent, "final assistant event")).toMatchObject({
+      stream: "assistant",
+      data: { text },
+    });
+  });
+
   it("keeps stripped reply directives out of later plain deltas", () => {
     const onAgentEvent = vi.fn();
     const context = createMessageUpdateContext({ onAgentEvent });
diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.ts b/src/agents/embedded-agent-subscribe.handlers.messages.ts
index 8f4fbc7d82c7..debeee775e91 100644
--- a/src/agents/embedded-agent-subscribe.handlers.messages.ts
+++ b/src/agents/embedded-agent-subscribe.handlers.messages.ts
@@ -616,10 +616,6 @@ function mergeReplyDirectiveResults(
   };
 }
 
-function parseFullStreamingReplyText(text: string): string {
-  return parseReplyDirectives(splitTrailingDirective(text).text).text;
-}
-
 function containsCompleteMediaDirectiveLine(text: string): boolean {
   return /(?:^|\n)\s*MEDIA:\s*\S[^\n]*(?:\n|$)/i.test(text);
 }
@@ -664,13 +660,16 @@ function resolveStreamingReplyText(params: {
   parsedStreamDirectives: ReplyDirectiveParseResult | null;
   shouldUsePhaseAwareBlockReply: boolean;
 }): string {
-  if (!params.parsedStreamDirectives) {
-    return params.evtType === "text_delta"
-      ? params.previousCleaned
-      : parseFullStreamingReplyText(params.next);
+  if (!params.parsedStreamDirectives && params.evtType === "text_delta") {
+    return params.previousCleaned;
   }
 
-  return resolveIncrementalStreamingReplyText(params) ?? parseFullStreamingReplyText(params.next);
+  return (
+    resolveIncrementalStreamingReplyText(params) ??
+    parseReplyDirectives(
+      params.evtType === "text_end" ? params.next : splitTrailingDirective(params.next).text,
+    ).text
+  );
 }
 
 /** Records parsed reply directives until a sendable reply payload is built. */