fix(agents): advance lightweight text deltas

vincentkoc · vincentkoc · commit 9569e072e535 · 2026-06-09T06:05:26.000+02:00
diff --git a/packages/agent-core/src/agent-loop.test.ts b/packages/agent-core/src/agent-loop.test.ts
@@ -76,12 +76,12 @@ describe("agentLoop EventStream failures", () => {
 });
 
 describe("agentLoop streaming updates", () => {
-  it("reuses the current assistant message for text deltas without partial snapshots", async () => {
+  it("rebuilds assistant message snapshots for text deltas without partial snapshots", async () => {
     const streamFn: StreamFn = async () => {
       const stream = createAssistantMessageEventStream();
-      const message: AssistantMessage = {
+      const startMessage: AssistantMessage = {
         role: "assistant",
-        content: [{ type: "text", text: "" }],
+        content: [],
         api: model.api,
         provider: model.provider,
         model: model.id,
@@ -96,17 +96,24 @@ describe("agentLoop streaming updates", () => {
         stopReason: "stop",
         timestamp: 1,
       };
+      const textStartMessage: AssistantMessage = { ...startMessage, content: [] };
+      const finalMessage: AssistantMessage = {
+        ...startMessage,
+        content: [{ type: "text", text: "Hello world" }],
+      };
 
       queueMicrotask(() => {
-        stream.push({ type: "start", partial: message });
-        stream.push({ type: "text_start", contentIndex: 0, partial: message });
-        const textBlock = message.content[0];
-        if (textBlock?.type === "text") {
-          textBlock.text = "Hello";
-        }
+        stream.push({ type: "start", partial: startMessage });
+        stream.push({ type: "text_start", contentIndex: 0, partial: textStartMessage });
         stream.push({ type: "text_delta", contentIndex: 0, delta: "Hello" });
-        stream.push({ type: "text_end", contentIndex: 0, content: "Hello", partial: message });
-        stream.push({ type: "done", reason: "stop", message });
+        stream.push({ type: "text_delta", contentIndex: 0, delta: " world" });
+        stream.push({
+          type: "text_end",
+          contentIndex: 0,
+          content: "Hello world",
+          partial: finalMessage,
+        });
+        stream.push({ type: "done", reason: "stop", message: finalMessage });
       });
 
       return stream;
@@ -121,15 +128,17 @@ describe("agentLoop streaming updates", () => {
     );
     const events = await collectEvents(stream);
 
-    const deltaUpdate = events.find(
+    const deltaUpdates = events.filter(
       (event): event is Extract<AgentEvent, { type: "message_update" }> =>
         event.type === "message_update" && event.assistantMessageEvent.type === "text_delta",
     );
-    expect(deltaUpdate).toMatchObject({
-      type: "message_update",
-      message: { role: "assistant", content: [{ type: "text", text: "Hello" }] },
-      assistantMessageEvent: { type: "text_delta", delta: "Hello" },
-    });
-    expect(deltaUpdate?.assistantMessageEvent).not.toHaveProperty("partial");
+    expect(deltaUpdates).toHaveLength(2);
+    expect(deltaUpdates.map((event) => event.message)).toMatchObject([
+      { role: "assistant", content: [{ type: "text", text: "Hello" }] },
+      { role: "assistant", content: [{ type: "text", text: "Hello world" }] },
+    ]);
+    for (const update of deltaUpdates) {
+      expect(update.assistantMessageEvent).not.toHaveProperty("partial");
+    }
   });
 });
diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts
@@ -3,6 +3,7 @@
 import { EventStream as LlmEventStream } from "@openclaw/llm-core";
 import type {
   AssistantMessage,
+  AssistantMessageEvent,
   Context,
   EventStream,
   ToolResultMessage,
@@ -35,6 +36,49 @@ const EMPTY_USAGE = {
 
 const EventStreamConstructor: typeof SourceEventStream = LlmEventStream;
 
+type AssistantMessageUpdateEvent = Extract<
+  AssistantMessageEvent,
+  {
+    type:
+      | "text_start"
+      | "text_delta"
+      | "text_end"
+      | "thinking_start"
+      | "thinking_delta"
+      | "thinking_end"
+      | "toolcall_start"
+      | "toolcall_delta"
+      | "toolcall_end";
+  }
+>;
+
+function appendTextDeltaToAssistantMessage(
+  message: AssistantMessage,
+  contentIndex: number,
+  delta: string,
+): AssistantMessage {
+  const content = [...message.content];
+  const currentContent = content[contentIndex];
+  content[contentIndex] =
+    currentContent?.type === "text"
+      ? { ...currentContent, text: currentContent.text + delta }
+      : { type: "text", text: delta };
+  return { ...message, content };
+}
+
+function resolveAssistantMessageUpdate(
+  event: AssistantMessageUpdateEvent,
+  currentMessage: AssistantMessage,
+): AssistantMessage {
+  if ("partial" in event && event.partial) {
+    return event.partial;
+  }
+  if (event.type === "text_delta") {
+    return appendTextDeltaToAssistantMessage(currentMessage, event.contentIndex, event.delta);
+  }
+  return currentMessage;
+}
+
 /**
  * Start an agent loop with a new prompt message.
  * The prompt is added to the context and events are emitted for it.
@@ -402,7 +446,7 @@ async function streamAssistantResponse(
       case "toolcall_delta":
       case "toolcall_end":
         if (partialMessage) {
-          const message: AssistantMessage = event.partial ?? partialMessage;
+          const message = resolveAssistantMessageUpdate(event, partialMessage);
           partialMessage = message;
           context.messages[context.messages.length - 1] = message;
           await emit({
diff --git a/packages/llm-core/src/types.ts b/packages/llm-core/src/types.ts
@@ -366,6 +366,11 @@ export interface Context {
 export type AssistantMessageEvent =
   | { type: "start"; partial: AssistantMessage }
   | { type: "text_start"; contentIndex: number; partial: AssistantMessage }
+  /**
+   * Plain text deltas may omit `partial` to avoid retaining one full assistant
+   * snapshot per token. Consumers that need current text should replay `delta`
+   * from the latest start/end partial checkpoint.
+   */
   | { type: "text_delta"; contentIndex: number; delta: string; partial?: AssistantMessage }
   | { type: "text_end"; contentIndex: number; content: string; partial: AssistantMessage }
   | { type: "thinking_start"; contentIndex: number; partial: AssistantMessage }
