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, + ); } }