fix(handoff): set mode and repositoryId together on cloud hand off

apps/code/src/main/db/repositories/workspace-repository.mock.ts

@@ -68,6 +68,17 @@ export function createMockWorkspaceRepository(): MockWorkspaceRepository {
     updateLastViewedAt: () => {},
     updateLastActivityAt: () => {},
     updateMode: () => {},
+    setModeAndRepository: (taskId, mode, repositoryId) => {
+      const id = taskIndex.get(taskId);
+      const existing = id ? workspaces.get(id) : undefined;
+      if (!id || !existing) return;
+      workspaces.set(id, {
+        ...existing,
+        mode,
+        repositoryId,
+        updatedAt: new Date().toISOString(),
+      });
+    },
     deleteAll: () => {
       workspaces.clear();
       taskIndex.clear();

apps/code/src/main/db/repositories/workspace-repository.ts

@@ -28,6 +28,11 @@ export interface IWorkspaceRepository {
   updateLastActivityAt(taskId: string, lastActivityAt: string): void;
   updateLinkedBranch(taskId: string, linkedBranch: string | null): void;
   updateMode(taskId: string, mode: WorkspaceMode): void;
+  setModeAndRepository(
+    taskId: string,
+    mode: WorkspaceMode,
+    repositoryId: string | null,
+  ): void;
   deleteAll(): void;
 }
 
@@ -141,6 +146,18 @@ export class WorkspaceRepository implements IWorkspaceRepository {
       .run();
   }
 
+  setModeAndRepository(
+    taskId: string,
+    mode: WorkspaceMode,
+    repositoryId: string | null,
+  ): void {
+    this.db
+      .update(workspaces)
+      .set({ mode, repositoryId, updatedAt: now() })
+      .where(byTaskId(taskId))
+      .run();
+  }
+
   deleteAll(): void {
     this.db.delete(workspaces).run();
   }

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

@@ -64,6 +64,7 @@ function createDeps(overrides: Partial<HandoffSagaDeps> = {}): HandoffSagaDeps {
     }),
     applyGitCheckpoint: vi.fn().mockResolvedValue(undefined),
     updateWorkspaceMode: vi.fn(),
+    attachWorkspaceToFolder: vi.fn().mockReturnValue({ revert: vi.fn() }),
     reconnectSession: vi.fn().mockResolvedValue({
       sessionId: "session-1",
       channel: "ch-1",
@@ -247,9 +248,11 @@ describe("HandoffSaga", () => {
   });
 
   describe("rollbacks", () => {
-    it("rolls back workspace mode when spawn_agent fails", async () => {
+    it("reverts workspace attachment when spawn_agent fails", async () => {
+      const revert = vi.fn();
       const { deps, result } = await runSaga({
         deps: {
+          attachWorkspaceToFolder: vi.fn().mockReturnValue({ revert }),
           reconnectSession: vi
             .fn()
             .mockRejectedValue(new Error("spawn failed")),
@@ -259,7 +262,11 @@ describe("HandoffSaga", () => {
       expect(result.success).toBe(false);
       if (result.success) return;
       expect(result.failedStep).toBe("spawn_agent");
-      expect(deps.updateWorkspaceMode).toHaveBeenCalledWith("task-1", "cloud");
+      expect(deps.attachWorkspaceToFolder).toHaveBeenCalledWith(
+        "task-1",
+        "/repo",
+      );
+      expect(revert).toHaveBeenCalledTimes(1);
     });
 
     it("kills session on rollback if spawn partially succeeded", async () => {
@@ -274,7 +281,7 @@ describe("HandoffSaga", () => {
       expect(result.failedStep).toBe("spawn_agent");
     });
 
-    it("fails at fetch_and_rebuild without rolling back workspace", async () => {
+    it("fails at fetch_and_rebuild without touching workspace state", async () => {
       mockResumeFromLog.mockRejectedValue(new Error("API down"));
 
       const deps = createDeps();
@@ -284,7 +291,7 @@ describe("HandoffSaga", () => {
       expect(result.success).toBe(false);
       if (result.success) return;
       expect(result.failedStep).toBe("fetch_and_rebuild");
-      expect(deps.updateWorkspaceMode).not.toHaveBeenCalled();
+      expect(deps.attachWorkspaceToFolder).not.toHaveBeenCalled();
       expect(deps.reconnectSession).not.toHaveBeenCalled();
     });
   });

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

@@ -18,6 +18,10 @@ export interface HandoffSagaOutput {
 }
 
 export interface HandoffSagaDeps extends HandoffBaseDeps {
+  attachWorkspaceToFolder(
+    taskId: string,
+    repoPath: string,
+  ): { revert: () => void };
   applyGitCheckpoint(
     checkpoint: AgentTypes.GitCheckpointEvent,
     repoPath: string,
@@ -120,13 +124,11 @@ export class HandoffSaga extends Saga<HandoffSagaInput, HandoffSagaOutput> {
       });
     }
 
-    await this.step({
-      name: "update_workspace",
-      execute: async () => {
-        this.deps.updateWorkspaceMode(taskId, "local");
-      },
-      rollback: async () => {
-        this.deps.updateWorkspaceMode(taskId, "cloud");
+    await this.step<{ revert: () => void }>({
+      name: "attach_workspace_to_folder",
+      execute: async () => this.deps.attachWorkspaceToFolder(taskId, repoPath),
+      rollback: async ({ revert }) => {
+        revert();
       },
     });
 

apps/code/src/main/services/handoff/service.test.ts

@@ -86,7 +86,12 @@ function createService(): HandoffService {
   const agentAuthAdapter = {
     createPosthogConfig: mockCreatePosthogConfig,
   } as never;
-  const workspaceRepo = { updateMode: mockUpdateMode } as never;
+  const workspaceRepo = {
+    updateMode: mockUpdateMode,
+    findByTaskId: vi.fn().mockReturnValue(null),
+    setModeAndRepository: vi.fn(),
+  } as never;
+  const repositoryRepo = { findByPath: vi.fn().mockReturnValue(null) } as never;
   const dialog = { confirm: vi.fn().mockResolvedValue(1) } as never;
   const appLifecycle = {
     whenReady: vi.fn().mockResolvedValue(undefined),
@@ -98,6 +103,7 @@ function createService(): HandoffService {
     cloudTaskService,
     agentAuthAdapter,
     workspaceRepo,
+    repositoryRepo,
     dialog,
     appLifecycle,
   );

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

@@ -23,6 +23,7 @@ import { StashPushSaga } from "@posthog/git/sagas/stash";
 import type { IAppLifecycle } from "@posthog/platform/app-lifecycle";
 import type { IDialog } from "@posthog/platform/dialog";
 import { inject, injectable } from "inversify";
+import type { IRepositoryRepository } from "../../db/repositories/repository-repository";
 import type { IWorkspaceRepository } from "../../db/repositories/workspace-repository";
 import type { AgentAuthAdapter } from "../agent/auth-adapter";
 import type { AgentService } from "../agent/service";
@@ -61,6 +62,8 @@ export class HandoffService extends TypedEventEmitter<HandoffServiceEvents> {
     private readonly agentAuthAdapter: AgentAuthAdapter,
     @inject(MAIN_TOKENS.WorkspaceRepository)
     private readonly workspaceRepo: IWorkspaceRepository,
+    @inject(MAIN_TOKENS.RepositoryRepository)
+    private readonly repositoryRepo: IRepositoryRepository,
     @inject(MAIN_TOKENS.Dialog)
     private readonly dialog: IDialog,
     @inject(MAIN_TOKENS.AppLifecycle)
@@ -151,6 +154,35 @@ export class HandoffService extends TypedEventEmitter<HandoffServiceEvents> {
         this.workspaceRepo.updateMode(taskId, mode);
       },
 
+      attachWorkspaceToFolder: (taskId, repoPath) => {
+        const repository = this.repositoryRepo.findByPath(repoPath);
+        if (!repository) {
+          throw new Error(
+            `No registered folder for path '${repoPath}' — cannot attach workspace`,
+          );
+        }
+        const previous = this.workspaceRepo.findByTaskId(taskId);
+        if (!previous) {
+          throw new Error(`No workspace exists for task ${taskId}`);
+        }
+        if (
+          previous.mode === "local" &&
+          previous.repositoryId === repository.id
+        ) {
+          return { revert: () => {} };
+        }
+        this.workspaceRepo.setModeAndRepository(taskId, "local", repository.id);
+        return {
+          revert: () => {
+            this.workspaceRepo.setModeAndRepository(
+              taskId,
+              previous.mode,
+              previous.repositoryId,
+            );
+          },
+        };
+      },
+
       seedLocalLogs: async (runId: string, logUrl: string) => {
         const response = await fetch(logUrl);
         if (!response.ok) {
@@ -165,7 +197,11 @@ export class HandoffService extends TypedEventEmitter<HandoffServiceEvents> {
         const logDir = join(homedir(), ".posthog-code", "sessions", runId);
         mkdirSync(logDir, { recursive: true });
         const marker = JSON.stringify({ type: "seed_boundary" });
-        writeFileSync(join(logDir, "logs.ndjson"), `${content}\n${marker}`);
+        const trailingNewline = content.endsWith("\n") ? "" : "\n";
+        writeFileSync(
+          join(logDir, "logs.ndjson"),
+          `${content}${trailingNewline}${marker}\n`,
+        );
         log.info("Seeded local logs from cloud", {
           runId,
           bytes: content.length,

apps/code/src/renderer/features/sessions/service/localHandoffService.ts

@@ -17,29 +17,16 @@ async function resolveRepoPathFromRemote(
   return repo?.path ?? null;
 }
 
-async function resolveRepoPathFromPicker(
-  taskId: string,
-): Promise<string | null> {
+async function resolveRepoPathFromPicker(): Promise<string | null> {
   const selectedPath = await trpcClient.os.selectDirectory.query();
   if (!selectedPath) return null;
 
-  let folder = (await trpcClient.folders.getFolders.query()).find(
-    (f) => f.path === selectedPath,
-  );
+  const folders = await trpcClient.folders.getFolders.query();
+  const folder = folders.find((f) => f.path === selectedPath);
   if (!folder) {
-    folder = await trpcClient.folders.addFolder.mutate({
-      folderPath: selectedPath,
-    });
+    await trpcClient.folders.addFolder.mutate({ folderPath: selectedPath });
   }
 
-  await trpcClient.workspace.create.mutate({
-    taskId,
-    mainRepoPath: selectedPath,
-    folderId: folder.id,
-    folderPath: selectedPath,
-    mode: "local",
-  });
-
   return selectedPath;
 }
 
@@ -79,7 +66,7 @@ export class LocalHandoffService {
     try {
       const targetPath =
         (await resolveRepoPathFromRemote(task.repository)) ??
-        (await resolveRepoPathFromPicker(taskId));
+        (await resolveRepoPathFromPicker());
 
       if (!targetPath) return;