From 7fed6d144df852e339215e3a5eca8ce8422ed83f Mon Sep 17 00:00:00 2001 From: Basit Balogun Date: Mon, 1 Jun 2026 12:55:35 +0100 Subject: [PATCH 1/3] feat: restore checkpoint after any AI response turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add checkpoint tRPC router with restore procedure that reverts git state, truncates session JSONL to the restore point, and deletes orphaned checkpoint refs for abandoned future turns - Track lastCheckpointId per turn in buildConversationItems so each completed agent turn knows its git ref - Show per-turn restore button in AgentMessage (disabled with tooltip when no checkpoint exists for that turn) - Add CheckpointTimelineModal (mod+shift+h) — command-palette-style list of all checkpoints in the session, newest first, with user message snippet and relative timestamp; shortcut is user-remappable via keybindings store - Add RestoreCheckpointDialog with confirmation warning before reverting - Add useRestoreCheckpoint hook to wire restore flow end-to-end - Register checkpoint-timeline as a configurable shortcut Closes #2328 --- apps/code/src/main/services/agent/service.ts | 8 + apps/code/src/main/trpc/router.ts | 2 + apps/code/src/main/trpc/routers/checkpoint.ts | 219 ++++++++++++++++++ .../renderer/constants/keyboard-shortcuts.ts | 70 ++++++ .../components/CheckpointTimelineModal.tsx | 173 ++++++++++++++ .../sessions/components/ConversationView.tsx | 54 ++++- .../components/RestoreCheckpointDialog.tsx | 54 +++++ .../components/buildConversationItems.ts | 15 ++ .../session-update/AgentMessage.tsx | 36 ++- .../session-update/SessionUpdateView.tsx | 10 +- .../sessions/hooks/useRestoreCheckpoint.ts | 69 ++++++ .../features/sessions/stores/sessionStore.ts | 52 ++++- 12 files changed, 755 insertions(+), 7 deletions(-) create mode 100644 apps/code/src/main/trpc/routers/checkpoint.ts create mode 100644 apps/code/src/renderer/features/sessions/components/CheckpointTimelineModal.tsx create mode 100644 apps/code/src/renderer/features/sessions/components/RestoreCheckpointDialog.tsx create mode 100644 apps/code/src/renderer/features/sessions/hooks/useRestoreCheckpoint.ts diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index f6348191ad..d86221b9e5 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -993,6 +993,14 @@ When creating pull requests, add the following footer at the end of the PR descr return this.sessions.get(taskRunId); } + getSessionInfo( + taskRunId: string, + ): { sessionId: string; repoPath: string } | undefined { + const session = this.sessions.get(taskRunId); + if (!session?.config.sessionId) return undefined; + return { sessionId: session.config.sessionId, repoPath: session.repoPath }; + } + async setSessionConfigOption( sessionId: string, configId: string, diff --git a/apps/code/src/main/trpc/router.ts b/apps/code/src/main/trpc/router.ts index f0f8dd9eb5..08cf091103 100644 --- a/apps/code/src/main/trpc/router.ts +++ b/apps/code/src/main/trpc/router.ts @@ -1,6 +1,7 @@ import { additionalDirectoriesRouter } from "./routers/additional-directories"; import { agentRouter } from "./routers/agent"; import { analyticsRouter } from "./routers/analytics"; +import { checkpointRouter } from "./routers/checkpoint"; import { archiveRouter } from "./routers/archive"; import { authRouter } from "./routers/auth"; import { cloudTaskRouter } from "./routers/cloud-task"; @@ -46,6 +47,7 @@ export const trpcRouter = router({ analytics: analyticsRouter, archive: archiveRouter, auth: authRouter, + checkpoint: checkpointRouter, cloudTask: cloudTaskRouter, connectivity: connectivityRouter, contextMenu: contextMenuRouter, diff --git a/apps/code/src/main/trpc/routers/checkpoint.ts b/apps/code/src/main/trpc/routers/checkpoint.ts new file mode 100644 index 0000000000..aae95e49ea --- /dev/null +++ b/apps/code/src/main/trpc/routers/checkpoint.ts @@ -0,0 +1,219 @@ +import fs from "node:fs/promises"; +import { isNotification, POSTHOG_NOTIFICATIONS } from "@posthog/agent"; +import { getSessionJsonlPath } from "@posthog/agent/adapters/claude/session/jsonl-hydration"; +import { createGitClient } from "@posthog/git/client"; +import { + deleteCheckpoint, + RevertCheckpointSaga, +} from "@posthog/git/sagas/checkpoint"; +import { z } from "zod"; +import { container } from "../../di/container"; +import { MAIN_TOKENS } from "../../di/tokens"; +import type { AgentService } from "../../services/agent/service"; +import { logger } from "../../utils/logger"; +import { publicProcedure, router } from "../trpc"; + +const log = logger.scope("checkpoint-router"); + +const getAgentService = () => + container.get(MAIN_TOKENS.AgentService); + +const restoreInput = z.object({ + checkpointId: z.string(), + repoPath: z.string(), + taskRunId: z.string().optional(), +}); + +const restoreOutput = z.object({ + checkpointId: z.string(), + commit: z.string(), + head: z.string().nullable(), + branch: z.string().nullable(), +}); + +interface TruncateResult { + truncated: boolean; + /** Checkpoint IDs that appear in the discarded portion (orphaned refs). */ + orphanedCheckpointIds: string[]; +} + +/** + * Truncate a session JSONL file at the turn containing the given checkpoint. + * Finds the `_posthog/git_checkpoint` entry with matching checkpointId, then + * includes all entries up to (but not including) the next user message group. + * Returns orphaned checkpoint IDs from the discarded lines for cleanup. + */ +async function truncateSessionJsonl( + jsonlPath: string, + checkpointId: string, +): Promise { + let content: string; + try { + content = await fs.readFile(jsonlPath, "utf-8"); + } catch { + return { truncated: false, orphanedCheckpointIds: [] }; + } + + const lines = content.split("\n").filter((l) => l.trim()); + let checkpointLineIdx = -1; + + for (let i = 0; i < lines.length; i++) { + try { + const entry = JSON.parse(lines[i]); + const method = entry.notification?.method; + if (!method) continue; + if (!isNotification(method, POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT)) + continue; + const params = entry.notification?.params; + if (params?.checkpointId === checkpointId) { + checkpointLineIdx = i; + break; + } + } catch { + // skip malformed lines + } + } + + if (checkpointLineIdx === -1) + return { truncated: false, orphanedCheckpointIds: [] }; + + // Find the end of the current turn: scan forward for the next user message + // group start (a user_message_chunk after non-user content). + let cutoff = lines.length; + let inUserMessage = false; + let passedNonUser = false; + + for (let i = checkpointLineIdx + 1; i < lines.length; i++) { + try { + const entry = JSON.parse(lines[i]); + const method = entry.notification?.method; + const params = entry.notification?.params as + | Record + | undefined; + + if (method === "session/update" && params?.update) { + const update = params.update as { sessionUpdate?: string }; + const isUserChunk = + update.sessionUpdate === "user_message" || + update.sessionUpdate === "user_message_chunk"; + + if (isUserChunk) { + if (passedNonUser && !inUserMessage) { + // Start of a new user message group — stop here + cutoff = i; + break; + } + inUserMessage = true; + } else { + if (inUserMessage) { + passedNonUser = true; + } + inUserMessage = false; + } + } else if (method === "session/prompt") { + // session/prompt request also marks a turn boundary + cutoff = i; + break; + } else { + if (inUserMessage) { + passedNonUser = true; + } + inUserMessage = false; + } + } catch { + // skip malformed + } + } + + // Collect checkpoint IDs from the discarded lines so their refs can be cleaned up + const orphanedCheckpointIds: string[] = []; + for (let i = cutoff; i < lines.length; i++) { + try { + const entry = JSON.parse(lines[i]); + const method = entry.notification?.method; + if (!method) continue; + if (!isNotification(method, POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT)) + continue; + const id = entry.notification?.params?.checkpointId; + if (id) orphanedCheckpointIds.push(id); + } catch { + // skip malformed + } + } + + const truncatedLines = lines.slice(0, cutoff); + const tmpPath = `${jsonlPath}.tmp.${Date.now()}`; + await fs.writeFile(tmpPath, `${truncatedLines.join("\n")}\n`, "utf-8"); + await fs.rename(tmpPath, jsonlPath); + + log.info("Truncated session JSONL", { + checkpointId, + originalLines: lines.length, + truncatedLines: truncatedLines.length, + orphanedCheckpointIds, + }); + return { truncated: true, orphanedCheckpointIds }; +} + +export const checkpointRouter = router({ + restore: publicProcedure + .input(restoreInput) + .output(restoreOutput) + .mutation(async ({ input }) => { + // 1. Revert git files to checkpoint state + const saga = new RevertCheckpointSaga(); + const result = await saga.run({ + baseDir: input.repoPath, + checkpointId: input.checkpointId, + }); + if (!result.success) { + throw new Error(result.error ?? "Failed to revert checkpoint"); + } + + // 2. Truncate agent's session JSONL, clean up orphaned checkpoint refs, and restart agent + if (input.taskRunId) { + try { + const agentService = getAgentService(); + const info = agentService.getSessionInfo(input.taskRunId); + if (info) { + const jsonlPath = getSessionJsonlPath( + info.sessionId, + info.repoPath, + ); + const { truncated, orphanedCheckpointIds } = + await truncateSessionJsonl(jsonlPath, input.checkpointId); + if (truncated) { + // Delete git refs for checkpoints in the abandoned future turns + if (orphanedCheckpointIds.length > 0) { + const git = createGitClient(input.repoPath); + await Promise.all( + orphanedCheckpointIds.map((id) => + deleteCheckpoint(git, id).catch(() => {}), + ), + ); + log.info("Deleted orphaned checkpoint refs", { + orphanedCheckpointIds, + }); + } + // Cancel the agent session — the renderer will automatically + // reconnect and resume from the truncated JSONL + await agentService.cancelSession(input.taskRunId); + log.info("Agent session cancelled for checkpoint restore", { + taskRunId: input.taskRunId, + checkpointId: input.checkpointId, + }); + } + } + } catch (err) { + // Non-fatal: git files were already reverted successfully. + // The UI will truncate its events regardless. + log.warn("Failed to truncate agent session", { + taskRunId: input.taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + } + } + + return result; + }), +}); diff --git a/apps/code/src/renderer/constants/keyboard-shortcuts.ts b/apps/code/src/renderer/constants/keyboard-shortcuts.ts index b162013bbc..bc737d4653 100644 --- a/apps/code/src/renderer/constants/keyboard-shortcuts.ts +++ b/apps/code/src/renderer/constants/keyboard-shortcuts.ts @@ -22,6 +22,8 @@ export const SHORTCUTS = { SPACE_UP: "mod+up", SPACE_DOWN: "mod+down", FIND_IN_CONVERSATION: "mod+f", + FILE_PICKER: "mod+p", + CHECKPOINT_TIMELINE: "mod+shift+h", BLUR: "escape", SUBMIT_BLUR: "mod+enter", } as const; @@ -160,6 +162,22 @@ export const KEYBOARD_SHORTCUTS: KeyboardShortcut[] = [ category: "panels", context: "Task detail", }, + { + id: "file-picker", + keys: SHORTCUTS.FILE_PICKER, + description: "Open file picker", + category: "panels", + context: "Task detail", + configurable: true, + }, + { + id: "checkpoint-timeline", + keys: SHORTCUTS.CHECKPOINT_TIMELINE, + description: "Open checkpoint timeline", + category: "panels", + context: "Task detail", + configurable: true, + }, { id: "paste-as-file", keys: SHORTCUTS.PASTE_AS_FILE, @@ -218,6 +236,58 @@ export const CATEGORY_LABELS: Record = { editor: "Editor", }; +export const CONFIGURABLE_SHORTCUT_IDS = [ + "command-menu", + "new-task", + "settings", + "shortcuts", + "inbox", + "prev-task", + "next-task", + "space-up", + "space-down", + "go-back", + "go-forward", + "toggle-left-sidebar", + "toggle-review-panel", + "close-tab", + "open-in-editor", + "copy-path", + "toggle-focus", + "file-picker", + "checkpoint-timeline", + "paste-as-file", + "prompt-history-prev", + "prompt-history-next", +] as const; + +export type ConfigurableShortcutId = (typeof CONFIGURABLE_SHORTCUT_IDS)[number]; + +export const DEFAULT_KEYBINDINGS: Record = { + "command-menu": SHORTCUTS.COMMAND_MENU, + "new-task": SHORTCUTS.NEW_TASK, + settings: SHORTCUTS.SETTINGS, + shortcuts: SHORTCUTS.SHORTCUTS_SHEET, + inbox: SHORTCUTS.INBOX, + "prev-task": SHORTCUTS.PREV_TASK, + "next-task": SHORTCUTS.NEXT_TASK, + "space-up": SHORTCUTS.SPACE_UP, + "space-down": SHORTCUTS.SPACE_DOWN, + "go-back": SHORTCUTS.GO_BACK, + "go-forward": SHORTCUTS.GO_FORWARD, + "toggle-left-sidebar": SHORTCUTS.TOGGLE_LEFT_SIDEBAR, + "toggle-review-panel": SHORTCUTS.TOGGLE_REVIEW_PANEL, + "close-tab": SHORTCUTS.CLOSE_TAB, + "open-in-editor": SHORTCUTS.OPEN_IN_EDITOR, + "copy-path": SHORTCUTS.COPY_PATH, + "toggle-focus": SHORTCUTS.TOGGLE_FOCUS, + "file-picker": SHORTCUTS.FILE_PICKER, + "checkpoint-timeline": SHORTCUTS.CHECKPOINT_TIMELINE, + "paste-as-file": SHORTCUTS.PASTE_AS_FILE, + "prompt-history-prev": "shift+up", + "prompt-history-next": "shift+down", +}; + export function getShortcutsByCategory(): Record< ShortcutCategory, KeyboardShortcut[] diff --git a/apps/code/src/renderer/features/sessions/components/CheckpointTimelineModal.tsx b/apps/code/src/renderer/features/sessions/components/CheckpointTimelineModal.tsx new file mode 100644 index 0000000000..73d5811f5c --- /dev/null +++ b/apps/code/src/renderer/features/sessions/components/CheckpointTimelineModal.tsx @@ -0,0 +1,173 @@ +import { + ArrowCounterClockwise, + ClockCounterClockwise, +} from "@phosphor-icons/react"; +import { isNotification, POSTHOG_NOTIFICATIONS } from "@posthog/agent"; +import { Button, Dialog, Flex, ScrollArea, Text } from "@radix-ui/themes"; +import type { AcpMessage } from "@shared/types/session-events"; +import { + isJsonRpcNotification, + isJsonRpcRequest, +} from "@shared/types/session-events"; +import { formatRelativeTimeShort } from "@utils/time"; +import { useMemo } from "react"; + +interface CheckpointEntry { + checkpointId: string; + timestamp: number; + userMessageSnippet: string; + turnIndex: number; +} + +function parseTimeline(events: AcpMessage[]): CheckpointEntry[] { + const entries: CheckpointEntry[] = []; + let turnIndex = 0; + let currentUserMessage = ""; + let currentCheckpointId: string | null = null; + let currentCheckpointTs = 0; + + for (const event of events) { + const msg = event.message; + + if (isJsonRpcRequest(msg) && msg.method === "session/prompt") { + if (currentCheckpointId) { + entries.push({ + checkpointId: currentCheckpointId, + timestamp: currentCheckpointTs, + userMessageSnippet: currentUserMessage, + turnIndex, + }); + } + turnIndex++; + const p = msg.params as { + prompt?: Array<{ type: string; text?: string }>; + }; + currentUserMessage = + p?.prompt + ?.filter((b) => b.type === "text") + .map((b) => b.text ?? "") + .join("") ?? ""; + currentCheckpointId = null; + currentCheckpointTs = 0; + } + + if ( + isJsonRpcNotification(msg) && + isNotification(msg.method, POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT) + ) { + const params = msg.params as { checkpointId?: string }; + if (params?.checkpointId) { + currentCheckpointId = params.checkpointId; + currentCheckpointTs = event.ts; + } + } + } + + if (currentCheckpointId) { + entries.push({ + checkpointId: currentCheckpointId, + timestamp: currentCheckpointTs, + userMessageSnippet: currentUserMessage, + turnIndex, + }); + } + + return entries.reverse(); +} + +interface CheckpointTimelineModalProps { + open: boolean; + onOpenChange: (open: boolean) => void; + events: AcpMessage[]; + onRestore: (checkpointId: string) => void; +} + +export function CheckpointTimelineModal({ + open, + onOpenChange, + events, + onRestore, +}: CheckpointTimelineModalProps) { + const entries = useMemo(() => parseTimeline(events), [events]); + + return ( + + + + + + + Checkpoint timeline + + + {entries.length === 0 ? ( + + No checkpoints in this session yet. + + ) : ( + + + {entries.map((entry) => ( + { + onOpenChange(false); + onRestore(entry.checkpointId); + }} + /> + ))} + + + )} + + + + + + + + + ); +} + +function CheckpointRow({ + entry, + onRestore, +}: { + entry: CheckpointEntry; + onRestore: () => void; +}) { + const snippet = entry.userMessageSnippet.trim().slice(0, 120); + const label = snippet || `Turn ${entry.turnIndex}`; + + return ( + + + + {label} + + + {formatRelativeTimeShort(entry.timestamp)} + + + + + ); +} diff --git a/apps/code/src/renderer/features/sessions/components/ConversationView.tsx b/apps/code/src/renderer/features/sessions/components/ConversationView.tsx index 4afb50fd67..60d74d4fd2 100644 --- a/apps/code/src/renderer/features/sessions/components/ConversationView.tsx +++ b/apps/code/src/renderer/features/sessions/components/ConversationView.tsx @@ -1,6 +1,7 @@ import { CHAT_CONTENT_MAX_WIDTH } from "@features/sessions/constants"; import { useContextUsage } from "@features/sessions/hooks/useContextUsage"; import { useConversationSearch } from "@features/sessions/hooks/useConversationSearch"; +import { useRestoreCheckpoint } from "@features/sessions/hooks/useRestoreCheckpoint"; import { SessionTaskIdProvider } from "@features/sessions/hooks/useSessionTaskId"; import { sessionStoreSetters, @@ -11,6 +12,7 @@ import { } from "@features/sessions/stores/sessionStore"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { SkillButtonActionMessage } from "@features/skill-buttons/components/SkillButtonActionMessage"; +import { useShortcut } from "@hooks/useShortcut"; import { ArrowDown, XCircle } from "@phosphor-icons/react"; import { WorkerPoolContextProvider } from "@pierre/diffs/react"; import WorkerUrl from "@pierre/diffs/worker/worker.js?worker&url"; @@ -18,15 +20,18 @@ import { Box, Button, Flex, Text } from "@radix-ui/themes"; import type { Task } from "@shared/types"; import type { AcpMessage } from "@shared/types/session-events"; import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { useHotkeys } from "react-hotkeys-hook"; import { buildConversationItems, type ConversationItem, type TurnContext, } from "./buildConversationItems"; +import { CheckpointTimelineModal } from "./CheckpointTimelineModal"; import { ConversationSearchBar } from "./ConversationSearchBar"; import { GitActionMessage } from "./GitActionMessage"; import { GitActionResult } from "./GitActionResult"; import { mergeConversationItems } from "./mergeConversationItems"; +import { RestoreCheckpointDialog } from "./RestoreCheckpointDialog"; import { SessionFooter } from "./SessionFooter"; import { QueuedMessageView } from "./session-update/QueuedMessageView"; import { @@ -74,6 +79,7 @@ export function ConversationView({ const listRef = useRef(null); const isAtBottomRef = useRef(true); const [showScrollButton, setShowScrollButton] = useState(false); + const [timelineOpen, setTimelineOpen] = useState(false); const debugLogsCloudRuns = useSettingsStore((s) => s.debugLogsCloudRuns); const showDebugLogs = debugLogsCloudRuns; const contextUsage = useContextUsage(events); @@ -126,6 +132,17 @@ export function ConversationView({ const isCloud = session?.isCloud ?? false; + const restore = useRestoreCheckpoint({ + repoPath: repoPath ?? undefined, + taskId, + taskRunId: session?.taskRunId, + }); + + const checkpointTimelineKey = useShortcut("checkpoint-timeline"); + useHotkeys(checkpointTimelineKey, () => setTimelineOpen((o) => !o), { + preventDefault: true, + }); + const items = useMemo( () => mergeConversationItems({ @@ -209,6 +226,16 @@ export function ConversationView({ update={item.update} turnContext={item.turnContext} thoughtComplete={item.thoughtComplete} + showRestoreButton={item.turnContext.turnComplete} + onRestoreCheckpoint={ + item.turnContext.turnComplete && + item.turnContext.lastCheckpointId + ? () => + restore.requestRestore( + item.turnContext.lastCheckpointId as string, + ) + : undefined + } /> ); case "git_action_result": @@ -240,7 +267,14 @@ export function ConversationView({ ); } }, - [repoPath, taskId, slackThreadUrl, firstUserMessageId, initialItemIds], + [ + repoPath, + taskId, + slackThreadUrl, + firstUserMessageId, + initialItemIds, + restore.requestRestore, + ], ); const getItemKey = useCallback((item: ConversationItem) => item.id, []); @@ -310,6 +344,18 @@ export function ConversationView({ )} + + ); } @@ -318,10 +364,14 @@ const SessionUpdateRow = memo(function SessionUpdateRow({ update, turnContext, thoughtComplete, + showRestoreButton, + onRestoreCheckpoint, }: { update: RenderItem; turnContext: TurnContext; thoughtComplete?: boolean; + showRestoreButton?: boolean; + onRestoreCheckpoint?: () => void; }) { return ( ); }); diff --git a/apps/code/src/renderer/features/sessions/components/RestoreCheckpointDialog.tsx b/apps/code/src/renderer/features/sessions/components/RestoreCheckpointDialog.tsx new file mode 100644 index 0000000000..164aa7269a --- /dev/null +++ b/apps/code/src/renderer/features/sessions/components/RestoreCheckpointDialog.tsx @@ -0,0 +1,54 @@ +import { ArrowCounterClockwise, Warning } from "@phosphor-icons/react"; +import { Button, Dialog, Flex, Text } from "@radix-ui/themes"; + +interface RestoreCheckpointDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + onConfirm: () => void; + isLoading: boolean; +} + +export function RestoreCheckpointDialog({ + open, + onOpenChange, + onConfirm, + isLoading, +}: RestoreCheckpointDialogProps) { + return ( + + + + + + + Restore checkpoint + + + + + + This will revert all file changes made after this point. This + action cannot be undone. + + + + + + + + + + + + ); +} diff --git a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts index fbd0d1ee4b..31d396a6b8 100644 --- a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts +++ b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts @@ -28,6 +28,7 @@ export interface TurnContext { childItems: Map; turnCancelled: boolean; turnComplete: boolean; + lastCheckpointId: string | null; } export type ConversationItem = @@ -92,6 +93,7 @@ interface TurnState { context: TurnContext; gitAction: ReturnType; itemCount: number; + lastCheckpointId: string | null; } interface ItemBuilder { @@ -248,6 +250,7 @@ function handlePromptRequest( childItems, turnCancelled: false, turnComplete: false, + lastCheckpointId: null, }; b.currentTurn = { @@ -259,6 +262,7 @@ function handlePromptRequest( context, gitAction, itemCount: 0, + lastCheckpointId: null, }; b.pendingPrompts.set(msg.id, b.currentTurn); @@ -439,6 +443,15 @@ function handleNotification( }); return; } + + if (isNotification(msg.method, POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT)) { + const params = msg.params as { checkpointId?: string }; + if (params?.checkpointId && b.currentTurn) { + b.currentTurn.lastCheckpointId = params.checkpointId; + b.currentTurn.context.lastCheckpointId = params.checkpointId; + } + return; + } } function ensureProgressCardForGroup( @@ -533,6 +546,7 @@ function ensureImplicitTurn(b: ItemBuilder, ts: number) { childItems, turnCancelled: false, turnComplete: false, + lastCheckpointId: null, }; b.currentTurn = { @@ -544,6 +558,7 @@ function ensureImplicitTurn(b: ItemBuilder, ts: number) { context, gitAction: { isGitAction: false, actionType: null, prompt: "" }, itemCount: 0, + lastCheckpointId: null, }; } diff --git a/apps/code/src/renderer/features/sessions/components/session-update/AgentMessage.tsx b/apps/code/src/renderer/features/sessions/components/session-update/AgentMessage.tsx index 0cf7d00083..dae3716ae6 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/AgentMessage.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/AgentMessage.tsx @@ -7,8 +7,8 @@ import { useSessionTaskId } from "@features/sessions/hooks/useSessionTaskId"; import { useCwd } from "@features/sidebar/hooks/useCwd"; import type { FileItem } from "@hooks/useRepoFiles"; import { useRepoFiles } from "@hooks/useRepoFiles"; -import { Check, Copy } from "@phosphor-icons/react"; -import { Box, Code, IconButton } from "@radix-ui/themes"; +import { ArrowCounterClockwise, Check, Copy } from "@phosphor-icons/react"; +import { Box, Code, Flex, IconButton } from "@radix-ui/themes"; import { memo, useCallback, useMemo, useState } from "react"; import type { Components } from "react-markdown"; @@ -135,12 +135,17 @@ const agentComponents: Partial = { interface AgentMessageProps { content: string; + showRestoreButton?: boolean; + onRestoreCheckpoint?: () => void; } export const AgentMessage = memo(function AgentMessage({ content, + showRestoreButton, + onRestoreCheckpoint, }: AgentMessageProps) { const [copied, setCopied] = useState(false); + const canRestore = !!onRestoreCheckpoint; const handleCopy = useCallback(() => { navigator.clipboard.writeText(content); @@ -154,7 +159,10 @@ export const AgentMessage = memo(function AgentMessage({ content={content} componentsOverride={agentComponents} /> - + : } - + {showRestoreButton && ( + + + + + + )} + ); }); diff --git a/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx b/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx index 90eebd85bf..9563e1f488 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx @@ -57,6 +57,8 @@ interface SessionUpdateViewProps { turnCancelled?: boolean; turnComplete?: boolean; thoughtComplete?: boolean; + showRestoreButton?: boolean; + onRestoreCheckpoint?: () => void; } export const SessionUpdateView = memo(function SessionUpdateView({ @@ -66,13 +68,19 @@ export const SessionUpdateView = memo(function SessionUpdateView({ turnCancelled, turnComplete, thoughtComplete, + showRestoreButton, + onRestoreCheckpoint, }: SessionUpdateViewProps) { switch (item.sessionUpdate) { case "user_message_chunk": return null; case "agent_message_chunk": return item.content.type === "text" ? ( - + ) : null; case "agent_thought_chunk": return item.content.type === "text" ? ( diff --git a/apps/code/src/renderer/features/sessions/hooks/useRestoreCheckpoint.ts b/apps/code/src/renderer/features/sessions/hooks/useRestoreCheckpoint.ts new file mode 100644 index 0000000000..a3895ecaf0 --- /dev/null +++ b/apps/code/src/renderer/features/sessions/hooks/useRestoreCheckpoint.ts @@ -0,0 +1,69 @@ +import { sessionStoreSetters } from "@features/sessions/stores/sessionStore"; +import { trpcClient } from "@renderer/trpc"; +import { useCallback, useState } from "react"; +import { toast } from "sonner"; + +interface UseRestoreCheckpointOptions { + repoPath: string | undefined; + taskId: string | undefined; + taskRunId: string | undefined; +} + +export function useRestoreCheckpoint({ + repoPath, + taskId, + taskRunId, +}: UseRestoreCheckpointOptions) { + const [dialogOpen, setDialogOpen] = useState(false); + const [pendingCheckpointId, setPendingCheckpointId] = useState( + null, + ); + const [isRestoring, setIsRestoring] = useState(false); + + const requestRestore = useCallback((checkpointId: string) => { + setPendingCheckpointId(checkpointId); + setDialogOpen(true); + }, []); + + const confirmRestore = useCallback(async () => { + if (!pendingCheckpointId || !repoPath) return; + + setIsRestoring(true); + try { + await trpcClient.checkpoint.restore.mutate({ + checkpointId: pendingCheckpointId, + repoPath, + taskRunId, + }); + if (taskId) { + sessionStoreSetters.truncateEventsToCheckpoint( + taskId, + pendingCheckpointId, + ); + } + toast.success("Checkpoint restored successfully"); + setDialogOpen(false); + setPendingCheckpointId(null); + } catch (error) { + const message = + error instanceof Error ? error.message : "Failed to restore checkpoint"; + toast.error(message); + } finally { + setIsRestoring(false); + } + }, [pendingCheckpointId, repoPath, taskId, taskRunId]); + + const cancelRestore = useCallback(() => { + setDialogOpen(false); + setPendingCheckpointId(null); + }, []); + + return { + dialogOpen, + setDialogOpen, + isRestoring, + requestRestore, + confirmRestore, + cancelRestore, + }; +} diff --git a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts index 718206228b..da27e152dc 100644 --- a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts +++ b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts @@ -5,9 +5,14 @@ import type { SessionConfigSelectOption, SessionConfigSelectOptions, } from "@agentclientprotocol/sdk"; +import { isNotification, POSTHOG_NOTIFICATIONS } from "@posthog/agent"; import type { ExecutionMode, TaskRunStatus } from "@shared/types"; import type { SkillButtonId } from "@shared/types/analytics"; -import type { AcpMessage } from "@shared/types/session-events"; +import { + type AcpMessage, + isJsonRpcNotification, + isJsonRpcRequest, +} from "@shared/types/session-events"; import { create } from "zustand"; import { immer } from "zustand/middleware/immer"; import type { PermissionRequest } from "../utils/parseSessionLogs"; @@ -498,6 +503,51 @@ export const sessionStoreSetters = { return useSessionStore.getState().sessions; }, + truncateEventsToCheckpoint: ( + taskId: string, + checkpointId: string, + ): boolean => { + const state = useSessionStore.getState(); + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return false; + const session = state.sessions[taskRunId]; + if (!session) return false; + + const events = session.events; + let checkpointEventIdx = -1; + for (let i = 0; i < events.length; i++) { + const msg = events[i].message; + if (!isJsonRpcNotification(msg)) continue; + if (!isNotification(msg.method, POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT)) + continue; + const params = msg.params as { checkpointId?: string } | undefined; + if (params?.checkpointId === checkpointId) { + checkpointEventIdx = i; + break; + } + } + if (checkpointEventIdx === -1) return false; + + let cutoff = events.length; + for (let i = checkpointEventIdx + 1; i < events.length; i++) { + const msg = events[i].message; + if (isJsonRpcRequest(msg) && msg.method === "session/prompt") { + cutoff = i; + break; + } + } + + useSessionStore.setState((draft) => { + const trid = draft.taskIdIndex[taskId]; + if (!trid) return; + const s = draft.sessions[trid]; + if (s) { + s.events = s.events.slice(0, cutoff); + } + }); + return true; + }, + clearAll: () => { useSessionStore.setState((state) => { state.sessions = {}; From 84a2fec05525ec327121b51c576cc61ae6cea9b0 Mon Sep 17 00:00:00 2001 From: Basit Balogun Date: Mon, 1 Jun 2026 12:59:39 +0100 Subject: [PATCH 2/3] fix: capture git checkpoint for local tasks after each turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cloud path (agent-server.ts) guards checkpoint capture on posthogAPI being configured, so local tasks never emit _posthog/git_checkpoint. Hook into extNotification in the local AgentService: on TURN_COMPLETE, run CaptureCheckpointSaga, then emit a synthetic _posthog/git_checkpoint ACP message to the renderer and append it to the session JSONL so it survives reload. The renderer's buildConversationItems already handles the notification correctly — it just wasn't arriving. Add console logs in buildConversationItems and structured logs in service.ts for visibility during debugging. --- apps/code/src/main/services/agent/service.ts | 110 +++++++++++++++++- .../components/buildConversationItems.ts | 11 ++ 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index d86221b9e5..f10cb940dd 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -1,4 +1,4 @@ -import fs, { mkdirSync, symlinkSync } from "node:fs"; +import fs, { mkdirSync, promises as fsPromises, symlinkSync } from "node:fs"; import { homedir, tmpdir } from "node:os"; import { isAbsolute, join, relative, resolve, sep } from "node:path"; import { @@ -18,7 +18,10 @@ import { POSTHOG_NOTIFICATIONS, } from "@posthog/agent"; import type { McpToolApprovals } from "@posthog/agent/adapters/claude/mcp/tool-metadata"; -import { hydrateSessionJsonl } from "@posthog/agent/adapters/claude/session/jsonl-hydration"; +import { + getSessionJsonlPath, + hydrateSessionJsonl, +} from "@posthog/agent/adapters/claude/session/jsonl-hydration"; import { getReasoningEffortOptions } from "@posthog/agent/adapters/reasoning-effort"; import { Agent } from "@posthog/agent/agent"; import { @@ -38,6 +41,7 @@ import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; import { extractCreatedPrUrl } from "@posthog/agent/pr-url-detector"; import type * as AgentTypes from "@posthog/agent/types"; import { getCurrentBranch } from "@posthog/git/queries"; +import { CaptureCheckpointSaga } from "@posthog/git/sagas/checkpoint"; import type { IAppMeta } from "@posthog/platform/app-meta"; import type { IBundledResources } from "@posthog/platform/bundled-resources"; import type { IPowerManager } from "@posthog/platform/power-manager"; @@ -1493,6 +1497,33 @@ For git operations while detached: } } + if (isNotification(method, POSTHOG_NOTIFICATIONS.TURN_COMPLETE)) { + const turnSession = service.sessions.get(taskRunId); + if (turnSession?.config.repoPath) { + log.debug("TURN_COMPLETE — capturing local checkpoint", { + taskRunId, + repoPath: turnSession.config.repoPath, + }); + service + .captureLocalCheckpoint( + taskRunId, + turnSession.config.repoPath, + turnSession.config.sessionId, + emitToRenderer, + ) + .catch((err) => { + log.warn("Local checkpoint capture failed", { + taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + }); + } else { + log.debug("TURN_COMPLETE — no repoPath, skipping checkpoint", { + taskRunId, + }); + } + } + if (isNotification(method, POSTHOG_NOTIFICATIONS.USAGE_UPDATE)) { this.emit(AgentServiceEvent.LlmActivity, undefined); } @@ -1737,6 +1768,81 @@ For git operations while detached: }); } + /** + * Capture a local git checkpoint after a turn completes, emit the + * `_posthog/git_checkpoint` notification to the renderer, and append it to + * the session JSONL so it survives page reload. + */ + private async captureLocalCheckpoint( + taskRunId: string, + repoPath: string, + sessionId: string | undefined, + emitToRenderer: (payload: unknown) => void, + ): Promise { + log.info("Capturing local checkpoint after turn", { taskRunId, repoPath }); + + const saga = new CaptureCheckpointSaga(); + let result: Awaited>; + try { + result = await saga.execute({ baseDir: repoPath }); + } catch (err) { + log.warn("CaptureCheckpointSaga failed — no checkpoint for this turn", { + taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + return; + } + + log.info("Local checkpoint captured", { + taskRunId, + checkpointId: result.checkpointId, + commit: result.commit, + branch: result.branch, + }); + + const notification = { + jsonrpc: "2.0" as const, + method: POSTHOG_NOTIFICATIONS.GIT_CHECKPOINT, + params: { checkpointId: result.checkpointId }, + }; + + // Emit to renderer so the restore button activates on the completed turn + const acpMessage: AcpMessage = { + type: "acp_message", + ts: Date.now(), + message: notification as AcpMessage["message"], + }; + emitToRenderer(acpMessage); + + log.info("Emitted GIT_CHECKPOINT notification to renderer", { + taskRunId, + checkpointId: result.checkpointId, + }); + + // Append to the session JSONL so restore can find the checkpoint on reload + if (sessionId) { + try { + const jsonlPath = getSessionJsonlPath(sessionId, repoPath); + const line = `${JSON.stringify({ notification })}\n`; + await fsPromises.appendFile(jsonlPath, line, "utf-8"); + log.info("Checkpoint appended to JSONL", { + taskRunId, + checkpointId: result.checkpointId, + jsonlPath, + }); + } catch (err) { + log.warn("Failed to append checkpoint to JSONL (restore may not survive reload)", { + taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + } + } else { + log.warn("No sessionId yet — checkpoint not written to JSONL", { + taskRunId, + }); + } + } + async getGatewayModels(apiHost: string) { const gatewayUrl = getLlmGatewayUrl(apiHost); const models = await fetchGatewayModels({ gatewayUrl }); diff --git a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts index 31d396a6b8..4f1cf10012 100644 --- a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts +++ b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts @@ -449,6 +449,17 @@ function handleNotification( if (params?.checkpointId && b.currentTurn) { b.currentTurn.lastCheckpointId = params.checkpointId; b.currentTurn.context.lastCheckpointId = params.checkpointId; + console.log("[checkpoint] GIT_CHECKPOINT set on turn", { + checkpointId: params.checkpointId, + turnId: b.currentTurn.id, + turnComplete: b.currentTurn.isComplete, + }); + } else { + console.warn("[checkpoint] GIT_CHECKPOINT received but could not be applied", { + checkpointId: params?.checkpointId, + hasCurrentTurn: !!b.currentTurn, + currentTurnId: b.currentTurn?.id, + }); } return; } From 158a49ca17536829869909e9024ad6bb7eabe758 Mon Sep 17 00:00:00 2001 From: Basit Balogun Date: Mon, 1 Jun 2026 13:19:44 +0100 Subject: [PATCH 3/3] fix: intercept TURN_COMPLETE in raw stream tap, not extNotification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit extNotification is not reliably called by the ACP SDK for _posthog/ notifications in the local path. The raw stream tap (onAcpMessage) is guaranteed to fire for every ndjson frame — move the TURN_COMPLETE checkpoint hook there instead. --- apps/code/src/main/services/agent/service.ts | 66 ++++++++++++-------- 1 file changed, 39 insertions(+), 27 deletions(-) diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index f10cb940dd..461152bd71 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -1285,6 +1285,11 @@ For git operations while detached: // Inspect tool call updates for PR URLs and file activity this.handleToolCallUpdate(taskRunId, message as AcpMessage["message"]); + + // Capture a local git checkpoint when a turn completes. + // Intercepted here (raw stream tap) rather than extNotification because + // the ACP SDK does not reliably route _posthog/ notifications to that callback. + this.handleTurnCompleteForCheckpoint(taskRunId, message, emitToRenderer); }; const tappedReadable = createTappedReadableStream( @@ -1497,33 +1502,6 @@ For git operations while detached: } } - if (isNotification(method, POSTHOG_NOTIFICATIONS.TURN_COMPLETE)) { - const turnSession = service.sessions.get(taskRunId); - if (turnSession?.config.repoPath) { - log.debug("TURN_COMPLETE — capturing local checkpoint", { - taskRunId, - repoPath: turnSession.config.repoPath, - }); - service - .captureLocalCheckpoint( - taskRunId, - turnSession.config.repoPath, - turnSession.config.sessionId, - emitToRenderer, - ) - .catch((err) => { - log.warn("Local checkpoint capture failed", { - taskRunId, - error: err instanceof Error ? err.message : String(err), - }); - }); - } else { - log.debug("TURN_COMPLETE — no repoPath, skipping checkpoint", { - taskRunId, - }); - } - } - if (isNotification(method, POSTHOG_NOTIFICATIONS.USAGE_UPDATE)) { this.emit(AgentServiceEvent.LlmActivity, undefined); } @@ -1768,6 +1746,40 @@ For git operations while detached: }); } + private handleTurnCompleteForCheckpoint( + taskRunId: string, + message: unknown, + emitToRenderer: (payload: unknown) => void, + ): void { + const msg = message as { method?: string }; + if (!isNotification(msg.method, POSTHOG_NOTIFICATIONS.TURN_COMPLETE)) return; + + const session = this.sessions.get(taskRunId); + if (!session?.config.repoPath) { + log.debug("TURN_COMPLETE in stream — no repoPath, skipping checkpoint", { + taskRunId, + }); + return; + } + + log.info("TURN_COMPLETE in stream — capturing local checkpoint", { + taskRunId, + repoPath: session.config.repoPath, + }); + + this.captureLocalCheckpoint( + taskRunId, + session.config.repoPath, + session.config.sessionId, + emitToRenderer, + ).catch((err) => { + log.warn("Local checkpoint capture failed", { + taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + }); + } + /** * Capture a local git checkpoint after a turn completes, emit the * `_posthog/git_checkpoint` notification to the renderer, and append it to