From 42cbec5bd8481337b5d272907317a17ee1becbf0 Mon Sep 17 00:00:00 2001 From: "paseo-bot[bot]" <266920839+paseo-bot[bot]@users.noreply.github.com> Date: Fri, 6 Nov 2026 06:18:06 +0000 Subject: [PATCH] fix(app): preserve inherited context through rewind and replay Restore fork history alongside the editable prompt after conversation rewind, including history-only prompts. Keep locally assigned submission IDs distinct from provider IDs so replay without that local ID can recover the snapshot. --- .../e2e/browser/assistant-fork-menu.spec.ts | 34 ++++++++++ .../app/e2e/support/helpers/chat-history.ts | 16 ++++- .../src/attachments/chat-history/index.tsx | 15 +++++ packages/app/src/attachments/types.ts | 2 +- .../attachments/workspace-attachment-utils.ts | 1 + packages/app/src/components/message.tsx | 22 +++++-- .../rewind/composer-restore.test.ts | 34 +++++++--- .../components/rewind/composer-restore.tsx | 34 +++++----- .../app/src/components/rewind/rewind-menu.tsx | 13 +++- .../rewind/use-rewind-agent-mutation.ts | 7 ++- .../src/composer/attachments/workspace.tsx | 1 + packages/app/src/panels/agent-panel.tsx | 2 + .../src/server/agent/agent-manager.test.ts | 62 +++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 5 +- 14 files changed, 211 insertions(+), 37 deletions(-) diff --git a/packages/app/e2e/browser/assistant-fork-menu.spec.ts b/packages/app/e2e/browser/assistant-fork-menu.spec.ts index 84308d5ab3..0ec214eb48 100644 --- a/packages/app/e2e/browser/assistant-fork-menu.spec.ts +++ b/packages/app/e2e/browser/assistant-fork-menu.spec.ts @@ -1,5 +1,6 @@ import { captureChatHistory, + rewindForkPrompt, closeChatHistory, followChatHistorySource, openChatHistory, @@ -88,6 +89,39 @@ test.describe("Assistant fork menu", () => { await expectActiveForkAgent(page, forkAgentId); }); + for (const prompt of ["", "Continue the fork"]) { + test(`retains inherited context when rewinding the first fork prompt (${prompt || "history only"})`, async ({ + page, + seedForkWorkspace, + }) => { + const initialPrompt = "emit 1 coalesced agent stream updates for rewindable fork history."; + const session = await seedForkWorkspace({ + repoPrefix: "fork-rewind-", + title: "Rewind source", + initialPrompt, + model: "ten-second-stream", + }); + await openAgentRoute(page, session); + await expectComposerVisible(page); + await awaitAssistantMessage(page); + await session.client.waitForFinish(session.agentId, 45_000); + await forkMostRecentAssistantTurnToNewTab(page); + await submitMessage(page, prompt); + const forkId = await getActiveForkAgentId(page); + await session.client.waitForFinish(forkId, 45_000); + await page.reload(); + await openChatHistory(page, "sent", [initialPrompt]); + await closeChatHistory(page); + await rewindForkPrompt(page); + await openChatHistory(page, "draft", [initialPrompt]); + await closeChatHistory(page); + await submitMessage(page, prompt); + await session.client.waitForFinish(forkId, 45_000); + await page.reload(); + await openChatHistory(page, "sent", [initialPrompt]); + }); + } + test("shows history sent from another connected client without reloading", async ({ page, seedForkWorkspace, diff --git a/packages/app/e2e/support/helpers/chat-history.ts b/packages/app/e2e/support/helpers/chat-history.ts index c1ec5ff64f..99b49ba6ac 100644 --- a/packages/app/e2e/support/helpers/chat-history.ts +++ b/packages/app/e2e/support/helpers/chat-history.ts @@ -13,7 +13,9 @@ export async function openChatHistory(page: Page, location: "draft" | "sent", co location === "draft" ? "composer-chat-history-attachment-pill" : "message-chat-history-attachment"; - await page.getByTestId(testId).click(); + const attachment = page.getByTestId(testId).filter({ visible: true }).last(); + await expect(attachment).toBeVisible({ timeout: 10_000 }); + await attachment.click(); for (const text of contents) { await expect(page.getByTestId("chat-history-content")).toContainText(text); } @@ -56,3 +58,15 @@ export async function sendChatHistoryFromAnotherClient( ], }); } + +export async function rewindForkPrompt(page: Page) { + const prompt = page.getByTestId("user-message").filter({ visible: true }).last(); + await prompt.hover(); + const trigger = prompt.getByRole("button", { name: "Rewind to this message", exact: true }); + await expect(trigger).toBeVisible({ timeout: 10_000 }); + await trigger.click(); + const option = page.getByRole("menuitem", { name: "Rewind conversation", exact: true }); + await expect(option).toBeInViewport({ ratio: 1 }); + await option.click(); + await expect(option).not.toBeVisible(); +} diff --git a/packages/app/src/attachments/chat-history/index.tsx b/packages/app/src/attachments/chat-history/index.tsx index 6237ffc5aa..75278e2824 100644 --- a/packages/app/src/attachments/chat-history/index.tsx +++ b/packages/app/src/attachments/chat-history/index.tsx @@ -1,3 +1,6 @@ +import { generateMessageId } from "@/types/stream"; +import { useWorkspaceAttachmentsStore } from "@/attachments/workspace-attachments-store"; +import type { ChatHistoryAttachment as HistoryAttachment } from "@getpaseo/protocol/messages"; import { extractChatHistoryBody } from "@getpaseo/protocol/chat-history"; import { useCallback, useMemo, useState } from "react"; import { useTranslation } from "react-i18next"; @@ -108,3 +111,15 @@ const styles = StyleSheet.create((theme) => ({ color: theme.colors.foreground, }, })); + +export function restoreChatHistoryAttachment(scopeKey: string, attachment: HistoryAttachment) { + useWorkspaceAttachmentsStore.getState().addWorkspaceAttachment({ + scopeKey, + attachment: { + kind: "chat_history", + id: generateMessageId(), + attachment, + source: attachment.chatHistorySource, + }, + }); +} diff --git a/packages/app/src/attachments/types.ts b/packages/app/src/attachments/types.ts index c32d430969..16ea4977f0 100644 --- a/packages/app/src/attachments/types.ts +++ b/packages/app/src/attachments/types.ts @@ -83,7 +83,7 @@ export interface ChatHistoryContextAttachment { kind: "chat_history"; id: string; attachment: Extract; - source: ChatHistorySource & { + source?: ChatHistorySource & { boundaryMessageId?: string | null; boundaryCursor?: { epoch: string; seq: number } | null; itemCount?: number; diff --git a/packages/app/src/attachments/workspace-attachment-utils.ts b/packages/app/src/attachments/workspace-attachment-utils.ts index 821fd4ee58..b1aa760503 100644 --- a/packages/app/src/attachments/workspace-attachment-utils.ts +++ b/packages/app/src/attachments/workspace-attachment-utils.ts @@ -62,6 +62,7 @@ export function workspaceAttachmentToSubmitAttachment( }; } if (attachment.kind === "chat_history") { + if (!attachment.source) return attachment.attachment; const { serverId, agentId, title } = attachment.source; return { ...attachment.attachment, diff --git a/packages/app/src/components/message.tsx b/packages/app/src/components/message.tsx index 4c0d29ac63..f84ce28e9d 100644 --- a/packages/app/src/components/message.tsx +++ b/packages/app/src/components/message.tsx @@ -34,7 +34,10 @@ import { useTargetAgentTitle } from "@/tool-calls/use-target-agent-title"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; import type { AgentMessage } from "@getpaseo/protocol/agent-message"; import type { AgentCapabilityFlags, ToolCallDetail } from "@getpaseo/protocol/agent-types"; -import type { AgentAttachment } from "@getpaseo/protocol/messages"; +import type { + AgentAttachment, + ChatHistoryAttachment as ChatHistorySnapshot, +} from "@getpaseo/protocol/messages"; import MaskedView from "@react-native-masked-view/masked-view"; import { Check, @@ -420,7 +423,12 @@ export const UserMessage = memo(function UserMessage({ const hasText = message.trim().length > 0; const hasImages = images.length > 0; const hasAttachments = attachments.length > 0; - const showTrailingRow = !isPending && hasText && (isCompact || isNative || isHovered); + const chatHistory = attachments.find( + (attachment): attachment is ChatHistorySnapshot => + attachment.type === "text" && attachment.contextKind === "chat_history", + ); + const showTrailingRow = + !isPending && (hasText || Boolean(chatHistory)) && (isCompact || isNative || isHovered); const formattedTimestamp = useMemo( () => formatMessageTimestamp(new Date(timestamp)), [timestamp], @@ -429,9 +437,12 @@ export const UserMessage = memo(function UserMessage({ const handlePointerEnter = useCallback(() => setIsHovered(true), []); const handlePointerLeave = useCallback(() => setIsHovered(false), []); - const getMessageContent = useCallback(() => message, [message]); + const getMessageContent = useCallback( + () => message || chatHistory?.text || "", + [message, chatHistory], + ); const handleRewind = useCallback( - (input: { mode: RewindMode; rewoundText: string }) => { + (input: { mode: RewindMode; rewoundText: string; chatHistory?: ChatHistorySnapshot }) => { return rewindMutation.rewindAgent(input); }, [rewindMutation], @@ -527,7 +538,7 @@ export const UserMessage = memo(function UserMessage({ ) : null} - {hasText ? ( + {hasText || chatHistory ? ( ) : null} diff --git a/packages/app/src/components/rewind/composer-restore.test.ts b/packages/app/src/components/rewind/composer-restore.test.ts index 8ee2c6f246..51c490857d 100644 --- a/packages/app/src/components/rewind/composer-restore.test.ts +++ b/packages/app/src/components/rewind/composer-restore.test.ts @@ -1,27 +1,45 @@ import { describe, expect, test } from "vitest"; -import { restoreComposerTextIfEmpty } from "./composer-restore"; +import { restoreComposerPromptIfEmpty } from "./composer-restore"; import { shouldRestoreComposerForRewindMode } from "./rewind-mode"; -describe("restoreComposerTextIfEmpty", () => { +describe("restoreComposerPromptIfEmpty", () => { test("restores the rewound message when the composer is empty", () => { expect( - restoreComposerTextIfEmpty({ + restoreComposerPromptIfEmpty({ currentText: "", - rewoundText: "message before rewind", + rewoundPrompt: { text: "message before rewind" }, }), - ).toBe("message before rewind"); + ).toEqual({ text: "message before rewind" }); }); test("preserves an existing composer draft", () => { expect( - restoreComposerTextIfEmpty({ + restoreComposerPromptIfEmpty({ currentText: "keep this draft", - rewoundText: "message before rewind", + rewoundPrompt: { text: "message before rewind" }, }), - ).toBe("keep this draft"); + ).toEqual({ text: "keep this draft" }); }); }); +test.each(["", "Continue here"])("restores inherited history with rewound prompt %s", (text) => { + const chatHistory = { + type: "text", + mimeType: "text/plain", + contextKind: "chat_history", + text: "Earlier context", + } as const; + expect( + restoreComposerPromptIfEmpty({ currentText: "", rewoundPrompt: { text, chatHistory } }), + ).toEqual({ text, chatHistory }); + expect( + restoreComposerPromptIfEmpty({ + currentText: "Existing draft", + rewoundPrompt: { text, chatHistory }, + }), + ).toEqual({ text: "Existing draft" }); +}); + describe("shouldRestoreComposerForRewindMode", () => { test("restores only conversation-mutating rewind modes", () => { expect(shouldRestoreComposerForRewindMode("conversation")).toBe(true); diff --git a/packages/app/src/components/rewind/composer-restore.tsx b/packages/app/src/components/rewind/composer-restore.tsx index 07e645d225..6b74a69ccc 100644 --- a/packages/app/src/components/rewind/composer-restore.tsx +++ b/packages/app/src/components/rewind/composer-restore.tsx @@ -1,11 +1,14 @@ +import type { ChatHistoryAttachment } from "@getpaseo/protocol/messages"; +import { restoreChatHistoryAttachment } from "@/attachments/chat-history"; import type { ComposerTextSource } from "@/composer/text-source"; import { createContext, useCallback, useContext, useMemo, type ReactNode } from "react"; interface RewindComposerRestoreContextValue { - completeRewind: (text: string) => void; + completeRewind: (prompt: { text: string; chatHistory?: ChatHistoryAttachment }) => void; } interface RewindComposerRestoreProviderProps { + attachmentScopeKey: string; textSource: ComposerTextSource; setText: (text: string) => void; onRewindComplete: () => void; @@ -14,35 +17,34 @@ interface RewindComposerRestoreProviderProps { const RewindComposerRestoreContext = createContext(null); -export function restoreComposerTextIfEmpty(input: { +export function restoreComposerPromptIfEmpty(input: { currentText: string; - rewoundText: string; -}): string { - if (input.currentText.length > 0) { - return input.currentText; - } - return input.rewoundText; + rewoundPrompt: { text: string; chatHistory?: ChatHistoryAttachment }; +}): { text: string; chatHistory?: ChatHistoryAttachment } { + if (input.currentText.length > 0) return { text: input.currentText }; + return input.rewoundPrompt; } export function RewindComposerRestoreProvider({ + attachmentScopeKey, textSource, setText, onRewindComplete, children, }: RewindComposerRestoreProviderProps) { const completeRewind = useCallback( - (rewoundText: string) => { + (prompt: { text: string; chatHistory?: ChatHistoryAttachment }) => { const currentText = textSource.getSnapshot(); - const nextText = restoreComposerTextIfEmpty({ - currentText: currentText, - rewoundText, - }); - if (nextText !== currentText) { - setText(nextText); + const restored = restoreComposerPromptIfEmpty({ currentText, rewoundPrompt: prompt }); + if (restored.chatHistory) { + restoreChatHistoryAttachment(attachmentScopeKey, restored.chatHistory); + } + if (restored.text !== currentText) { + setText(restored.text); } onRewindComplete(); }, - [onRewindComplete, setText, textSource], + [attachmentScopeKey, onRewindComplete, setText, textSource], ); const value = useMemo(() => ({ completeRewind }), [completeRewind]); diff --git a/packages/app/src/components/rewind/rewind-menu.tsx b/packages/app/src/components/rewind/rewind-menu.tsx index ca13d0ad2d..daca12cf80 100644 --- a/packages/app/src/components/rewind/rewind-menu.tsx +++ b/packages/app/src/components/rewind/rewind-menu.tsx @@ -1,3 +1,4 @@ +import type { ChatHistoryAttachment } from "@getpaseo/protocol/messages"; import { memo, useCallback, useMemo, useState, type ReactElement } from "react"; import { useTranslation } from "react-i18next"; import { Text, View } from "react-native"; @@ -19,7 +20,12 @@ export type { RewindMode }; interface RewindMenuProps { capabilities: AgentCapabilityFlags; rewoundText: string; - onRewind: (input: { mode: RewindMode; rewoundText: string }) => Promise | void; + chatHistory?: ChatHistoryAttachment; + onRewind: (input: { + mode: RewindMode; + rewoundText: string; + chatHistory?: ChatHistoryAttachment; + }) => Promise | void; isPending?: boolean; testID?: string; } @@ -38,6 +44,7 @@ function getIcon(mode: RewindMode, color: string): ReactElement { export const RewindMenu = memo(function RewindMenu({ capabilities, rewoundText, + chatHistory, onRewind, isPending: isPendingProp = false, testID = "rewind-menu", @@ -70,7 +77,7 @@ export const RewindMenu = memo(function RewindMenu({ if (isLocked) return; setPendingMode(mode); try { - await onRewind({ mode, rewoundText }); + await onRewind({ mode, rewoundText, chatHistory }); } catch { // useRewindAgentMutation owns the toast; the menu only owns flow state. } finally { @@ -78,7 +85,7 @@ export const RewindMenu = memo(function RewindMenu({ setIsOpen(false); } }, - [isLocked, onRewind, rewoundText], + [isLocked, onRewind, rewoundText, chatHistory], ); const triggerStyle = useCallback( diff --git a/packages/app/src/components/rewind/use-rewind-agent-mutation.ts b/packages/app/src/components/rewind/use-rewind-agent-mutation.ts index f6d3d0325c..2ff444baa0 100644 --- a/packages/app/src/components/rewind/use-rewind-agent-mutation.ts +++ b/packages/app/src/components/rewind/use-rewind-agent-mutation.ts @@ -1,3 +1,4 @@ +import type { ChatHistoryAttachment } from "@getpaseo/protocol/messages"; import { planTimelineTailFetch } from "@/timeline/timeline-sync-plan"; import { useCallback } from "react"; import { useTranslation } from "react-i18next"; @@ -20,6 +21,7 @@ interface UseRewindAgentMutationInput { interface RewindAgentInput { mode: RewindMode; rewoundText: string; + chatHistory?: ChatHistoryAttachment; } export function useRewindAgentMutation(input: UseRewindAgentMutationInput): { @@ -52,7 +54,10 @@ export function useRewindAgentMutation(input: UseRewindAgentMutationInput): { if (!shouldRestoreComposerForRewindMode(variables.mode)) { return; } - composerRestore?.completeRewind(variables.rewoundText); + composerRestore?.completeRewind({ + text: variables.rewoundText, + chatHistory: variables.chatHistory, + }); }, onError: (error) => { toast.error(error instanceof Error ? error.message : t("rewind.errors.failed")); diff --git a/packages/app/src/composer/attachments/workspace.tsx b/packages/app/src/composer/attachments/workspace.tsx index 80befd2be3..af2ee0d312 100644 --- a/packages/app/src/composer/attachments/workspace.tsx +++ b/packages/app/src/composer/attachments/workspace.tsx @@ -258,6 +258,7 @@ function WorkspaceAttachmentPill({ const content = getWorkspaceAttachmentPillContent(attachment, t); const chatHistory = useMemo(() => { if (attachment.kind !== "chat_history") return null; + if (!attachment.source) return attachment.attachment; const { serverId, agentId, title } = attachment.source; return { ...attachment.attachment, chatHistorySource: { serverId, agentId, title } }; }, [attachment]); diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index 32ad1aedbb..48dd5ff89c 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -1112,6 +1112,7 @@ const ChatAgentReadyContent = memo(function ChatAgentReadyContent({ const parentAgentId = useParentAgentId(serverId, agentId); // With a composer the parent is a pill above it; without one, the box takes the composer's place. const parentPillAgentId = showsComposer ? parentAgentId : null; + const rewindAttachmentScopeKey = useWorkspaceAttachmentScopeKey({ serverId, workspaceId, cwd }); const rawAgentInputDraft = useAgentInputDraft({ draftKey: buildDraftStoreKey({ serverId, @@ -1245,6 +1246,7 @@ const ChatAgentReadyContent = memo(function ChatAgentReadyContent({ const content = ( { + const workdir = mkdtempSync(join(tmpdir(), "fork-local-echo-")); + const clientMessageId = "local-fork-prompt"; + const attachment = { + type: "text", + mimeType: "text/plain", + contextKind: "chat_history", + text: "[User] Original question", + } as const; + class LocalEchoSession extends TestAgentSession { + override async startTurn(): Promise<{ turnId: string }> { + this.pushEvent({ + type: "timeline", + provider: "codex", + item: { + type: "user_message", + text: attachment.text, + messageId: clientMessageId, + clientMessageId, + }, + }); + return super.startTurn(); + } + override async *streamHistory(): AsyncGenerator { + yield { + type: "timeline", + provider: "codex", + item: { + type: "user_message", + text: attachment.text, + ...(replayId ? { messageId: replayId } : {}), + }, + }; + } + } + class LocalEchoClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new LocalEchoSession(config); + } + } + const manager = new AgentManager({ clients: { codex: new LocalEchoClient() }, logger }); + const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, { + workspaceId: undefined, + }); + try { + await drainAsyncGenerator(manager.streamAgent(agent.id, [attachment], { clientMessageId })); + await manager.hydrateTimelineFromProvider(agent.id, { force: true }); + expect((await manager.getTimelineRows(agent.id))[0]?.item).toMatchObject({ + type: "user_message", + text: "", + chatHistory: attachment, + }); + } finally { + await manager.closeAgent(agent.id); + await manager.flush(); + rmSync(workdir, { recursive: true, force: true }); + } + }, +); + test("refreshing an agent replaces the injected timeline store instead of appending a second copy", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-timeline-rehydrate-")); const store = new RecordingTimelineStore(); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 0f791b3e68..f1b5fdc82c 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -5157,7 +5157,8 @@ export class AgentManager { if (attachment?.type === "text" && "mimeType" in attachment) { agent.forkContext = { clientMessageId, - providerMessageId: options?.providerMessageId, + providerMessageId: + options?.providerMessageId !== clientMessageId ? options?.providerMessageId : undefined, text, attachment: { ...attachment, contextKind: "chat_history" }, }; @@ -5213,7 +5214,7 @@ export class AgentManager { } if (!existing || existing.item.type !== "user_message") return null; if (messageId) { - if (agent.forkContext?.clientMessageId === clientMessageId) { + if (agent.forkContext?.clientMessageId === clientMessageId && messageId !== clientMessageId) { agent.forkContext.providerMessageId = messageId; this.enqueueBackgroundPersist(agent); }