Fix draft agent handoff continuity

Move optimistic user message reconciliation into the stream reducers and keep draft-created agents in a handoff lifecycle until authoritative history lands.
2026-05-25 20:53:49 +07:00
parent e5658654c0
commit 53faba64ca
21 changed files with 677 additions and 458 deletions
@@ -1,6 +1,11 @@
 import { describe, expect, it } from "vitest";
 import type { AgentStreamEventPayload } from "@server/shared/messages";
-import { hydrateStreamState, type AgentToolCallItem, type StreamItem } from "@/types/stream";
+import {
+  buildOptimisticUserMessage,
+  hydrateStreamState,
+  type AgentToolCallItem,
+  type StreamItem,
+} from "@/types/stream";
 import {
  createAgentStreamReducerQueue,
  processTimelineResponse,
@@ -103,6 +108,17 @@ function makeAssistantItem(text: string, id = `assistant-${text.length}`): Strea
  };
 }

+function makeOptimisticUserMessage(
+  text: string,
+  id = `optimistic-${text.length}`,
+): Extract<StreamItem, { kind: "user_message" }> {
+  return buildOptimisticUserMessage({
+    id,
+    text,
+    timestamp: new Date(1000),
+  });
+}
+
 function getAssistantTexts(items: StreamItem[]): string[] {
  return items
    .filter((item): item is Extract<StreamItem, { kind: "assistant_message" }> => {
@@ -275,6 +291,76 @@ describe("processTimelineResponse", () => {
    expect(assistant?.timestamp.toISOString()).toBe("2025-01-01T12:00:04.000Z");
  });

+  it("reconciles an optimistic user message during tail replacement", () => {
+    const image = {
+      id: "optimistic-image",
+      mimeType: "image/png",
+      storageType: "web-indexeddb" as const,
+      storageKey: "optimistic-image",
+      createdAt: 1000,
+    };
+    const attachment = {
+      type: "text" as const,
+      mimeType: "text/plain" as const,
+      text: "attached context",
+      title: "context.txt",
+    };
+    const optimistic = buildOptimisticUserMessage({
+      id: "optimistic-create-user",
+      text: "Analyze this",
+      timestamp: new Date(1000),
+      images: [image],
+      attachments: [attachment],
+    });
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentTail: [optimistic],
+      payload: {
+        ...baseTimelineInput.payload,
+        reset: true,
+        startCursor: { seq: 1 },
+        endCursor: { seq: 1 },
+        entries: [
+          {
+            ...makeTimelineEntry(1, "Analyze this", "user_message"),
+            item: {
+              type: "user_message",
+              text: "Analyze this",
+              messageId: "canonical-create-user",
+            },
+          },
+        ],
+      },
+    });
+
+    const userMessages = result.tail.filter((item) => item.kind === "user_message");
+    expect(userMessages).toHaveLength(1);
+    expect(userMessages[0]).toMatchObject({
+      id: "canonical-create-user",
+      text: "Analyze this",
+      images: [image],
+      attachments: [attachment],
+    });
+    expect(userMessages[0]?.optimistic).toBeUndefined();
+  });
+
+  it("keeps an unmatched optimistic user message during tail replacement", () => {
+    const optimistic = makeOptimisticUserMessage("still sending", "optimistic-unmatched");
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentTail: [optimistic],
+      payload: {
+        ...baseTimelineInput.payload,
+        reset: true,
+        entries: [],
+      },
+    });
+
+    expect(result.tail).toEqual([optimistic]);
+  });
+
  it("sets cursor to null when reset=true but no cursors in payload", () => {
    const existingTail: StreamItem[] = [
      {
@@ -363,6 +449,40 @@ describe("processTimelineResponse", () => {
    expect(result.error).toBe(null);
  });

+  it("reconciles an optimistic user message during an after-page response", () => {
+    const existingCursor: TimelineCursor = {
+      epoch: "epoch-1",
+      startSeq: 1,
+      endSeq: 1,
+    };
+    const optimistic = makeOptimisticUserMessage("sent while catching up", "optimistic-after");
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentTail: [optimistic],
+      currentCursor: existingCursor,
+      payload: {
+        ...baseTimelineInput.payload,
+        epoch: "epoch-1",
+        entries: [
+          {
+            ...makeTimelineEntry(2, "sent while catching up", "user_message"),
+            item: {
+              type: "user_message",
+              text: "sent while catching up",
+              messageId: "canonical-after",
+            },
+          },
+        ],
+      },
+    });
+
+    const userMessages = result.tail.filter((item) => item.kind === "user_message");
+    expect(userMessages).toHaveLength(1);
+    expect(userMessages[0]?.id).toBe("canonical-after");
+    expect(userMessages[0]?.optimistic).toBeUndefined();
+  });
+
  it("keeps an active assistant head live when an incremental fetch accepts same-turn assistant text", () => {
    const existingCursor: TimelineCursor = {
      epoch: "epoch-1",
@@ -483,6 +603,43 @@ describe("processTimelineResponse", () => {
    });
  });

+  it("does not reconcile an active optimistic user message from a before-page response", () => {
+    const optimistic = makeOptimisticUserMessage("active prompt", "optimistic-active");
+    const existingCursor: TimelineCursor = {
+      epoch: "epoch-1",
+      startSeq: 3,
+      endSeq: 5,
+    };
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentTail: [optimistic],
+      currentCursor: existingCursor,
+      payload: {
+        ...baseTimelineInput.payload,
+        direction: "before",
+        epoch: "epoch-1",
+        startCursor: { seq: 1 },
+        endCursor: { seq: 2 },
+        entries: [
+          {
+            ...makeTimelineEntry(1, "older prompt", "user_message"),
+            item: {
+              type: "user_message",
+              text: "older prompt",
+              messageId: "canonical-before",
+            },
+          },
+        ],
+      },
+    });
+
+    const userMessages = result.tail.filter((item) => item.kind === "user_message");
+    expect(userMessages).toHaveLength(2);
+    expect(userMessages.map((item) => item.id)).toEqual(["canonical-before", "optimistic-active"]);
+    expect(userMessages[1]?.optimistic).toBe(true);
+  });
+
  it("leaves the cursor alone when a before page makes no progress", () => {
    const currentTail: StreamItem[] = [
      {