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[] = [
      {
@@ -2,7 +2,7 @@ import type { AgentStreamEventPayload } from "@server/shared/messages";
 import type { AgentLifecycleStatus } from "@server/shared/agent-lifecycle";
 import type { Agent } from "@/stores/session-store";
 import { useSessionStore } from "@/stores/session-store";
-import type { StreamItem } from "@/types/stream";
+import type { StreamItem, UserMessageItem } from "@/types/stream";
 import {
  applyStreamEvent,
  hydrateStreamState,
@@ -240,15 +240,22 @@ function applyTimelineReplacePath(args: {
  timelineUnits: TimelineUnit[];
  payload: ProcessTimelineResponseInput["payload"];
  bootstrapPolicy: ReturnType<typeof deriveBootstrapTailTimelinePolicy>;
+  currentTail: StreamItem[];
  currentHead: StreamItem[];
  toHydratedEvents: (
    units: TimelineUnit[],
  ) => Array<{ event: AgentStreamEventPayload; timestamp: Date }>;
 }): TimelinePathResult {
-  const { timelineUnits, payload, bootstrapPolicy, currentHead, toHydratedEvents } = args;
+  const { timelineUnits, payload, bootstrapPolicy, currentTail, currentHead, toHydratedEvents } =
+    args;
  const hydratedTail = hydrateStreamState(toHydratedEvents(timelineUnits), { source: "canonical" });
+  const reconciledTail = reconcileOptimisticUsersAfterReplace({
+    canonicalTail: hydratedTail,
+    previousTail: currentTail,
+    previousHead: currentHead,
+  });
  const { tail, head } = preserveReplacePathAssistantHead({
-    tail: hydratedTail,
+    tail: reconciledTail,
    currentHead,
  });
  const cursor: TimelineCursor | null =
@@ -266,6 +273,93 @@ function applyTimelineReplacePath(args: {
  return { tail, head, cursor, cursorChanged: true, sideEffects };
 }

+function collectOptimisticUserMessages(items: StreamItem[]): Array<{
+  ordinal: number;
+  item: UserMessageItem;
+}> {
+  const optimistic: Array<{ ordinal: number; item: UserMessageItem }> = [];
+  let ordinal = 0;
+  for (const item of items) {
+    if (item.kind !== "user_message") {
+      continue;
+    }
+    if (item.optimistic) {
+      optimistic.push({ ordinal, item });
+    }
+    ordinal += 1;
+  }
+  return optimistic;
+}
+
+function mergeCanonicalUserWithOptimistic(
+  canonical: UserMessageItem,
+  optimistic: UserMessageItem,
+): UserMessageItem {
+  const shouldPreserveImages =
+    (!canonical.images || canonical.images.length === 0) &&
+    Boolean(optimistic.images && optimistic.images.length > 0);
+  const shouldPreserveAttachments =
+    (!canonical.attachments || canonical.attachments.length === 0) &&
+    Boolean(optimistic.attachments && optimistic.attachments.length > 0);
+  return {
+    ...canonical,
+    ...(shouldPreserveImages ? { images: optimistic.images } : {}),
+    ...(shouldPreserveAttachments ? { attachments: optimistic.attachments } : {}),
+  };
+}
+
+function reconcileOptimisticUsersAfterReplace(params: {
+  canonicalTail: StreamItem[];
+  previousTail: StreamItem[];
+  previousHead: StreamItem[];
+}): StreamItem[] {
+  const optimisticUsers = collectOptimisticUserMessages([
+    ...params.previousTail,
+    ...params.previousHead,
+  ]);
+  if (optimisticUsers.length === 0) {
+    return params.canonicalTail;
+  }
+
+  const canonicalUserIndexes: number[] = [];
+  params.canonicalTail.forEach((item, index) => {
+    if (item.kind === "user_message") {
+      canonicalUserIndexes.push(index);
+    }
+  });
+
+  let changed = false;
+  const nextTail = [...params.canonicalTail];
+  let searchFromOrdinal = 0;
+  const unmatched: UserMessageItem[] = [];
+
+  for (const optimistic of optimisticUsers) {
+    const canonicalOrdinal = canonicalUserIndexes.findIndex(
+      (_index, ordinal) => ordinal >= Math.max(optimistic.ordinal, searchFromOrdinal),
+    );
+    if (canonicalOrdinal < 0) {
+      unmatched.push(optimistic.item);
+      continue;
+    }
+
+    const canonicalIndex = canonicalUserIndexes[canonicalOrdinal];
+    const canonicalItem = canonicalIndex !== undefined ? nextTail[canonicalIndex] : undefined;
+    if (!canonicalItem || canonicalItem.kind !== "user_message") {
+      unmatched.push(optimistic.item);
+      continue;
+    }
+    nextTail[canonicalIndex] = mergeCanonicalUserWithOptimistic(canonicalItem, optimistic.item);
+    searchFromOrdinal = canonicalOrdinal + 1;
+    changed = true;
+  }
+
+  if (unmatched.length === 0) {
+    return changed ? nextTail : params.canonicalTail;
+  }
+
+  return [...nextTail, ...unmatched];
+}
+
 interface IncrementalAcceptResult {
  acceptedUnits: TimelineUnit[];
  cursor: TimelineCursor | undefined;
@@ -531,6 +625,7 @@ export function processTimelineResponse(
        timelineUnits,
        payload,
        bootstrapPolicy,
+        currentTail,
        currentHead,
        toHydratedEvents,
      })