fix(gateway): preserve streamed prefixes across tool boundaries

src/gateway/server-chat.agent-events.test.ts

@@ -310,6 +310,98 @@ describe("agent event handler", () => {
     nowSpy.mockRestore();
   });
 
+  it("preserves pre-tool assistant text when later segments stream as non-prefix snapshots", () => {
+    let now = 10_500;
+    const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now);
+    const { broadcast, nodeSendToSession, chatRunState, handler } = createHarness();
+    chatRunState.registry.add("run-segmented", {
+      sessionKey: "session-segmented",
+      clientRunId: "client-segmented",
+    });
+
+    handler({
+      runId: "run-segmented",
+      seq: 1,
+      stream: "assistant",
+      ts: Date.now(),
+      data: { text: "Before tool call", delta: "Before tool call" },
+    });
+
+    now = 10_700;
+    handler({
+      runId: "run-segmented",
+      seq: 2,
+      stream: "assistant",
+      ts: Date.now(),
+      data: { text: "After tool call", delta: "\nAfter tool call" },
+    });
+
+    emitLifecycleEnd(handler, "run-segmented", 3);
+
+    const chatCalls = chatBroadcastCalls(broadcast);
+    expect(chatCalls).toHaveLength(3);
+    const secondPayload = chatCalls[1]?.[1] as {
+      state?: string;
+      message?: { content?: Array<{ text?: string }> };
+    };
+    const finalPayload = chatCalls[2]?.[1] as {
+      state?: string;
+      message?: { content?: Array<{ text?: string }> };
+    };
+    expect(secondPayload.state).toBe("delta");
+    expect(secondPayload.message?.content?.[0]?.text).toBe("Before tool call\nAfter tool call");
+    expect(finalPayload.state).toBe("final");
+    expect(finalPayload.message?.content?.[0]?.text).toBe("Before tool call\nAfter tool call");
+    expect(sessionChatCalls(nodeSendToSession)).toHaveLength(3);
+    nowSpy.mockRestore();
+  });
+
+  it("flushes merged segmented text before final when latest segment is throttled", () => {
+    let now = 10_800;
+    const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now);
+    const { broadcast, nodeSendToSession, chatRunState, handler } = createHarness();
+    chatRunState.registry.add("run-segmented-flush", {
+      sessionKey: "session-segmented-flush",
+      clientRunId: "client-segmented-flush",
+    });
+
+    handler({
+      runId: "run-segmented-flush",
+      seq: 1,
+      stream: "assistant",
+      ts: Date.now(),
+      data: { text: "Before tool call", delta: "Before tool call" },
+    });
+
+    now = 10_860;
+    handler({
+      runId: "run-segmented-flush",
+      seq: 2,
+      stream: "assistant",
+      ts: Date.now(),
+      data: { text: "After tool call", delta: "\nAfter tool call" },
+    });
+
+    emitLifecycleEnd(handler, "run-segmented-flush", 3);
+
+    const chatCalls = chatBroadcastCalls(broadcast);
+    expect(chatCalls).toHaveLength(3);
+    const flushPayload = chatCalls[1]?.[1] as {
+      state?: string;
+      message?: { content?: Array<{ text?: string }> };
+    };
+    const finalPayload = chatCalls[2]?.[1] as {
+      state?: string;
+      message?: { content?: Array<{ text?: string }> };
+    };
+    expect(flushPayload.state).toBe("delta");
+    expect(flushPayload.message?.content?.[0]?.text).toBe("Before tool call\nAfter tool call");
+    expect(finalPayload.state).toBe("final");
+    expect(finalPayload.message?.content?.[0]?.text).toBe("Before tool call\nAfter tool call");
+    expect(sessionChatCalls(nodeSendToSession)).toHaveLength(3);
+    nowSpy.mockRestore();
+  });
+
   it("does not flush an extra delta when the latest text already broadcast", () => {
     let now = 11_000;
     const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now);

src/gateway/server-chat.ts

@@ -89,6 +89,48 @@ function isSilentReplyLeadFragment(text: string): boolean {
   return SILENT_REPLY_TOKEN.startsWith(normalized);
 }
 
+function appendUniqueSuffix(base: string, suffix: string): string {
+  if (!suffix) {
+    return base;
+  }
+  if (!base) {
+    return suffix;
+  }
+  if (base.endsWith(suffix)) {
+    return base;
+  }
+  const maxOverlap = Math.min(base.length, suffix.length);
+  for (let overlap = maxOverlap; overlap > 0; overlap -= 1) {
+    if (base.slice(-overlap) === suffix.slice(0, overlap)) {
+      return base + suffix.slice(overlap);
+    }
+  }
+  return base + suffix;
+}
+
+function resolveMergedAssistantText(params: {
+  previousText: string;
+  nextText: string;
+  nextDelta: string;
+}) {
+  const { previousText, nextText, nextDelta } = params;
+  if (nextText && previousText) {
+    if (nextText.startsWith(previousText)) {
+      return nextText;
+    }
+    if (previousText.startsWith(nextText) && !nextDelta) {
+      return previousText;
+    }
+  }
+  if (nextDelta) {
+    return appendUniqueSuffix(previousText, nextDelta);
+  }
+  if (nextText) {
+    return nextText;
+  }
+  return previousText;
+}
+
 export type ChatRunEntry = {
   sessionKey: string;
   clientRunId: string;
@@ -302,16 +344,25 @@ export function createAgentEventHandler({
     sourceRunId: string,
     seq: number,
     text: string,
+    delta?: unknown,
   ) => {
-    const cleaned = stripInlineDirectiveTagsForDisplay(text).text;
-    if (!cleaned) {
+    const cleanedText = stripInlineDirectiveTagsForDisplay(text).text;
+    const cleanedDelta =
+      typeof delta === "string" ? stripInlineDirectiveTagsForDisplay(delta).text : "";
+    const previousText = chatRunState.buffers.get(clientRunId) ?? "";
+    const mergedText = resolveMergedAssistantText({
+      previousText,
+      nextText: cleanedText,
+      nextDelta: cleanedDelta,
+    });
+    if (!mergedText) {
       return;
     }
-    chatRunState.buffers.set(clientRunId, cleaned);
-    if (isSilentReplyText(cleaned, SILENT_REPLY_TOKEN)) {
+    chatRunState.buffers.set(clientRunId, mergedText);
+    if (isSilentReplyText(mergedText, SILENT_REPLY_TOKEN)) {
       return;
     }
-    if (isSilentReplyLeadFragment(cleaned)) {
+    if (isSilentReplyLeadFragment(mergedText)) {
       return;
     }
     if (shouldHideHeartbeatChatOutput(clientRunId, sourceRunId)) {
@@ -323,7 +374,7 @@ export function createAgentEventHandler({
       return;
     }
     chatRunState.deltaSentAt.set(clientRunId, now);
-    chatRunState.deltaLastBroadcastLen.set(clientRunId, cleaned.length);
+    chatRunState.deltaLastBroadcastLen.set(clientRunId, mergedText.length);
     const payload = {
       runId: clientRunId,
       sessionKey,
@@ -331,7 +382,7 @@ export function createAgentEventHandler({
       state: "delta" as const,
       message: {
         role: "assistant",
-        content: [{ type: "text", text: cleaned }],
+        content: [{ type: "text", text: mergedText }],
         timestamp: now,
       },
     };
@@ -512,7 +563,7 @@ export function createAgentEventHandler({
         nodeSendToSession(sessionKey, "agent", isToolEvent ? toolPayload : agentPayload);
       }
       if (!isAborted && evt.stream === "assistant" && typeof evt.data?.text === "string") {
-        emitChatDelta(sessionKey, clientRunId, evt.runId, evt.seq, evt.data.text);
+        emitChatDelta(sessionKey, clientRunId, evt.runId, evt.seq, evt.data.text, evt.data.delta);
       } else if (!isAborted && (lifecyclePhase === "end" || lifecyclePhase === "error")) {
         const evtStopReason =
           typeof evt.data?.stopReason === "string" ? evt.data.stopReason : undefined;