From c3ae54a1a2de2e7e44b3280883de20977dfda2c4 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Fri, 21 Aug 2026 19:09:56 +0200 Subject: [PATCH] fix(rewind): replace long timelines without replay (#3642) * fix(rewind): refresh canonical timeline after rewind A rewind replaces the daemon timeline and epoch, but replaying the rebuilt history made long conversations visibly reload from the beginning. Capable clients now refresh the existing bounded tail path after a replacement signal while legacy clients retain replay compatibility. * test(rewind): cover coherent timeline replacement Hold rewind completion in the browser regression so it can verify the existing pending state and detect any reconstructed timeline rows before the normal tail refresh is allowed to finish. --- .../e2e/browser/agent-message-rewind.spec.ts | 35 +++++- .../support/helpers/agent-timeline-gate.ts | 63 ++++++++++ .../agent-stream/bottom-anchor-controller.ts | 2 +- .../components/rewind/composer-restore.tsx | 11 +- .../rewind/use-rewind-agent-mutation.ts | 2 +- packages/app/src/contexts/session-context.tsx | 15 +++ packages/app/src/panels/agent-panel.tsx | 8 ++ packages/app/src/runtime/host-runtime.ts | 1 + .../src/timeline/timeline-replacement.test.ts | 21 ++++ .../app/src/timeline/timeline-replacement.ts | 13 ++ packages/protocol/src/client-capabilities.ts | 3 + packages/protocol/src/messages.ts | 10 ++ .../protocol/src/messages.wire-compat.test.ts | 23 ++++ .../server/src/server/agent/agent-manager.ts | 18 ++- .../server/src/server/agent/agent-prompt.ts | 4 + .../src/server/agent/rewind/rewind.test.ts | 21 ++++ .../selective-timeline-delivery.e2e.test.ts | 89 +++++++++++++- packages/server/src/server/session.ts | 113 +++++++++++++++--- 18 files changed, 418 insertions(+), 34 deletions(-) create mode 100644 packages/app/src/timeline/timeline-replacement.test.ts create mode 100644 packages/app/src/timeline/timeline-replacement.ts diff --git a/packages/app/e2e/browser/agent-message-rewind.spec.ts b/packages/app/e2e/browser/agent-message-rewind.spec.ts index 2cb01f89e1..990c922c30 100644 --- a/packages/app/e2e/browser/agent-message-rewind.spec.ts +++ b/packages/app/e2e/browser/agent-message-rewind.spec.ts @@ -3,10 +3,13 @@ import { expect, test, type Page } from "../support/fixtures"; import { expectAgentIdle, expectAgentReadyToInterrupt } from "../support/helpers/agent-stream"; import { cancelAgent, + fillComposerDraft, expectComposerDraft, expectComposerVisible, submitMessage, } from "../support/helpers/composer"; +import { expectNearBottom, scrollChatAwayFromBottom } from "../support/helpers/agent-bottom-anchor"; +import { holdRewindCompletion } from "../support/helpers/agent-timeline-gate"; import { openAgentRoute, seedMockAgentWorkspace } from "../support/helpers/mock-agent"; async function completeSubmittedTurn( @@ -42,9 +45,13 @@ async function expectTurnCompletesNormally( } test.describe("Agent message rewind", () => { - test("rewinds a submitted prompt and restores it to the composer", async ({ page }) => { + test("rewinds a submitted prompt without replaying history and preserves a human draft", async ({ + page, + }) => { test.setTimeout(90_000); - const prompt = "Restore this submitted prompt after rewind."; + const retainedPrompt = "Keep this completed turn after rewind."; + const rewoundPrompt = "Remove this completed turn after rewind."; + const preservedDraft = "Keep this human draft after rewind."; const agent = await seedMockAgentWorkspace({ repoPrefix: "message-rewind-roundtrip-", title: "Message rewind roundtrip", @@ -52,14 +59,30 @@ test.describe("Agent message rewind", () => { }); try { + const rewindCompletion = await holdRewindCompletion(page, agent.agentId); await openAgentRoute(page, agent); await expectComposerVisible(page); - const userMessage = await completeSubmittedTurn(page, agent, prompt); + await completeSubmittedTurn(page, agent, retainedPrompt); + const userMessage = await completeSubmittedTurn(page, agent, rewoundPrompt); + await scrollChatAwayFromBottom(page, { + deltaY: -900, + minDistanceFromBottom: 300, + }); + await fillComposerDraft(page, preservedDraft); + rewindCompletion.clearTimelineStreamCount(); - await rewindConversation(page, userMessage, prompt); + const rewind = rewindConversation(page, userMessage, rewoundPrompt); + await rewindCompletion.waitForDelayedResponse(); + await expect(page.getByTestId("rewind-menu-conversation")).toBeDisabled(); + expect(rewindCompletion.timelineStreamCount()).toBe(0); + rewindCompletion.release(); + await rewind; - await expect(page.getByText("(end of synthetic stream)", { exact: true })).toHaveCount(0); - await expectComposerDraft(page, prompt); + await expect(page.getByTestId("user-message").filter({ hasText: rewoundPrompt })).toHaveCount( + 0, + ); + await expectComposerDraft(page, preservedDraft); + await expectNearBottom(page); } finally { await agent.cleanup(); } diff --git a/packages/app/e2e/support/helpers/agent-timeline-gate.ts b/packages/app/e2e/support/helpers/agent-timeline-gate.ts index eca355698b..91b6867cf0 100644 --- a/packages/app/e2e/support/helpers/agent-timeline-gate.ts +++ b/packages/app/e2e/support/helpers/agent-timeline-gate.ts @@ -27,6 +27,13 @@ export interface DaemonHydrationGate { release(): void; } +export interface RewindCompletionGate { + clearTimelineStreamCount(): void; + release(): void; + timelineStreamCount(): number; + waitForDelayedResponse(): Promise; +} + export interface PromptJumpRequestTracker { requests(): Array<{ cursorSeq: number | null; limit: number | null; mergeWindow: boolean }>; } @@ -255,6 +262,62 @@ export async function holdDaemonHydration(page: Page): Promise { + let released = false; + let timelineStreamCount = 0; + const delayedForwards: Array<() => void> = []; + let resolveDelayedResponse: (() => void) | null = null; + const delayedResponse = new Promise((resolve) => { + resolveDelayedResponse = resolve; + }); + + await page.routeWebSocket(daemonWsRoutePattern(), (ws) => { + const server = ws.connectToServer(); + ws.onMessage((message) => server.send(message)); + server.onMessage((message) => { + const sessionMessage = getSessionMessage(message); + const payload = sessionMessage ? getPayload(sessionMessage) : null; + const streamEvent = payload?.event; + if ( + sessionMessage?.type === "agent_stream" && + payload?.agentId === agentId && + streamEvent && + typeof streamEvent === "object" && + (streamEvent as { type?: unknown }).type === "timeline" + ) { + timelineStreamCount += 1; + } + if ( + !released && + sessionMessage?.type === "agent.rewind.response" && + payload?.agentId === agentId + ) { + delayedForwards.push(() => ws.send(message)); + resolveDelayedResponse?.(); + return; + } + ws.send(message); + }); + }); + + return { + clearTimelineStreamCount() { + timelineStreamCount = 0; + }, + release() { + released = true; + for (const forward of delayedForwards.splice(0)) { + forward(); + } + }, + timelineStreamCount: () => timelineStreamCount, + waitForDelayedResponse: () => delayedResponse, + }; +} + export async function delayCreatedAgentInitialTailResponse( page: Page, ): Promise { diff --git a/packages/app/src/agent-stream/bottom-anchor-controller.ts b/packages/app/src/agent-stream/bottom-anchor-controller.ts index 35bd6e6b6f..7533699815 100644 --- a/packages/app/src/agent-stream/bottom-anchor-controller.ts +++ b/packages/app/src/agent-stream/bottom-anchor-controller.ts @@ -10,7 +10,7 @@ export interface BottomAnchorRouteRequest { } export interface BottomAnchorLocalRequest { - reason: "jump-to-bottom" | "message-sent"; + reason: "jump-to-bottom" | "message-sent" | "rewind"; agentId: string; } diff --git a/packages/app/src/components/rewind/composer-restore.tsx b/packages/app/src/components/rewind/composer-restore.tsx index 286a516acc..e615db5912 100644 --- a/packages/app/src/components/rewind/composer-restore.tsx +++ b/packages/app/src/components/rewind/composer-restore.tsx @@ -9,12 +9,13 @@ import { } from "react"; interface RewindComposerRestoreContextValue { - restoreTextIfComposerEmpty: (text: string) => void; + completeRewind: (text: string) => void; } interface RewindComposerRestoreProviderProps { text: string; setText: (text: string) => void; + onRewindComplete: () => void; children: ReactNode; } @@ -33,6 +34,7 @@ export function restoreComposerTextIfEmpty(input: { export function RewindComposerRestoreProvider({ text, setText, + onRewindComplete, children, }: RewindComposerRestoreProviderProps) { const textRef = useRef(text); @@ -41,7 +43,7 @@ export function RewindComposerRestoreProvider({ textRef.current = text; }, [text]); - const restoreTextIfComposerEmpty = useCallback( + const completeRewind = useCallback( (rewoundText: string) => { const nextText = restoreComposerTextIfEmpty({ currentText: textRef.current, @@ -50,11 +52,12 @@ export function RewindComposerRestoreProvider({ if (nextText !== textRef.current) { setText(nextText); } + onRewindComplete(); }, - [setText], + [onRewindComplete, setText], ); - const value = useMemo(() => ({ restoreTextIfComposerEmpty }), [restoreTextIfComposerEmpty]); + const value = useMemo(() => ({ completeRewind }), [completeRewind]); return ( 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 f1483808cf..d482bb2d42 100644 --- a/packages/app/src/components/rewind/use-rewind-agent-mutation.ts +++ b/packages/app/src/components/rewind/use-rewind-agent-mutation.ts @@ -52,7 +52,7 @@ export function useRewindAgentMutation(input: UseRewindAgentMutationInput): { if (!shouldRestoreComposerForRewindMode(variables.mode)) { return; } - composerRestore?.restoreTextIfComposerEmpty(variables.rewoundText); + composerRestore?.completeRewind(variables.rewoundText); }, onError: (error) => { toast.error(error instanceof Error ? error.message : t("rewind.errors.failed")); diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 6cf80708e9..5e937c61bb 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -23,6 +23,7 @@ import { isTimelineResumeSnapshotAuthoritative, planTimelineTailFetch, } from "@/timeline/timeline-sync-plan"; +import { requestTimelineReplacement } from "@/timeline/timeline-replacement"; import { createViewedTimelineSync, type TimelineDeliveryMode, @@ -788,6 +789,19 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider applyTimelineResponse(message.payload); }); + const unsubTimelineReplacement = client.on("agent.timeline.replacement", (message) => { + if (message.type !== "agent.timeline.replacement") return; + void requestTimelineReplacement( + { + fetchAgentTimeline: (agentId, request) => + getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, request), + }, + message.payload.agentId, + ).catch((error: unknown) => { + console.warn("[Session] timeline replacement refresh failed", { serverId, error }); + }); + }); + const unsubProviderSubagentUpdate = client.on("agent.provider_subagents.update", (message) => { if (message.type !== "agent.provider_subagents.update") return; useProviderSubagentStore.getState().applyUpdate(serverId, message.payload); @@ -983,6 +997,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider }); return () => { + unsubTimelineReplacement(); unsubAgentStream(); unsubAgentTimeline(); unsubProviderSubagentUpdate(); diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index 1f4c4feca7..49e0473414 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -1039,6 +1039,10 @@ function ChatAgentContent({ streamViewRef.current?.scrollToBottom("message-sent"); }, [agentId]); + const handleRewindComplete = useCallback(() => { + streamViewRef.current?.scrollToBottom("rewind"); + }, []); + const retryAgentLoad = useCallback(() => setMissingAgentState({ kind: "idle" }), []); const retryTimelineSync = useCallback(() => { @@ -1178,6 +1182,7 @@ function ChatAgentContent({ animatedContentStyle={animatedContentStyle} handleComposerHeightChange={handleComposerHeightChange} handleMessageSent={handleMessageSent} + handleRewindComplete={handleRewindComplete} showHistorySyncOverlay={showHistorySyncOverlay} showHistorySyncError={showHistorySyncError} isRetryingHistorySync={isRetryingHistorySync} @@ -1206,6 +1211,7 @@ const ChatAgentReadyContent = memo(function ChatAgentReadyContent({ animatedContentStyle, handleComposerHeightChange, handleMessageSent, + handleRewindComplete, showHistorySyncOverlay, showHistorySyncError, isRetryingHistorySync, @@ -1230,6 +1236,7 @@ const ChatAgentReadyContent = memo(function ChatAgentReadyContent({ animatedContentStyle: object[]; handleComposerHeightChange: (height: number) => void; handleMessageSent: () => void; + handleRewindComplete: () => void; showHistorySyncOverlay: boolean; showHistorySyncError: boolean; isRetryingHistorySync: boolean; @@ -1355,6 +1362,7 @@ const ChatAgentReadyContent = memo(function ChatAgentReadyContent({ diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index 4d27f200c4..9a5e0338b3 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -480,6 +480,7 @@ function createDefaultDeps(): HostRuntimeControllerDeps { : undefined; const appCapabilities = { [CLIENT_CAPS.selectiveAgentTimeline]: true, + [CLIENT_CAPS.timelineReplacementInvalidation]: true, ...browserAutomationCapabilities, }; diff --git a/packages/app/src/timeline/timeline-replacement.test.ts b/packages/app/src/timeline/timeline-replacement.test.ts new file mode 100644 index 0000000000..f457ea7daa --- /dev/null +++ b/packages/app/src/timeline/timeline-replacement.test.ts @@ -0,0 +1,21 @@ +import { describe, expect, test } from "vitest"; +import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import { requestTimelineReplacement } from "./timeline-replacement"; + +describe("requestTimelineReplacement", () => { + test("requests the existing projected tail without receiving a timeline installer", async () => { + const requests: Parameters[] = []; + const client = { + fetchAgentTimeline: async (...request: Parameters) => { + requests.push(request); + throw new Error("fetch complete"); + }, + }; + + await expect(requestTimelineReplacement(client, "agent-1")).rejects.toThrow("fetch complete"); + + expect(requests).toEqual([ + ["agent-1", { direction: "tail", projection: "projected", limit: 40 }], + ]); + }); +}); diff --git a/packages/app/src/timeline/timeline-replacement.ts b/packages/app/src/timeline/timeline-replacement.ts new file mode 100644 index 0000000000..29bbe7df97 --- /dev/null +++ b/packages/app/src/timeline/timeline-replacement.ts @@ -0,0 +1,13 @@ +import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import { planTimelineTailFetch } from "./timeline-sync-plan"; + +interface TimelineReplacementFetch { + fetchAgentTimeline: DaemonClient["fetchAgentTimeline"]; +} + +export function requestTimelineReplacement( + client: TimelineReplacementFetch, + agentId: string, +): ReturnType { + return client.fetchAgentTimeline(agentId, planTimelineTailFetch()); +} diff --git a/packages/protocol/src/client-capabilities.ts b/packages/protocol/src/client-capabilities.ts index dd2b90369b..b9c405b4f1 100644 --- a/packages/protocol/src/client-capabilities.ts +++ b/packages/protocol/src/client-capabilities.ts @@ -24,6 +24,9 @@ export const CLIENT_CAPS = { // provider catalogs with shared thinking sets and may revalidate by content hash. // Remove the legacy snapshot encoding after 2027-02-04. compactProviderSnapshots: "compact_provider_snapshots", + // COMPAT(timelineReplacementInvalidation): added in v0.5.0, remove legacy + // reconstructed timeline replay after 2027-02-21 once the client floor supports invalidation. + timelineReplacementInvalidation: "timeline_replacement_invalidation", browserHost: "browser_host", } as const; diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index bcc9c58450..89e9bc2730 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -4266,6 +4266,14 @@ export const FetchAgentTimelineResponseMessageSchema = z.object({ }), }); +export const AgentTimelineReplacementMessageSchema = z.object({ + type: z.literal("agent.timeline.replacement"), + payload: z.object({ + agentId: z.string(), + epoch: z.string(), + }), +}); + export const AgentTimelineListPromptsResponseMessageSchema = z.object({ type: z.literal("agent.timeline.list_prompts.response"), payload: z.object({ @@ -6222,6 +6230,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ ArchiveWorkspaceResponseMessageSchema, FetchAgentResponseMessageSchema, FetchAgentTimelineResponseMessageSchema, + AgentTimelineReplacementMessageSchema, AgentTimelineListPromptsResponseMessageSchema, ProviderSubagentListResponseMessageSchema, ProviderSubagentTimelineResponseMessageSchema, @@ -6829,6 +6838,7 @@ export const WSHelloMessageSchema = z.object({ [CLIENT_CAPS.providerSubagents]: z.boolean().optional(), [CLIENT_CAPS.projectUpdates]: z.boolean().optional(), [CLIENT_CAPS.compactProviderSnapshots]: z.boolean().optional(), + [CLIENT_CAPS.timelineReplacementInvalidation]: z.boolean().optional(), [CLIENT_CAPS.browserHost]: BrowserAutomationHostCapabilitySchema.optional(), }) .passthrough() diff --git a/packages/protocol/src/messages.wire-compat.test.ts b/packages/protocol/src/messages.wire-compat.test.ts index 22d22acacf..6bd948e28d 100644 --- a/packages/protocol/src/messages.wire-compat.test.ts +++ b/packages/protocol/src/messages.wire-compat.test.ts @@ -4,6 +4,7 @@ import { AgentSnapshotPayloadSchema, AgentTimelineItemPayloadSchema, ServerInfoStatusPayloadSchema, + SessionOutboundMessageSchema, WSHelloMessageSchema, } from "./messages.js"; @@ -75,6 +76,28 @@ describe("wire schema compatibility", () => { ]); }); + test("timeline replacement invalidation is opt-in and carries no timeline rows", () => { + expect( + WSHelloMessageSchema.parse({ + type: "hello", + clientId: "capable-client", + clientType: "mobile", + protocolVersion: 1, + capabilities: { timeline_replacement_invalidation: true }, + }).capabilities, + ).toEqual({ timeline_replacement_invalidation: true }); + + expect( + SessionOutboundMessageSchema.parse({ + type: "agent.timeline.replacement", + payload: { agentId: "agent-1", epoch: "epoch-2" }, + }), + ).toEqual({ + type: "agent.timeline.replacement", + payload: { agentId: "agent-1", epoch: "epoch-2" }, + }); + }); + test("server info strips unknown legacy features while accepting former turn identity", () => { const parsed = ServerInfoStatusPayloadSchema.parse({ status: "server_info", diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 2634324505..509d6f4982 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -185,6 +185,7 @@ export type { export type AgentManagerEvent = | { type: "agent_state"; agent: ManagedAgent } | { type: "provider_subagent"; event: ProviderSubagentStoreEvent } + | { type: "timeline_replacement"; agentId: string; epoch: string } | { type: "agent_stream"; agentId: string; @@ -204,6 +205,7 @@ export interface SubscribeOptions { interface HydrateTimelineOptions { force?: boolean; broadcast?: boolean | (() => boolean); + broadcastTimeline?: boolean; } export type ImportablePersistedAgentQueryOptions = ListImportableSessionsOptions & { @@ -2822,7 +2824,16 @@ export class AgentManager { ); await invokeRewindCapability(agent.session, { messageId: providerMessageId, mode }); if (mode !== "files") { - await this.hydrateTimelineFromProvider(agentId, { force: true, broadcast: true }); + await this.hydrateTimelineFromProvider(agentId, { + force: true, + broadcast: true, + broadcastTimeline: false, + }); + this.dispatch({ + type: "timeline_replacement", + agentId, + epoch: this.timelineStore.getEpoch(agentId), + }); } await this.refreshRuntimeInfo(agent); await this.persistSnapshot(agent); @@ -3589,11 +3600,13 @@ export class AgentManager { } const broadcast = options?.broadcast ?? false; + const broadcastTimeline = options?.broadcastTimeline ?? broadcast; if (options?.force) { await this.forceHydrateTimelineFromLegacyProviderHistory( agent, typeof broadcast === "function" ? broadcast() : broadcast, + typeof broadcastTimeline === "function" ? broadcastTimeline() : broadcastTimeline, ); return; } @@ -3604,6 +3617,7 @@ export class AgentManager { private async forceHydrateTimelineFromLegacyProviderHistory( agent: ActiveManagedAgent, broadcast: boolean, + broadcastTimeline: boolean, ): Promise { const historyEvents: Extract[] = []; const providerSubagentEvents: Extract[] = []; @@ -3642,7 +3656,7 @@ export class AgentManager { event.item, event.timestamp ? { timestamp: event.timestamp } : undefined, ); - if (broadcast) { + if (broadcastTimeline) { this.dispatchStream(agent.id, event, { seq: row.seq, epoch: this.timelineStore.getEpoch(agent.id), diff --git a/packages/server/src/server/agent/agent-prompt.ts b/packages/server/src/server/agent/agent-prompt.ts index 07c7b4f57a..f78cdf0ef1 100644 --- a/packages/server/src/server/agent/agent-prompt.ts +++ b/packages/server/src/server/agent/agent-prompt.ts @@ -462,6 +462,10 @@ export function setupFinishNotification(params: SetupFinishNotificationParams): return; } + if (event.type === "timeline_replacement") { + return; + } + if (event.event.type === "permission_requested") { // A permission pause is an intermediate checkpoint. Forget the run // observed before it so an idle state during follow-up startup cannot diff --git a/packages/server/src/server/agent/rewind/rewind.test.ts b/packages/server/src/server/agent/rewind/rewind.test.ts index a321fb79fd..7040c42e3e 100644 --- a/packages/server/src/server/agent/rewind/rewind.test.ts +++ b/packages/server/src/server/agent/rewind/rewind.test.ts @@ -80,6 +80,27 @@ describe("AgentManager rewind", () => { ]); }); + test("replaces the canonical epoch without replaying reconstructed parent rows", async () => { + const { manager, session, agentId } = await createRewindHarness(); + session.history = Array.from({ length: 250 }, (_, index) => ({ + type: "assistant_message" as const, + text: `rewound ${index}`, + })); + const epochBefore = manager.fetchTimeline(agentId, { limit: 0 }).epoch; + const events: string[] = []; + const unsubscribe = manager.subscribe((event) => events.push(event.type), { + replayState: false, + }); + + await manager.rewind(agentId, "message-1", "conversation"); + unsubscribe(); + + const replacement = manager.fetchTimeline(agentId, { limit: 0 }); + expect(replacement.epoch).not.toBe(epochBefore); + expect(replacement.rows).toHaveLength(250); + expect(events.filter((type) => type === "agent_stream")).toEqual([]); + }); + test("rewinds files without rehydrating the conversation timeline", async () => { const { manager, session, agentId } = await createRewindHarness(); diff --git a/packages/server/src/server/selective-timeline-delivery.e2e.test.ts b/packages/server/src/server/selective-timeline-delivery.e2e.test.ts index 8933def320..c9d57c4b59 100644 --- a/packages/server/src/server/selective-timeline-delivery.e2e.test.ts +++ b/packages/server/src/server/selective-timeline-delivery.e2e.test.ts @@ -3,6 +3,7 @@ import { CLIENT_CAPS } from "@getpaseo/protocol/client-capabilities"; import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; import { DaemonClient } from "./test-utils/daemon-client.js"; import { createTestPaseoDaemon, type TestPaseoDaemon } from "./test-utils/paseo-daemon.js"; +import { MockLoadTestAgentClient } from "./agent/providers/mock-load-test-agent.js"; interface MessageWaiter { predicate(message: SessionOutboundMessage): boolean; @@ -124,11 +125,20 @@ afterEach(async () => { await daemon.close(); }, 30_000); -async function connect(input: { clientId: string; selective: boolean }): Promise { +async function connect(input: { + clientId: string; + selective: boolean; + timelineReplacementInvalidation?: boolean; +}): Promise { const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws`, clientId: input.clientId, - capabilities: { [CLIENT_CAPS.selectiveAgentTimeline]: input.selective }, + capabilities: { + [CLIENT_CAPS.selectiveAgentTimeline]: input.selective, + ...(input.timelineReplacementInvalidation + ? { [CLIENT_CAPS.timelineReplacementInvalidation]: true } + : {}), + }, reconnect: { enabled: false }, }); await client.connect(); @@ -137,6 +147,81 @@ async function connect(input: { clientId: string; selective: boolean }): Promise return connected; } +test("rewind routes replacement completion by source capability and subscription", async () => { + await daemon.close(); + daemon = await createTestPaseoDaemon({ + isDev: true, + agentClients: { mock: new MockLoadTestAgentClient() }, + }); + const initiating = await connect({ + clientId: "rewind-initiating", + selective: true, + timelineReplacementInvalidation: true, + }); + const passive = await connect({ + clientId: "rewind-passive", + selective: true, + timelineReplacementInvalidation: true, + }); + const unrelated = await connect({ + clientId: "rewind-unrelated", + selective: true, + timelineReplacementInvalidation: true, + }); + const legacy = await connect({ clientId: "rewind-legacy", selective: false }); + const agent = await initiating.client.createAgent({ + provider: "mock", + cwd: "/tmp", + title: "Rewind routing", + model: "ten-second-stream", + }); + + await Promise.all([ + initiating.client.setAgentTimelineSubscription([agent.id]), + passive.client.setAgentTimelineSubscription([agent.id]), + unrelated.client.setAgentTimelineSubscription([]), + ]); + await initiating.client.sendMessage(agent.id, "Rewind this synthetic prompt"); + await initiating.client.cancelAgent(agent.id); + const timeline = await initiating.client.fetchAgentTimeline(agent.id, { + direction: "tail", + limit: 0, + projection: "canonical", + }); + const target = timeline.entries.find( + (entry) => + entry.item.type === "user_message" && entry.item.text === "Rewind this synthetic prompt", + ); + if (!target || target.item.type !== "user_message" || !target.item.messageId) { + throw new Error("Expected rewindable canonical user message"); + } + + for (const connected of [initiating, passive, unrelated, legacy]) connected.clear(); + await initiating.client.rewindAgent(agent.id, target.item.messageId, "conversation"); + await Promise.all([ + initiating.barrier("rewind-initiator"), + passive.barrier("rewind-passive"), + unrelated.barrier("rewind-unrelated"), + legacy.barrier("rewind-legacy"), + ]); + + expect( + initiating.messages.filter((message) => message.type === "agent.rewind.response"), + ).toHaveLength(1); + expect( + initiating.messages.filter((message) => message.type === "agent.timeline.replacement"), + ).toHaveLength(0); + expect( + passive.messages.filter((message) => message.type === "agent.timeline.replacement"), + ).toHaveLength(1); + expect(passive.hasTimeline(agent.id)).toBe(false); + expect(legacy.messages.filter(isAgentStream(agent.id)).length).toBeGreaterThan(0); + expect( + unrelated.messages.filter((message) => message.type === "agent.timeline.replacement"), + ).toHaveLength(0); + expect(unrelated.hasTimeline(agent.id)).toBe(false); +}, 30_000); + test("subscription acknowledgements stay on the requesting socket of a retained session", async () => { const legacy = await connect({ clientId: "shared-client", selective: false }); const capable = await connect({ clientId: "shared-client", selective: true }); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index b0ebb7f248..af3902898a 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -661,6 +661,7 @@ export class Session { private readonly paseoHome: string; private readonly projectIcons: ProjectIconReader; private readonly worktreesRoot: string | undefined; + private readonly rewindInitiators = new Map(); private agentManager: AgentManager; private readonly agentStorage: AgentStorage; @@ -1620,6 +1621,11 @@ export class Session { this.unsubscribeAgentEvents = this.agentManager.subscribe( (event) => { + if (event.type === "timeline_replacement") { + this.deliverTimelineReplacement(event.agentId, this.rewindInitiators.get(event.agentId)); + return; + } + if (event.type === "agent_state") { this.sessionLogger.trace( { @@ -1907,7 +1913,7 @@ export class Session { private async dispatchInboundMessage(msg: SessionInboundMessage, source?: object): Promise { const promise = this.dispatchVoiceAndControlMessage(msg) ?? - this.dispatchAgentRewindMessage(msg) ?? + this.dispatchAgentRewindMessage(msg, source) ?? this.dispatchAgentRelationshipMessage(msg) ?? this.dispatchAgentTimelineMessage(msg, source) ?? this.dispatchHubExecutionMessage(msg) ?? @@ -2153,10 +2159,13 @@ export class Session { } } - private dispatchAgentRewindMessage(msg: SessionInboundMessage): Promise | undefined { + private dispatchAgentRewindMessage( + msg: SessionInboundMessage, + source?: object, + ): Promise | undefined { switch (msg.type) { case "agent.rewind.request": - return this.handleAgentRewindRequest(msg); + return this.handleAgentRewindRequest(msg, source); default: return undefined; } @@ -3872,28 +3881,96 @@ export class Session { private async handleAgentRewindRequest( msg: Extract, + source?: object, ): Promise { try { + this.rewindInitiators.set(msg.agentId, source); await this.agentManager.rewind(msg.agentId, msg.messageId, msg.mode); - this.emit({ - type: "agent.rewind.response", - payload: { - requestId: msg.requestId, - agentId: msg.agentId, - ok: true, - error: null, + this.emitForSource( + { + type: "agent.rewind.response", + payload: { + requestId: msg.requestId, + agentId: msg.agentId, + ok: true, + error: null, + }, }, - }); + source, + ); } catch (error) { - this.emit({ - type: "agent.rewind.response", - payload: { - requestId: msg.requestId, - agentId: msg.agentId, - ok: false, - error: error instanceof Error ? error.message : "Failed to rewind agent", + this.emitForSource( + { + type: "agent.rewind.response", + payload: { + requestId: msg.requestId, + agentId: msg.agentId, + ok: false, + error: error instanceof Error ? error.message : "Failed to rewind agent", + }, }, + source, + ); + } finally { + this.rewindInitiators.delete(msg.agentId); + } + } + + private deliverTimelineReplacement(agentId: string, initiatingSource?: object): void { + const agent = this.agentManager.getAgent(agentId); + if (!agent) return; + const timeline = this.agentManager.fetchTimeline(agentId, { limit: 0 }); + const epoch = timeline.epoch; + + if (this.clientCapabilitiesBySource.size === 0 || !this.onMessageToSource) { + if (!this.supports(CLIENT_CAPS.timelineReplacementInvalidation)) { + this.emitReconstructedTimelineRows(agentId, agent.provider, timeline.rows, epoch); + } + return; + } + + for (const [source, capabilities] of this.clientCapabilitiesBySource) { + const isInitiator = source === initiatingSource; + const supportsReplacement = capabilities.has(CLIENT_CAPS.timelineReplacementInvalidation); + const isSubscribed = this.viewedTimelineAgentIdsBySource.get(source)?.has(agentId) === true; + if (supportsReplacement) { + if (isSubscribed && !isInitiator) { + this.onMessageToSource(source, { + type: "agent.timeline.replacement", + payload: { agentId, epoch }, + }); + } + continue; + } + // COMPAT(timelineReplacementInvalidation): added in v0.5.0, replay reconstructed + // rows to legacy clients until the supported client floor is >= v0.5.0 after 2027-02-21. + this.emitReconstructedTimelineRows(agentId, agent.provider, timeline.rows, epoch, source); + } + } + + private emitReconstructedTimelineRows( + agentId: string, + provider: ManagedAgent["provider"], + rows: AgentTimelineFetchResult["rows"], + epoch: string, + source?: object, + ): void { + for (const row of rows) { + const event = serializeAgentStreamEvent({ + type: "timeline", + provider, + item: row.item, + ...(row.turnId ? { turnId: row.turnId } : {}), + timestamp: row.timestamp, }); + if (!event) continue; + this.emitForSource( + { + type: "agent_stream", + payload: { agentId, event, timestamp: row.timestamp, seq: row.seq, epoch }, + }, + source, + ); } }