Embedded: stop replaying historical session replies

Openbling · Openbling · commit 81db13631bad · 2026-03-23T11:43:13.000+08:00
diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts
@@ -38,6 +38,22 @@ const stripTrailingDirective = (text: string): string => {
   return text.slice(0, openIndex);
 };
 
+function resolveEligibleAssistantMessage(
+  ctx: EmbeddedPiSubscribeContext,
+  message: AgentMessage | undefined,
+) {
+  if (!message || message.role !== "assistant") {
+    return null;
+  }
+  if (ctx.state.initialReplayInProgress) {
+    return null;
+  }
+  if (ctx.state.preexistingMessages.has(message)) {
+    return null;
+  }
+  return message;
+}
+
 function emitReasoningEnd(ctx: EmbeddedPiSubscribeContext) {
   if (!ctx.state.reasoningStreamOpen) {
     return;
@@ -133,8 +149,8 @@ export function handleMessageStart(
   ctx: EmbeddedPiSubscribeContext,
   evt: AgentEvent & { message: AgentMessage },
 ) {
-  const msg = evt.message;
-  if (msg?.role !== "assistant") {
+  const msg = resolveEligibleAssistantMessage(ctx, evt.message);
+  if (!msg) {
     return;
   }
 
@@ -152,8 +168,8 @@ export function handleMessageUpdate(
   ctx: EmbeddedPiSubscribeContext,
   evt: AgentEvent & { message: AgentMessage; assistantMessageEvent?: unknown },
 ) {
-  const msg = evt.message;
-  if (msg?.role !== "assistant") {
+  const msg = resolveEligibleAssistantMessage(ctx, evt.message);
+  if (!msg) {
     return;
   }
 
@@ -322,8 +338,8 @@ export function handleMessageEnd(
   ctx: EmbeddedPiSubscribeContext,
   evt: AgentEvent & { message: AgentMessage },
 ) {
-  const msg = evt.message;
-  if (msg?.role !== "assistant") {
+  const msg = resolveEligibleAssistantMessage(ctx, evt.message);
+  if (!msg) {
     return;
   }
 
diff --git a/src/agents/pi-embedded-subscribe.handlers.ts b/src/agents/pi-embedded-subscribe.handlers.ts
@@ -21,6 +21,9 @@ import type {
 
 export function createEmbeddedPiSessionEventHandler(ctx: EmbeddedPiSubscribeContext) {
   return (evt: EmbeddedPiSubscribeEvent) => {
+    if (ctx.state.initialReplayInProgress) {
+      return;
+    }
     switch (evt.type) {
       case "message_start":
         handleMessageStart(ctx, evt as never);
diff --git a/src/agents/pi-embedded-subscribe.handlers.types.ts b/src/agents/pi-embedded-subscribe.handlers.types.ts
@@ -33,10 +33,12 @@ export type ToolCallSummary = {
 
 export type EmbeddedPiSubscribeState = {
   assistantTexts: string[];
+  preexistingMessages: Set<AgentMessage>;
   toolMetas: Array<{ toolName?: string; meta?: string }>;
   toolMetaById: Map<string, ToolCallSummary>;
   toolSummaryById: Set<string>;
   lastToolError?: ToolErrorSummary;
+  initialReplayInProgress: boolean;
 
   blockReplyBreak: "text_end" | "message_end";
   reasoningMode: ReasoningLevel;
diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.test.ts
@@ -12,6 +12,32 @@ import {
 import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
 
 describe("subscribeEmbeddedPiSession", () => {
+  function createReplayingAgentEventHarness(params: {
+    replayEvents: unknown[];
+    existingMessages?: unknown[];
+  }) {
+    let handler: ((evt: unknown) => void) | undefined;
+    const session = {
+      messages: params.existingMessages ?? [],
+      subscribe: (fn: (evt: unknown) => void) => {
+        handler = fn;
+        for (const evt of params.replayEvents) {
+          fn(evt);
+        }
+        return () => {};
+      },
+    } as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"];
+    const onAgentEvent = vi.fn();
+
+    subscribeEmbeddedPiSession({
+      session,
+      runId: "run",
+      onAgentEvent,
+    });
+
+    return { emit: (evt: unknown) => handler?.(evt), onAgentEvent };
+  }
+
   function createAgentEventHarness(options?: { runId?: string; sessionKey?: string }) {
     const { session, emit } = createStubSessionHarness();
     const onAgentEvent = vi.fn();
@@ -293,6 +319,28 @@ describe("subscribeEmbeddedPiSession", () => {
     expectSingleAgentEventText(onAgentEvent.mock.calls, "Hello world");
   });
 
+  it("does not replay historical assistant finals when subscribing to an existing session", () => {
+    const historicalAssistantMessage = {
+      role: "assistant",
+      content: [{ type: "text", text: "Earlier reply" }],
+    } as AssistantMessage;
+    const { emit, onAgentEvent } = createReplayingAgentEventHarness({
+      existingMessages: [historicalAssistantMessage],
+      replayEvents: [
+        { type: "agent_start" },
+        { type: "message_start", message: historicalAssistantMessage },
+        { type: "message_end", message: historicalAssistantMessage },
+        { type: "agent_end" },
+      ],
+    });
+
+    emitMessageStartAndEndForAssistantText({ emit, text: "Current reply" });
+
+    const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls);
+    expect(payloads).toHaveLength(1);
+    expect(payloads[0]?.text).toBe("Current reply");
+  });
+
   it("does not emit duplicate agent events when message_end repeats", () => {
     const { emit, onAgentEvent } = createAgentEventHarness();
 
diff --git a/src/agents/pi-embedded-subscribe.ts b/src/agents/pi-embedded-subscribe.ts
@@ -39,10 +39,12 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
   const useMarkdown = toolResultFormat === "markdown";
   const state: EmbeddedPiSubscribeState = {
     assistantTexts: [],
+    preexistingMessages: new Set(params.session.messages ?? []),
     toolMetas: [],
     toolMetaById: new Map(),
     toolSummaryById: new Set(),
     lastToolError: undefined,
+    initialReplayInProgress: true,
     blockReplyBreak: params.blockReplyBreak ?? "text_end",
     reasoningMode,
     includeReasoning: reasoningMode === "on",
@@ -652,6 +654,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
   };
 
   const sessionUnsubscribe = params.session.subscribe(createEmbeddedPiSessionEventHandler(ctx));
+  state.initialReplayInProgress = false;
 
   const unsubscribe = () => {
     if (state.unsubscribed) {
