fix: preserve rewritten stream snapshots in webchat (#58641) (thanks @neeravmakwana)

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

@@ -69,11 +69,13 @@ describe("buildAssistantStreamData", () => {
       buildAssistantStreamData({
         text: "hello",
         delta: "he",
+        replace: true,
         mediaUrl: "https://example.com/a.png",
       }),
     ).toEqual({
       text: "hello",
       delta: "he",
+      replace: true,
       mediaUrls: ["https://example.com/a.png"],
     });
   });

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

@@ -153,13 +153,15 @@ export function hasAssistantVisibleReply(params: {
 export function buildAssistantStreamData(params: {
   text?: string;
   delta?: string;
+  replace?: boolean;
   mediaUrls?: string[];
   mediaUrl?: string;
-}): { text: string; delta: string; mediaUrls?: string[] } {
+}): { text: string; delta: string; replace?: true; mediaUrls?: string[] } {
   const mediaUrls = resolveSendableOutboundReplyParts(params).mediaUrls;
   return {
     text: params.text ?? "",
     delta: params.delta ?? "",
+    replace: params.replace ? true : undefined,
     mediaUrls: mediaUrls.length ? mediaUrls : undefined,
   };
 }
@@ -310,13 +312,15 @@ export function handleMessageUpdate(
 
     let shouldEmit = false;
     let deltaText = "";
+    let replace = false;
     if (!hasAssistantVisibleReply({ text: cleanedText, mediaUrls, audioAsVoice: hasAudio })) {
       shouldEmit = false;
-    } else if (previousCleaned && !cleanedText.startsWith(previousCleaned)) {
-      shouldEmit = false;
     } else {
-      deltaText = cleanedText.slice(previousCleaned.length);
-      shouldEmit = Boolean(deltaText || hasMedia || hasAudio);
+      replace = Boolean(previousCleaned && !cleanedText.startsWith(previousCleaned));
+      deltaText = replace ? "" : cleanedText.slice(previousCleaned.length);
+      shouldEmit = replace
+        ? cleanedText !== previousCleaned || hasMedia || hasAudio
+        : Boolean(deltaText || hasMedia || hasAudio);
     }
 
     ctx.state.lastStreamedAssistant = next;
@@ -330,6 +334,7 @@ export function handleMessageUpdate(
       const data = buildAssistantStreamData({
         text: cleanedText,
         delta: deltaText,
+        replace,
         mediaUrls,
       });
       emitAgentEvent({

src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.test.ts

@@ -316,7 +316,7 @@ describe("subscribeEmbeddedPiSession", () => {
     expect(payloads).toHaveLength(1);
   });
 
-  it("skips agent events when cleaned text rewinds mid-stream", () => {
+  it("emits a replacement snapshot when cleaned text rewinds mid-stream", () => {
     const { emit, onAgentEvent } = createAgentEventHarness();
 
     emit({ type: "message_start", message: { role: "assistant" } });
@@ -324,8 +324,13 @@ describe("subscribeEmbeddedPiSession", () => {
     emitAssistantTextDelta(emit, " https://example.com/a.png\nCaption");
 
     const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls);
-    expect(payloads).toHaveLength(1);
+    expect(payloads).toHaveLength(2);
     expect(payloads[0]?.text).toBe("MEDIA:");
+    expect(payloads[0]?.delta).toBe("MEDIA:");
+    expect(payloads[0]?.replace).toBeUndefined();
+    expect(payloads[1]?.text).toBe("Caption");
+    expect(payloads[1]?.delta).toBe("");
+    expect(payloads[1]?.replace).toBe(true);
   });
 
   it("emits agent events when media arrives without text", () => {

src/gateway/openresponses-http.ts

@@ -920,6 +920,11 @@ export async function handleOpenResponsesHttpRequest(
     }
 
     if (evt.stream === "assistant") {
+      const text = evt.data?.text;
+      const replace = evt.data?.replace === true;
+      if (replace && typeof text === "string") {
+        accumulatedText = text;
+      }
       const content = resolveAssistantStreamDeltaText(evt);
       if (!content) {
         return;