feat: restore checkpoint after any AI response turn

apps/code/src/main/services/agent/service.ts

@@ -954,6 +954,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,

apps/code/src/main/trpc/router.ts

@@ -1,5 +1,6 @@
 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";
@@ -42,6 +43,7 @@ export const trpcRouter = router({
   analytics: analyticsRouter,
   archive: archiveRouter,
   auth: authRouter,
+  checkpoint: checkpointRouter,
   cloudTask: cloudTaskRouter,
   connectivity: connectivityRouter,
   contextMenu: contextMenuRouter,

apps/code/src/main/trpc/routers/checkpoint.ts

@@ -0,0 +1,216 @@
+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<AgentService>(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<TruncateResult> {
+  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<string, unknown>
+        | 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.execute({
+        baseDir: input.repoPath,
+        checkpointId: input.checkpointId,
+      });
+
+      // 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;
+    }),
+});

apps/code/src/renderer/components/GlobalEventHandlers.tsx

@@ -9,8 +9,10 @@ import { useSidebarStore } from "@features/sidebar/stores/sidebarStore";
 import { useTasks } from "@features/tasks/hooks/useTasks";
 import { useFocusWorkspace } from "@features/workspace/hooks/useFocusWorkspace";
 import { useWorkspaces } from "@features/workspace/hooks/useWorkspace";
+import { useShortcut } from "@hooks/useShortcut";
 import { SHORTCUTS } from "@renderer/constants/keyboard-shortcuts";
 import { useTRPC } from "@renderer/trpc";
+import { isMac } from "@renderer/utils/platform";
 import type { Task } from "@shared/types";
 import { useCommandMenuStore } from "@stores/commandMenuStore";
 import { useNavigationStore } from "@stores/navigationStore";
@@ -157,33 +159,43 @@ export function GlobalEventHandlers({
     preventDefault: true,
   } as const;
 
-  useHotkeys(SHORTCUTS.COMMAND_MENU, onToggleCommandMenu, {
+  const commandMenuKey = useShortcut("command-menu");
+  const newTaskKey = useShortcut("new-task");
+  const settingsKey = useShortcut("settings");
+  const goBackKey = useShortcut("go-back");
+  const goForwardKey = useShortcut("go-forward");
+  const toggleLeftSidebarKey = useShortcut("toggle-left-sidebar");
+  const toggleReviewPanelKey = useShortcut("toggle-review-panel");
+  const shortcutsSheetKey = useShortcut("shortcuts");
+  const inboxKey = useShortcut("inbox");
+  const prevTaskKey = useShortcut("prev-task");
+  const nextTaskKey = useShortcut("next-task");
+  const toggleFocusKey = useShortcut("toggle-focus");
+
+  useHotkeys(commandMenuKey, onToggleCommandMenu, {
     ...globalOptions,
     enabled: !commandMenuOpen,
   });
-  useHotkeys(SHORTCUTS.NEW_TASK, handleFocusTaskMode, globalOptions);
-  useHotkeys(SHORTCUTS.SETTINGS, handleOpenSettings, globalOptions);
-  useHotkeys(SHORTCUTS.GO_BACK, goBack, globalOptions);
-  useHotkeys(SHORTCUTS.GO_FORWARD, goForward, globalOptions);
+  useHotkeys(newTaskKey, handleFocusTaskMode, globalOptions);
+  useHotkeys(settingsKey, handleOpenSettings, globalOptions);
+  useHotkeys(goBackKey, goBack, globalOptions);
+  useHotkeys(goForwardKey, goForward, globalOptions);
+
   const handleToggleReview = useCallback(() => {
     if (!currentTaskId) return;
     const mode = getReviewMode(currentTaskId);
     setReviewMode(currentTaskId, mode === "closed" ? "split" : "closed");
   }, [currentTaskId, getReviewMode, setReviewMode]);
 
-  useHotkeys(SHORTCUTS.TOGGLE_LEFT_SIDEBAR, toggleLeftSidebar, globalOptions);
-  useHotkeys(SHORTCUTS.TOGGLE_REVIEW_PANEL, handleToggleReview, globalOptions);
-  useHotkeys(SHORTCUTS.SHORTCUTS_SHEET, onToggleShortcutsSheet, globalOptions);
-  useHotkeys(SHORTCUTS.INBOX, navigateToInbox, globalOptions);
-  useHotkeys(SHORTCUTS.PREV_TASK, handlePrevTask, globalOptions, [
-    handlePrevTask,
-  ]);
-  useHotkeys(SHORTCUTS.NEXT_TASK, handleNextTask, globalOptions, [
-    handleNextTask,
-  ]);
+  useHotkeys(toggleLeftSidebarKey, toggleLeftSidebar, globalOptions);
+  useHotkeys(toggleReviewPanelKey, handleToggleReview, globalOptions);
+  useHotkeys(shortcutsSheetKey, onToggleShortcutsSheet, globalOptions);
+  useHotkeys(inboxKey, navigateToInbox, globalOptions);
+  useHotkeys(prevTaskKey, handlePrevTask, globalOptions, [handlePrevTask]);
+  useHotkeys(nextTaskKey, handleNextTask, globalOptions, [handleNextTask]);
 
   useHotkeys(
-    SHORTCUTS.TOGGLE_FOCUS,
+    toggleFocusKey,
     handleToggleFocus,
     {
       ...globalOptions,
@@ -192,11 +204,13 @@ export function GlobalEventHandlers({
     [handleToggleFocus],
   );
 
-  // Task switching with mod+1-9
+  // Task switching with mod+1-9. On macOS, Ctrl+1..9 is reserved for
+  // SWITCH_TAB (panel tabs), so ignore plain-Ctrl there; on Windows/Linux,
+  // Ctrl IS mod, so the same event must trigger task switching.
   useHotkeys(
     SHORTCUTS.SWITCH_TASK,
     (event, handler) => {
-      if (event.ctrlKey && !event.metaKey) return;
+      if (isMac && event.ctrlKey && !event.metaKey) return;
 
       const keyPressed = handler.keys?.[0];
       if (!keyPressed) return;