Rewind chat or files from any user message (#1154)

* Add Pi extension-launch plumbing for in-process tree primitives * Add Rewind controls for agent sessions * Preserve SDK checkpoint env var so Claude file rewind works The SDK injects CLAUDE_CODE_ENABLE_SDK_FILE_CHECKPOINTING=true when file checkpointing is enabled. The earlier parent-env stripping from c4782ec71c was too broad and removed that child flag along with real parent-session markers. Keep stripping parent markers, but preserve the checkpoint flag so the SDK rewind path can run directly. With that boundary fixed, the JSONL file-rewind fallback is no longer needed. * Update OpenCode history replay tests for rewind state * Show Rewind actions in an adaptive sheet * Fix rewind semantics across providers * Refactor Pi IDs through extension capture * Restore user message to composer after rewind * Fix Claude rewind tool result mapping * Fix Claude rewind session switch roundtrip * Fix rewind and optimistic message regressions * Tighten rewind test discipline * Fix OpenCode optimistic message reconciliation * Reconcile optimistic user messages by marker * Reorder user message actions * Unify provider user message contract * Add Claude rewind flow Playwright contract * Add Codex rewind flow Playwright contract * Add OpenCode rewind flow Playwright contract * Add Pi rewind flow Playwright contract * Remove real-provider rewind Playwright contracts * Gate real-provider rewind specs from CI Playwright
2026-05-25 17:11:34 +08:00
parent e966f70322
commit 132e572d6e
99 changed files with 7106 additions and 3425 deletions
@@ -353,6 +353,7 @@ describe("dispatchComposerAgentMessage", () => {
    expect(userMessage.images).toEqual([image]);
    expect(userMessage.attachments).toEqual(call.options.attachments);
    expect(userMessage.id).toBe(call.options.messageId);
+    expect(userMessage.optimistic).toBe(true);
  });

  it("appends to the existing head when one is present", async () => {
@@ -135,6 +135,7 @@ export async function dispatchComposerAgentMessage(
    id: messageId,
    text: input.text,
    timestamp: new Date(),
+    optimistic: true,
    ...(wirePayload.images.length > 0 ? { images: wirePayload.images } : {}),
    ...(wirePayload.attachments.length > 0 ? { attachments: wirePayload.attachments } : {}),
  };
@@ -102,15 +102,23 @@ import { useAttachmentPreviewUrl } from "@/attachments/use-attachment-preview-ur
 import { persistAttachmentFromBytes, persistAttachmentFromDataUrl } from "@/attachments/service";
 import type { DaemonClient } from "@server/client/daemon-client";
 import { isWeb, isNative } from "@/constants/platform";
+import type { AgentCapabilityFlags } from "@server/server/agent/agent-sdk-types";
+import { RewindMenu, type RewindMode } from "@/components/rewind/rewind-menu";
+import { useRewindAgentMutation } from "@/components/rewind/use-rewind-agent-mutation";
 export type { InlinePathTarget } from "@/assistant-file-links";

 type MarkdownStyles = Record<string, TextStyle & ViewStyle & { [key: string]: unknown }>;

 interface UserMessageProps {
+  serverId?: string;
+  agentId?: string;
+  messageId?: string;
  message: string;
  images?: UserMessageImageAttachment[];
  attachments?: AgentAttachment[];
  timestamp: number;
+  capabilities?: AgentCapabilityFlags;
+  client?: DaemonClient | null;
  isFirstInGroup?: boolean;
  isLastInGroup?: boolean;
  disableOuterSpacing?: boolean;
@@ -450,10 +458,15 @@ function getUserMessageAttachmentLabel(attachment: AgentAttachment): string {
 }

 export const UserMessage = memo(function UserMessage({
+  serverId,
+  agentId,
+  messageId,
  message,
  images = [],
  attachments = [],
  timestamp,
+  capabilities,
+  client,
  isFirstInGroup = true,
  isLastInGroup = true,
  disableOuterSpacing,
@@ -469,10 +482,17 @@ export const UserMessage = memo(function UserMessage({
    () => formatMessageTimestamp(new Date(timestamp)),
    [timestamp],
  );
+  const rewindMutation = useRewindAgentMutation({ serverId, agentId, client, messageId });

  const handlePointerEnter = useCallback(() => setIsHovered(true), []);
  const handlePointerLeave = useCallback(() => setIsHovered(false), []);
  const getMessageContent = useCallback(() => message, [message]);
+  const handleRewind = useCallback(
+    (input: { mode: RewindMode; rewoundText: string }) => {
+      return rewindMutation.rewindAgent(input);
+    },
+    [rewindMutation],
+  );

  const containerStyle = useMemo(
    () => [
@@ -510,7 +530,7 @@ export const UserMessage = memo(function UserMessage({
  );

  return (
-    <View style={containerStyle}>
+    <View style={containerStyle} testID="user-message">
      <View
        style={userMessageStylesheet.content}
        onPointerEnter={handlePointerEnter}
@@ -549,6 +569,14 @@ export const UserMessage = memo(function UserMessage({
        {hasText ? (
          <View style={trailingRowStyle} pointerEvents={showTrailingRow ? "auto" : "none"}>
            <Text style={userMessageStylesheet.timestampText}>{formattedTimestamp}</Text>
+            {capabilities ? (
+              <RewindMenu
+                capabilities={capabilities}
+                isPending={rewindMutation.isPending}
+                rewoundText={message}
+                onRewind={handleRewind}
+              />
+            ) : null}
            <TurnCopyButton
              getContent={getMessageContent}
              containerStyle={userMessageStylesheet.copyButton}
@@ -0,0 +1,31 @@
+import { describe, expect, test } from "vitest";
+import { restoreComposerTextIfEmpty } from "./composer-restore";
+import { shouldRestoreComposerForRewindMode } from "./rewind-mode";
+
+describe("restoreComposerTextIfEmpty", () => {
+  test("restores the rewound message when the composer is empty", () => {
+    expect(
+      restoreComposerTextIfEmpty({
+        currentText: "",
+        rewoundText: "message before rewind",
+      }),
+    ).toBe("message before rewind");
+  });
+
+  test("preserves an existing composer draft", () => {
+    expect(
+      restoreComposerTextIfEmpty({
+        currentText: "keep this draft",
+        rewoundText: "message before rewind",
+      }),
+    ).toBe("keep this draft");
+  });
+});
+
+describe("shouldRestoreComposerForRewindMode", () => {
+  test("restores only conversation-mutating rewind modes", () => {
+    expect(shouldRestoreComposerForRewindMode("conversation")).toBe(true);
+    expect(shouldRestoreComposerForRewindMode("files")).toBe(false);
+    expect(shouldRestoreComposerForRewindMode("both")).toBe(true);
+  });
+});
@@ -0,0 +1,68 @@
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  type ReactNode,
+} from "react";
+
+interface RewindComposerRestoreContextValue {
+  restoreTextIfComposerEmpty: (text: string) => void;
+}
+
+interface RewindComposerRestoreProviderProps {
+  text: string;
+  setText: (text: string) => void;
+  children: ReactNode;
+}
+
+const RewindComposerRestoreContext = createContext<RewindComposerRestoreContextValue | null>(null);
+
+export function restoreComposerTextIfEmpty(input: {
+  currentText: string;
+  rewoundText: string;
+}): string {
+  if (input.currentText.length > 0) {
+    return input.currentText;
+  }
+  return input.rewoundText;
+}
+
+export function RewindComposerRestoreProvider({
+  text,
+  setText,
+  children,
+}: RewindComposerRestoreProviderProps) {
+  const textRef = useRef(text);
+
+  useEffect(() => {
+    textRef.current = text;
+  }, [text]);
+
+  const restoreTextIfComposerEmpty = useCallback(
+    (rewoundText: string) => {
+      const nextText = restoreComposerTextIfEmpty({
+        currentText: textRef.current,
+        rewoundText,
+      });
+      if (nextText !== textRef.current) {
+        setText(nextText);
+      }
+    },
+    [setText],
+  );
+
+  const value = useMemo(() => ({ restoreTextIfComposerEmpty }), [restoreTextIfComposerEmpty]);
+
+  return (
+    <RewindComposerRestoreContext.Provider value={value}>
+      {children}
+    </RewindComposerRestoreContext.Provider>
+  );
+}
+
+export function useRewindComposerRestore(): RewindComposerRestoreContextValue | null {
+  return useContext(RewindComposerRestoreContext);
+}
@@ -0,0 +1,164 @@
+import { memo, useCallback, useMemo, useState, type ComponentType } from "react";
+import { Text, View } from "react-native";
+import { FileText, Layers, MessageSquare, Undo2 } from "lucide-react-native";
+import { StyleSheet } from "react-native-unistyles";
+import { AdaptiveModalSheet, type SheetHeader } from "@/components/adaptive-modal-sheet";
+import { Button } from "@/components/ui/button";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+import {
+  type RewindMenuItem,
+  type RewindMode,
+  useRewindCapabilities,
+} from "./use-rewind-capabilities";
+import type { AgentCapabilityFlags } from "@server/server/agent/agent-sdk-types";
+
+export type { RewindMode };
+
+interface RewindMenuProps {
+  capabilities: AgentCapabilityFlags;
+  rewoundText: string;
+  onRewind: (input: { mode: RewindMode; rewoundText: string }) => Promise<void> | void;
+  isPending?: boolean;
+  testID?: string;
+}
+
+const REWIND_HEADER: SheetHeader = { title: "Rewind to this message" };
+const SNAP_POINTS: string[] = ["42%", "72%"];
+
+function iconForItem(item: RewindMenuItem): ComponentType<{ color: string; size: number }> {
+  switch (item.mode) {
+    case "conversation":
+      return MessageSquare;
+    case "files":
+      return FileText;
+    case "both":
+      return Layers;
+  }
+}
+
+export const RewindMenu = memo(function RewindMenu({
+  capabilities,
+  rewoundText,
+  onRewind,
+  isPending = false,
+  testID = "rewind-menu",
+}: RewindMenuProps) {
+  const items = useRewindCapabilities(capabilities);
+  const [isSheetVisible, setIsSheetVisible] = useState(false);
+  const [pendingMode, setPendingMode] = useState<RewindMode | null>(null);
+  const handleOpen = useCallback(() => setIsSheetVisible(true), []);
+  const handleClose = useCallback(() => {
+    if (!isPending) {
+      setIsSheetVisible(false);
+    }
+  }, [isPending]);
+  const handleSelect = useCallback(
+    (mode: RewindMode) => async () => {
+      if (isPending) {
+        return;
+      }
+      setPendingMode(mode);
+      try {
+        await onRewind({ mode, rewoundText });
+      } catch {
+        // useRewindAgentMutation owns the toast; the sheet only owns flow state.
+      } finally {
+        setPendingMode(null);
+        setIsSheetVisible(false);
+      }
+    },
+    [isPending, onRewind, rewoundText],
+  );
+  const triggerIcon = useCallback((color: string) => <Undo2 size={16} color={color} />, []);
+  const tooltipContent = useMemo(
+    () => (
+      <TooltipContent side="top" align="center" offset={8}>
+        <Text style={styles.tooltipText}>Rewind to this message</Text>
+      </TooltipContent>
+    ),
+    [],
+  );
+
+  if (items.length === 0) {
+    return null;
+  }
+
+  return (
+    <>
+      <Tooltip delayDuration={0} enabledOnDesktop enabledOnMobile={false}>
+        <TooltipTrigger asChild>
+          <Button
+            accessibilityLabel="Rewind to this message"
+            disabled={isPending}
+            leftIcon={triggerIcon}
+            onPress={handleOpen}
+            size="xs"
+            style={styles.trigger}
+            testID={`${testID}-trigger`}
+            variant="ghost"
+          />
+        </TooltipTrigger>
+        {tooltipContent}
+      </Tooltip>
+      <AdaptiveModalSheet
+        desktopMaxWidth={420}
+        header={REWIND_HEADER}
+        onClose={handleClose}
+        scrollable={false}
+        snapPoints={SNAP_POINTS}
+        testID={`${testID}-content`}
+        visible={isSheetVisible}
+      >
+        <View style={styles.content}>
+          <Text style={styles.warningText}>This action cannot be undone</Text>
+          <View style={styles.actions}>
+            {items.map((item) => (
+              <Button
+                key={item.mode}
+                leftIcon={iconForItem(item)}
+                loading={pendingMode === item.mode}
+                disabled={isPending}
+                onPress={handleSelect(item.mode)}
+                size="md"
+                style={styles.actionButton}
+                testID={item.testID}
+                variant="secondary"
+              >
+                {item.label}
+              </Button>
+            ))}
+          </View>
+        </View>
+      </AdaptiveModalSheet>
+    </>
+  );
+});
+
+const styles = StyleSheet.create((theme) => ({
+  trigger: {
+    width: 24,
+    height: 24,
+    minHeight: 24,
+    paddingHorizontal: 0,
+    paddingVertical: 0,
+    alignItems: "center",
+    justifyContent: "center",
+  },
+  tooltipText: {
+    color: theme.colors.foreground,
+    fontSize: theme.fontSize.xs,
+  },
+  content: {
+    gap: theme.spacing[4],
+  },
+  warningText: {
+    color: theme.colors.statusWarning,
+    fontSize: theme.fontSize.sm,
+  },
+  actions: {
+    gap: theme.spacing[2],
+  },
+  actionButton: {
+    justifyContent: "flex-start",
+  },
+}));
@@ -0,0 +1,5 @@
+import type { RewindMode } from "./use-rewind-capabilities";
+
+export function shouldRestoreComposerForRewindMode(mode: RewindMode): boolean {
+  return mode === "conversation" || mode === "both";
+}
@@ -0,0 +1,80 @@
+import { useCallback } from "react";
+import { useMutation } from "@tanstack/react-query";
+import { useToast } from "@/contexts/toast-context";
+import type { DaemonClient } from "@server/client/daemon-client";
+import type { RewindMode } from "./use-rewind-capabilities";
+import { useRewindComposerRestore } from "./composer-restore";
+import { useSessionStore } from "@/stores/session-store";
+import { shouldRestoreComposerForRewindMode } from "./rewind-mode";
+import { clearOptimisticUserMessages } from "@/types/stream";
+
+interface UseRewindAgentMutationInput {
+  serverId?: string;
+  agentId?: string;
+  messageId?: string;
+  client?: DaemonClient | null;
+}
+
+interface RewindAgentInput {
+  mode: RewindMode;
+  rewoundText: string;
+}
+
+export function useRewindAgentMutation(input: UseRewindAgentMutationInput): {
+  rewindAgent: (input: RewindAgentInput) => Promise<void>;
+  isPending: boolean;
+} {
+  const toast = useToast();
+  const composerRestore = useRewindComposerRestore();
+  const { isPending, mutateAsync } = useMutation({
+    mutationFn: async ({ mode }: RewindAgentInput) => {
+      if (!input.client || !input.agentId || !input.messageId) {
+        throw new Error("Daemon client not available");
+      }
+      await input.client.rewindAgent(input.agentId, input.messageId, mode);
+      if (mode !== "files") {
+        if (input.serverId) {
+          const session = useSessionStore.getState().sessions[input.serverId];
+          useSessionStore.getState().setAgentStreamState(input.serverId, input.agentId, {
+            tail: clearOptimisticUserMessages(session?.agentStreamTail.get(input.agentId) ?? []),
+            head: clearOptimisticUserMessages(session?.agentStreamHead.get(input.agentId) ?? []),
+          });
+        }
+        const cursor = input.serverId
+          ? useSessionStore
+              .getState()
+              .sessions[input.serverId]?.agentTimelineCursor.get(input.agentId)
+          : undefined;
+        await input.client.fetchAgentTimeline(input.agentId, {
+          direction: "tail",
+          projection: "canonical",
+          ...(cursor ? { cursor: { epoch: cursor.epoch, seq: cursor.endSeq } } : {}),
+        });
+      }
+    },
+    onSuccess: (_data, variables) => {
+      if (!shouldRestoreComposerForRewindMode(variables.mode)) {
+        return;
+      }
+      composerRestore?.restoreTextIfComposerEmpty(variables.rewoundText);
+    },
+    onError: (error) => {
+      toast.error(error instanceof Error ? error.message : "Failed to rewind agent");
+    },
+  });
+
+  const rewindAgent = useCallback(
+    async (rewindInput: RewindAgentInput) => {
+      if (isPending) {
+        return;
+      }
+      await mutateAsync(rewindInput);
+    },
+    [isPending, mutateAsync],
+  );
+
+  return {
+    rewindAgent,
+    isPending,
+  };
+}
@@ -0,0 +1,36 @@
+import { describe, expect, test } from "vitest";
+
+import { resolveRewindMenuItems } from "./use-rewind-capabilities";
+
+describe("resolveRewindMenuItems", () => {
+  test("returns no items when the provider declares no rewind capability", () => {
+    expect(
+      resolveRewindMenuItems({
+        supportsRewindConversation: false,
+        supportsRewindFiles: false,
+        supportsRewindBoth: false,
+      }),
+    ).toEqual([]);
+  });
+
+  test("returns only the capabilities declared by the provider", () => {
+    expect(
+      resolveRewindMenuItems({
+        supportsRewindConversation: true,
+        supportsRewindFiles: false,
+        supportsRewindBoth: true,
+      }),
+    ).toEqual([
+      {
+        mode: "conversation",
+        label: "Rewind conversation",
+        testID: "rewind-menu-conversation",
+      },
+      {
+        mode: "both",
+        label: "Rewind conversation and files",
+        testID: "rewind-menu-both",
+      },
+    ]);
+  });
+});
@@ -0,0 +1,53 @@
+import { useMemo } from "react";
+import type { AgentCapabilityFlags } from "@server/server/agent/agent-sdk-types";
+
+export type RewindMode = "conversation" | "files" | "both";
+
+export interface RewindMenuItem {
+  mode: RewindMode;
+  label: string;
+  testID: string;
+}
+
+export function resolveRewindMenuItems(
+  capabilities:
+    | Pick<
+        AgentCapabilityFlags,
+        "supportsRewindConversation" | "supportsRewindFiles" | "supportsRewindBoth"
+      >
+    | null
+    | undefined,
+): RewindMenuItem[] {
+  if (!capabilities) {
+    return [];
+  }
+  const items: RewindMenuItem[] = [];
+  if (capabilities.supportsRewindConversation) {
+    items.push({
+      mode: "conversation",
+      label: "Rewind conversation",
+      testID: "rewind-menu-conversation",
+    });
+  }
+  if (capabilities.supportsRewindFiles) {
+    items.push({
+      mode: "files",
+      label: "Rewind files",
+      testID: "rewind-menu-files",
+    });
+  }
+  if (capabilities.supportsRewindBoth) {
+    items.push({
+      mode: "both",
+      label: "Rewind conversation and files",
+      testID: "rewind-menu-both",
+    });
+  }
+  return items;
+}
+
+export function useRewindCapabilities(
+  capabilities: Parameters<typeof resolveRewindMenuItems>[0],
+): RewindMenuItem[] {
+  return useMemo(() => resolveRewindMenuItems(capabilities), [capabilities]);
+}