feat: remove tree snapshots, use checkpoints, fix handoff bugs

Author: jonathanlab

apps/code/src/main/services/cloud-task/service.test.ts

@@ -446,8 +446,9 @@ describe("CloudTaskService", () => {
     );
 
     expect(mockStreamFetch.mock.calls.length).toBe(6);
-    // 2 bootstrap calls + 6 handleStreamCompletion calls (one per stream error)
-    expect(mockNetFetch).toHaveBeenCalledTimes(8);
+    // 2 bootstrap calls + 1 post-bootstrap status verification + 6
+    // handleStreamCompletion calls (one per stream error)
+    expect(mockNetFetch).toHaveBeenCalledTimes(9);
     expect(updates).toContainEqual({
       taskId: "task-1",
       runId: "run-1",

apps/code/src/main/services/cloud-task/service.ts

@@ -548,6 +548,33 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
       watcher.needsPostBootstrapReconnect = false;
       this.scheduleReconnect(key);
     }
+
+    void this.verifyPostBootstrapStatus(key);
+  }
+
+  private async verifyPostBootstrapStatus(key: string): Promise<void> {
+    const watcher = this.watchers.get(key);
+    if (!watcher) return;
+    if (isTerminalStatus(watcher.lastStatus)) return;
+
+    const run = await this.fetchTaskRun(watcher);
+    const currentWatcher = this.watchers.get(key);
+    if (!currentWatcher || currentWatcher !== watcher) return;
+    if (!run) return;
+
+    if (!this.applyTaskRunState(watcher, run)) return;
+    if (isTerminalStatus(watcher.lastStatus)) return;
+
+    this.emit(CloudTaskEvent.Update, {
+      taskId: watcher.taskId,
+      runId: watcher.runId,
+      kind: "status",
+      status: watcher.lastStatus ?? undefined,
+      stage: watcher.lastStage,
+      output: watcher.lastOutput,
+      errorMessage: watcher.lastErrorMessage,
+      branch: watcher.lastBranch,
+    });
   }
 
   private async connectSse(

apps/code/src/main/services/handoff/handoff-saga.test.ts

@@ -33,19 +33,6 @@ function createInput(
   };
 }
 
-function createSnapshot(
-  overrides: Partial<AgentTypes.TreeSnapshotEvent> = {},
-): AgentTypes.TreeSnapshotEvent {
-  return {
-    treeHash: "abc123",
-    baseCommit: "def456",
-    archiveUrl: "https://s3.example.com/archive.tar.gz",
-    changes: [{ path: "test.txt", status: "A" }],
-    timestamp: "2026-04-07T00:00:00Z",
-    ...overrides,
-  };
-}
-
 function createCheckpoint(
   overrides: Partial<AgentTypes.GitCheckpointEvent> = {},
 ): AgentTypes.GitCheckpointEvent {
@@ -75,7 +62,6 @@ function createDeps(overrides: Partial<HandoffSagaDeps> = {}): HandoffSagaDeps {
       }),
       updateTaskRun: vi.fn().mockResolvedValue({}),
     }),
-    applyTreeSnapshot: vi.fn().mockResolvedValue(undefined),
     applyGitCheckpoint: vi.fn().mockResolvedValue(undefined),
     updateWorkspaceMode: vi.fn(),
     reconnectSession: vi.fn().mockResolvedValue({
@@ -96,7 +82,6 @@ function createResumeState(
 ): AgentResume.ResumeState {
   return {
     conversation: [],
-    latestSnapshot: null,
     latestGitCheckpoint: null,
     interrupted: false,
     logEntryCount: 0,
@@ -130,22 +115,22 @@ describe("HandoffSaga", () => {
     mockFormatConversation.mockReturnValue("conversation summary");
   });
 
-  it("completes happy path with snapshot", async () => {
-    const snapshot = createSnapshot();
+  it("completes happy path with checkpoint", async () => {
+    const checkpoint = createCheckpoint();
     const { result } = await runSaga({
       resumeState: {
         conversation: [
           { role: "user", content: [{ type: "text", text: "hello" }] },
         ],
-        latestSnapshot: snapshot,
+        latestGitCheckpoint: checkpoint,
         logEntryCount: 10,
       },
     });
 
     expect(result.success).toBe(true);
     if (!result.success) return;
     expect(result.data.sessionId).toBe("session-1");
-    expect(result.data.snapshotApplied).toBe(true);
+    expect(result.data.checkpointApplied).toBe(true);
     expect(result.data.conversationTurns).toBe(1);
   });
 
@@ -165,27 +150,15 @@ describe("HandoffSaga", () => {
     expect(closeOrder).toBeLessThan(fetchOrder);
   });
 
-  it("skips snapshot apply when no archiveUrl", async () => {
+  it("skips checkpoint apply when no checkpoint is present", async () => {
     const { deps, result } = await runSaga({
-      resumeState: {
-        latestSnapshot: createSnapshot({ archiveUrl: undefined }),
-        logEntryCount: 5,
-      },
+      resumeState: { logEntryCount: 5 },
     });
 
     expect(result.success).toBe(true);
     if (!result.success) return;
-    expect(result.data.snapshotApplied).toBe(false);
-    expect(deps.applyTreeSnapshot).not.toHaveBeenCalled();
-  });
-
-  it("skips snapshot apply when no snapshot at all", async () => {
-    const { deps, result } = await runSaga();
-
-    expect(result.success).toBe(true);
-    if (!result.success) return;
-    expect(result.data.snapshotApplied).toBe(false);
-    expect(deps.applyTreeSnapshot).not.toHaveBeenCalled();
+    expect(result.data.checkpointApplied).toBe(false);
+    expect(deps.applyGitCheckpoint).not.toHaveBeenCalled();
   });
 
   it("seeds local logs when cloudLogUrl is present", async () => {
@@ -232,16 +205,16 @@ describe("HandoffSaga", () => {
     );
   });
 
-  it("context mentions files restored when snapshot applied", async () => {
+  it("context mentions files restored when checkpoint applied", async () => {
     const { deps } = await runSaga({
       resumeState: {
-        latestSnapshot: createSnapshot(),
+        latestGitCheckpoint: createCheckpoint(),
       },
     });
 
     expect(deps.setPendingContext).toHaveBeenCalledWith(
       "run-1",
-      expect.stringContaining("fully restored"),
+      expect.stringContaining("restored from the cloud session checkpoint"),
     );
   });
 
@@ -262,14 +235,12 @@ describe("HandoffSaga", () => {
     const { deps } = await runSaga({
       resumeState: {
         latestGitCheckpoint: createCheckpoint(),
-        latestSnapshot: createSnapshot(),
       },
     });
 
     expect(getProgressSteps(deps)).toEqual([
       "fetching_logs",
       "applying_git_checkpoint",
-      "applying_snapshot",
       "spawning_agent",
       "complete",
     ]);
@@ -318,11 +289,10 @@ describe("HandoffSaga", () => {
     });
   });
 
-  it("applies git checkpoint before restoring the file snapshot", async () => {
+  it("applies git checkpoint with local git state during handoff", async () => {
     const { deps, result } = await runSaga({
       input: { localGitState: DEFAULT_LOCAL_GIT_STATE },
       resumeState: {
-        latestSnapshot: createSnapshot(),
         latestGitCheckpoint: createCheckpoint(),
       },
     });
@@ -338,6 +308,5 @@ describe("HandoffSaga", () => {
       expect.any(Object),
       DEFAULT_LOCAL_GIT_STATE,
     );
-    expect(deps.applyTreeSnapshot).toHaveBeenCalledTimes(1);
   });
 });

apps/code/src/main/services/handoff/handoff-saga.ts

@@ -13,18 +13,11 @@ export type HandoffSagaInput = HandoffExecuteInput;
 
 export interface HandoffSagaOutput {
   sessionId: string;
-  snapshotApplied: boolean;
+  checkpointApplied: boolean;
   conversationTurns: number;
 }
 
 export interface HandoffSagaDeps extends HandoffBaseDeps {
-  applyTreeSnapshot(
-    snapshot: AgentTypes.TreeSnapshotEvent,
-    repoPath: string,
-    taskId: string,
-    runId: string,
-    apiClient: PostHogAPIClient,
-  ): Promise<void>;
   applyGitCheckpoint(
     checkpoint: AgentTypes.GitCheckpointEvent,
     repoPath: string,
@@ -102,7 +95,7 @@ export class HandoffSaga extends Saga<HandoffSagaInput, HandoffSagaOutput> {
       },
     );
 
-    let filesRestored = false;
+    let checkpointApplied = false;
     const checkpoint = resumeState.latestGitCheckpoint;
     if (checkpoint) {
       this.deps.onProgress(
@@ -121,29 +114,7 @@ export class HandoffSaga extends Saga<HandoffSagaInput, HandoffSagaOutput> {
             apiClient,
             input.localGitState,
           );
-        },
-        rollback: async () => {},
-      });
-    }
-
-    const snapshot = resumeState.latestSnapshot;
-    if (snapshot?.archiveUrl) {
-      this.deps.onProgress(
-        "applying_snapshot",
-        "Applying cloud file state locally...",
-      );
-
-      await this.step({
-        name: "apply_snapshot",
-        execute: async () => {
-          await this.deps.applyTreeSnapshot(
-            snapshot,
-            repoPath,
-            taskId,
-            runId,
-            apiClient,
-          );
-          filesRestored = true;
+          checkpointApplied = true;
         },
         rollback: async () => {},
       });
@@ -198,7 +169,7 @@ export class HandoffSaga extends Saga<HandoffSagaInput, HandoffSagaOutput> {
     await this.readOnlyStep("set_context", async () => {
       const context = this.buildHandoffContext(
         resumeState.conversation,
-        filesRestored,
+        checkpointApplied,
       );
       this.deps.setPendingContext(runId, context);
     });
@@ -207,20 +178,20 @@ export class HandoffSaga extends Saga<HandoffSagaInput, HandoffSagaOutput> {
 
     return {
       sessionId: agentSessionId,
-      snapshotApplied: filesRestored,
+      checkpointApplied,
       conversationTurns: resumeState.conversation.length,
     };
   }
 
   private buildHandoffContext(
     conversation: AgentResume.ConversationTurn[],
-    snapshotApplied: boolean,
+    checkpointApplied: boolean,
   ): string {
     const conversationSummary = formatConversationForResume(conversation);
 
-    const fileStatus = snapshotApplied
-      ? "The workspace files have been fully restored from the cloud session."
-      : "The workspace files from the cloud session could not be restored. You are working with the local file state.";
+    const fileStatus = checkpointApplied
+      ? "The workspace git state and files have been restored from the cloud session checkpoint."
+      : "The workspace from the cloud session could not be restored from a checkpoint. You are working with the local file state.";
 
     return (
       `You are resuming a previous conversation that was running in a cloud sandbox. ` +

apps/code/src/main/services/handoff/handoff-to-cloud-saga.test.ts

@@ -0,0 +1,88 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  HandoffToCloudSaga,
+  type HandoffToCloudSagaDeps,
+} from "./handoff-to-cloud-saga";
+
+function createDeps(
+  overrides: Partial<HandoffToCloudSagaDeps> = {},
+): HandoffToCloudSagaDeps {
+  return {
+    captureGitCheckpoint: vi.fn().mockResolvedValue({
+      checkpointId: "checkpoint-1",
+      checkpointRef: "refs/posthog-code-checkpoint/checkpoint-1",
+    }),
+    persistCheckpointToLog: vi.fn().mockResolvedValue(undefined),
+    countLocalLogEntries: vi.fn().mockReturnValue(7),
+    resumeRunInCloud: vi.fn().mockResolvedValue(undefined),
+    killSession: vi.fn().mockResolvedValue(undefined),
+    updateWorkspaceMode: vi.fn(),
+    onProgress: vi.fn(),
+    ...overrides,
+  } as unknown as HandoffToCloudSagaDeps;
+}
+
+describe("HandoffToCloudSaga", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("persists the fresh checkpoint, starts cloud, then kills the local session", async () => {
+    const deps = createDeps();
+    const order: string[] = [];
+
+    vi.mocked(deps.persistCheckpointToLog).mockImplementation(async () => {
+      order.push("checkpoint");
+    });
+    vi.mocked(deps.resumeRunInCloud).mockImplementation(async () => {
+      order.push("resume");
+    });
+    vi.mocked(deps.killSession).mockImplementation(async () => {
+      order.push("kill");
+    });
+
+    const saga = new HandoffToCloudSaga(deps);
+    const result = await saga.run({
+      taskId: "task-1",
+      runId: "run-1",
+      repoPath: "/repo/path",
+      apiHost: "https://us.posthog.com",
+      teamId: 1,
+      localGitState: {
+        head: "head-1",
+        branch: "main",
+        upstreamHead: "upstream-head-1",
+        upstreamRemote: "origin",
+        upstreamMergeRef: "refs/heads/main",
+      },
+    });
+
+    expect(result.success).toBe(true);
+    expect(order).toEqual(["checkpoint", "resume", "kill"]);
+    expect(deps.countLocalLogEntries).toHaveBeenCalledWith("run-1");
+    if (result.success) {
+      expect(result.data.logEntryCount).toBe(7);
+      expect(result.data.checkpointCaptured).toBe(true);
+    }
+  });
+
+  it("reports logEntryCount of 0 when no local cache exists", async () => {
+    const deps = createDeps({
+      countLocalLogEntries: vi.fn().mockReturnValue(0),
+    });
+
+    const saga = new HandoffToCloudSaga(deps);
+    const result = await saga.run({
+      taskId: "task-1",
+      runId: "run-1",
+      repoPath: "/repo/path",
+      apiHost: "https://us.posthog.com",
+      teamId: 1,
+    });
+
+    expect(result.success).toBe(true);
+    if (result.success) {
+      expect(result.data.logEntryCount).toBe(0);
+    }
+  });
+});