feat: restore checkpoint after any AI response turn

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";
@@ -993,6 +997,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,
@@ -1273,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(
@@ -1729,6 +1746,115 @@ 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
+   * the session JSONL so it survives page reload.
+   */
+  private async captureLocalCheckpoint(
+    taskRunId: string,
+    repoPath: string,
+    sessionId: string | undefined,
+    emitToRenderer: (payload: unknown) => void,
+  ): Promise<void> {
+    log.info("Capturing local checkpoint after turn", { taskRunId, repoPath });
+
+    const saga = new CaptureCheckpointSaga();
+    let result: Awaited<ReturnType<typeof saga.execute>>;
+    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 });

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,