Allow failed agent turns to be forked (#2063)

* fix(chat): allow failed turns to be forked Failed turns can end in Paseo-generated timeline items without provider message IDs. Anchor forks to canonical timeline positions while retaining legacy boundaries for compatibility. Place the fork summary before the new prompt and delimit it as structured chat history. * fix(chat): tighten fork boundary handling
2026-07-15 16:30:42 +02:00
parent f06792ae89
commit 04e893417e
23 changed files with 549 additions and 108 deletions
@@ -1,6 +1,9 @@
 import { describe, expect, it } from "vitest";
 import type { StreamItem } from "@/types/stream";
-import { resolveAssistantTurnBoundaryMessageId } from "./turn-boundary";
+import {
+  resolveAssistantTurnBoundaryMessageId,
+  resolveAssistantTurnForkBoundary,
+} from "./turn-boundary";

 function timestamp(seed: number): Date {
  return new Date(`2026-01-01T00:00:${seed.toString().padStart(2, "0")}.000Z`);
@@ -62,3 +65,65 @@ describe("resolveAssistantTurnBoundaryMessageId", () => {
    ).toBeUndefined();
  });
 });
+
+describe("resolveAssistantTurnForkBoundary", () => {
+  it("forks a failed assistant turn from its Paseo timeline cursor without a provider message id", () => {
+    const failedTurn = {
+      ...assistantMessage("assistant-error", 2),
+      timelineCursor: { epoch: "timeline-1", seq: 42 },
+    };
+
+    expect(
+      resolveAssistantTurnForkBoundary({
+        items: [userMessage("user-1", 1), failedTurn],
+        startIndex: 1,
+        supportsTimelineCursor: true,
+      }),
+    ).toEqual({
+      boundaryCursor: { epoch: "timeline-1", seq: 42 },
+    });
+  });
+
+  it("includes the provider message id with a supported timeline cursor", () => {
+    const selected = {
+      ...assistantMessage("assistant-1", 2, "msg-assistant-1"),
+      timelineCursor: { epoch: "timeline-1", seq: 42 },
+    };
+
+    expect(
+      resolveAssistantTurnForkBoundary({
+        items: [selected],
+        startIndex: 0,
+        supportsTimelineCursor: true,
+      }),
+    ).toEqual({
+      boundaryCursor: { epoch: "timeline-1", seq: 42 },
+      boundaryMessageId: "msg-assistant-1",
+    });
+  });
+
+  it("falls back to the provider message id when timeline cursors are unsupported", () => {
+    const selected = {
+      ...assistantMessage("assistant-1", 2, "msg-assistant-1"),
+      timelineCursor: { epoch: "timeline-1", seq: 42 },
+    };
+
+    expect(
+      resolveAssistantTurnForkBoundary({
+        items: [selected],
+        startIndex: 0,
+        supportsTimelineCursor: false,
+      }),
+    ).toEqual({ boundaryMessageId: "msg-assistant-1" });
+  });
+
+  it("does not offer an unavailable boundary", () => {
+    expect(
+      resolveAssistantTurnForkBoundary({
+        items: [assistantMessage("assistant-1", 2)],
+        startIndex: 0,
+        supportsTimelineCursor: false,
+      }),
+    ).toBeUndefined();
+  });
+});
@@ -1,4 +1,8 @@
-import type { StreamItem } from "@/types/stream";
+import type { StreamItem, TimelinePosition } from "@/types/stream";
+
+export type AssistantTurnForkBoundary =
+  | { boundaryCursor: TimelinePosition; boundaryMessageId?: string }
+  | { boundaryCursor?: undefined; boundaryMessageId: string };

 export function resolveAssistantTurnBoundaryMessageId(input: {
  items: readonly StreamItem[];
@@ -11,3 +15,21 @@ export function resolveAssistantTurnBoundaryMessageId(input: {
  // Forking without the selected assistant's durable message id would send the wrong slice.
  return item.messageId || undefined;
 }
+
+export function resolveAssistantTurnForkBoundary(input: {
+  items: readonly StreamItem[];
+  startIndex: number;
+  supportsTimelineCursor: boolean;
+}): AssistantTurnForkBoundary | undefined {
+  const item = input.items[input.startIndex];
+  if (item?.kind !== "assistant_message") {
+    return undefined;
+  }
+  if (input.supportsTimelineCursor && item.timelineCursor) {
+    return {
+      boundaryCursor: item.timelineCursor,
+      ...(item.messageId ? { boundaryMessageId: item.messageId } : {}),
+    };
+  }
+  return item.messageId ? { boundaryMessageId: item.messageId } : undefined;
+}
@@ -9,7 +9,7 @@ import {
  collectAssistantTurnContentForStreamRenderStrategy,
  type StreamStrategy,
 } from "./strategy";
-import { resolveAssistantTurnBoundaryMessageId } from "./turn-boundary";
+import { resolveAssistantTurnForkBoundary, type AssistantTurnForkBoundary } from "./turn-boundary";
 import {
  AssistantTurnFooter,
  LiveElapsed,
@@ -31,7 +31,7 @@ const workingIndicatorColorMapping = (theme: Theme) => ({
 export type TurnContentStrategy = StreamStrategy;
 export type AssistantTurnForkHandler = (input: {
  target: AssistantForkTarget;
-  boundaryMessageId?: string;
+  boundary: AssistantTurnForkBoundary;
 }) => Promise<void> | void;

 export const TurnFooter = memo(function TurnFooter({
@@ -39,12 +39,14 @@ export const TurnFooter = memo(function TurnFooter({
  inFlightTurnStartedAt,
  host,
  strategy,
+  supportsTimelineCursor,
  onForkAssistantTurn,
 }: {
  isRunning: boolean;
  inFlightTurnStartedAt: Date | null;
  host: TurnFooterHost | null;
  strategy: TurnContentStrategy;
+  supportsTimelineCursor: boolean;
  onForkAssistantTurn?: AssistantTurnForkHandler;
 }) {
  if (isRunning) {
@@ -63,6 +65,7 @@ export const TurnFooter = memo(function TurnFooter({
      items={host.items}
      timing={host.timing}
      startIndex={host.startIndex}
+      supportsTimelineCursor={supportsTimelineCursor}
      onForkAssistantTurn={onForkAssistantTurn}
    />
  );
@@ -73,12 +76,14 @@ export const CompletedTurnFooterRow = memo(function CompletedTurnFooterRow({
  items,
  timing,
  startIndex,
+  supportsTimelineCursor,
  onForkAssistantTurn,
 }: {
  strategy: TurnContentStrategy;
  items: StreamItem[];
  timing?: TurnTiming;
  startIndex: number;
+  supportsTimelineCursor: boolean;
  onForkAssistantTurn?: AssistantTurnForkHandler;
 }) {
  return (
@@ -88,6 +93,7 @@ export const CompletedTurnFooterRow = memo(function CompletedTurnFooterRow({
        items={items}
        timing={timing}
        startIndex={startIndex}
+        supportsTimelineCursor={supportsTimelineCursor}
        onForkAssistantTurn={onForkAssistantTurn}
      />
    </TurnFooterRow>
@@ -130,12 +136,14 @@ function CompletedTurnFooter({
  items,
  timing,
  startIndex,
+  supportsTimelineCursor,
  onForkAssistantTurn,
 }: {
  strategy: TurnContentStrategy;
  items: StreamItem[];
  timing?: TurnTiming;
  startIndex: number;
+  supportsTimelineCursor: boolean;
  onForkAssistantTurn?: AssistantTurnForkHandler;
 }) {
  const getContent = useCallback(
@@ -147,18 +155,27 @@ function CompletedTurnFooter({
      }),
    [strategy, items, startIndex],
  );
-  const boundaryMessageId = resolveAssistantTurnBoundaryMessageId({
+  const boundary = resolveAssistantTurnForkBoundary({
    items,
    startIndex,
+    supportsTimelineCursor,
  });
+  const handleFork = useCallback(
+    (target: AssistantForkTarget) => {
+      if (!boundary) {
+        return;
+      }
+      return onForkAssistantTurn?.({ target, boundary });
+    },
+    [boundary, onForkAssistantTurn],
+  );
  return (
    <View style={stylesheet.turnFooterSlot}>
      <AssistantTurnFooter
        getContent={getContent}
        completedAt={timing?.completedAt}
        durationMs={timing?.durationMs}
-        forkBoundaryMessageId={boundaryMessageId}
-        onFork={onForkAssistantTurn}
+        onFork={boundary && onForkAssistantTurn ? handleFork : undefined}
      />
    </View>
  );
@@ -142,6 +142,7 @@ function renderStreamItemWithTurnFooter(input: {
  content: ReactNode;
  layoutItem: StreamLayoutItem;
  strategy: TurnContentStrategy;
+  supportsTimelineCursor: boolean;
  onForkAssistantTurn?: AssistantTurnForkHandler;
 }): ReactNode {
  if (!input.content) {
@@ -155,6 +156,7 @@ function renderStreamItemWithTurnFooter(input: {
      items={footerHost.items}
      timing={footerHost.timing}
      startIndex={footerHost.startIndex}
+      supportsTimelineCursor={input.supportsTimelineCursor}
      onForkAssistantTurn={input.onForkAssistantTurn}
    />
  ) : null;
@@ -282,6 +284,7 @@ function buildChatHistoryAttachment(input: {
      serverId: input.serverId,
      agentId: input.agentId,
      boundaryMessageId: input.payload.boundaryMessageId,
+      boundaryCursor: input.payload.boundaryCursor,
      itemCount: input.payload.itemCount,
    },
  };
@@ -369,6 +372,10 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
        !readOnly &&
        state.sessions[resolvedServerId]?.serverInfo?.features?.agentForkContext === true,
    );
+    const supportsAgentForkContextCursor = useSessionStore(
+      (state) =>
+        state.sessions[resolvedServerId]?.serverInfo?.features?.agentForkContextCursor === true,
+    );

    const workspaceRoot = context.cwd?.trim() || "";
    const { requestDirectoryListing } = useFileExplorerActions({
@@ -466,7 +473,7 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
    });

    const handleForkAssistantTurn: AssistantTurnForkHandler = useStableEvent(
-      async ({ target, boundaryMessageId }) => {
+      async ({ target, boundary }) => {
        try {
          if (!supportsAgentForkContext) {
            toast?.error(t("message.actions.forkUnavailable"));
@@ -478,10 +485,7 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
          const draftSetup = buildForkDraftSetup(context);
          const prepareForkDraft = async () => {
            const draftId = generateDraftId();
-            const payload = await client.buildAgentForkContext(
-              agentId,
-              boundaryMessageId ? { boundaryMessageId } : {},
-            );
+            const payload = await client.buildAgentForkContext(agentId, boundary);
            const attachment = buildChatHistoryAttachment({
              draftId,
              serverId: resolvedServerId,
@@ -854,10 +858,17 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
          content,
          layoutItem,
          strategy: streamRenderStrategy,
+          supportsTimelineCursor: supportsAgentForkContextCursor,
          onForkAssistantTurn: readOnly ? undefined : handleForkAssistantTurn,
        });
      },
-      [handleForkAssistantTurn, readOnly, renderStreamItemContent, streamRenderStrategy],
+      [
+        handleForkAssistantTurn,
+        readOnly,
+        renderStreamItemContent,
+        streamRenderStrategy,
+        supportsAgentForkContextCursor,
+      ],
    );

    const pendingPermissionItems = useMemo(
@@ -882,6 +893,7 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
            inFlightTurnStartedAt={baseRenderModel.turnTiming.runningStartedAt}
            host={bottomTurnFooterHost}
            strategy={streamRenderStrategy}
+            supportsTimelineCursor={supportsAgentForkContextCursor}
            onForkAssistantTurn={readOnly ? undefined : handleForkAssistantTurn}
          />
        ) : null,
@@ -892,6 +904,7 @@ const AgentStreamViewComponent = forwardRef<AgentStreamViewHandle, AgentStreamVi
        baseRenderModel.turnTiming.runningStartedAt,
        bottomTurnFooterHost,
        streamRenderStrategy,
+        supportsAgentForkContextCursor,
      ],
    );
    const renderModel = useMemo<AgentStreamRenderModel>(() => {