Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions apps/code/src/main/services/cloud-task/service.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
27 changes: 27 additions & 0 deletions apps/code/src/main/services/cloud-task/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
55 changes: 12 additions & 43 deletions apps/code/src/main/services/handoff/handoff-saga.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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({
Expand All @@ -96,7 +82,6 @@ function createResumeState(
): AgentResume.ResumeState {
return {
conversation: [],
latestSnapshot: null,
latestGitCheckpoint: null,
interrupted: false,
logEntryCount: 0,
Expand Down Expand Up @@ -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);
});

Expand All @@ -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 () => {
Expand Down Expand Up @@ -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"),
);
});

Expand All @@ -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",
]);
Expand Down Expand Up @@ -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(),
},
});
Expand All @@ -338,6 +308,5 @@ describe("HandoffSaga", () => {
expect.any(Object),
DEFAULT_LOCAL_GIT_STATE,
);
expect(deps.applyTreeSnapshot).toHaveBeenCalledTimes(1);
});
});
47 changes: 9 additions & 38 deletions apps/code/src/main/services/handoff/handoff-saga.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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 () => {},
});
Expand Down Expand Up @@ -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);
});
Expand All @@ -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. ` +
Expand Down
88 changes: 88 additions & 0 deletions apps/code/src/main/services/handoff/handoff-to-cloud-saga.test.ts
Original file line number Diff line number Diff line change
@@ -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);
}
});
});
Loading
Loading