From b306ae2a2664a583892ce9fe0af0396799419f01 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Apr 2026 18:37:15 +0200 Subject: [PATCH 1/2] fix(handoff): set mode and repositoryId together --- .../repositories/workspace-repository.mock.ts | 11 ++++++ .../db/repositories/workspace-repository.ts | 17 +++++++++ .../services/handoff/handoff-saga.test.ts | 15 ++++++-- .../src/main/services/handoff/handoff-saga.ts | 12 ++++-- .../src/main/services/handoff/service.test.ts | 2 + .../code/src/main/services/handoff/service.ts | 38 ++++++++++++++++++- .../sessions/service/localHandoffService.ts | 23 +++-------- 7 files changed, 92 insertions(+), 26 deletions(-) diff --git a/apps/code/src/main/db/repositories/workspace-repository.mock.ts b/apps/code/src/main/db/repositories/workspace-repository.mock.ts index 804af19ba..7be3ade37 100644 --- a/apps/code/src/main/db/repositories/workspace-repository.mock.ts +++ b/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(); diff --git a/apps/code/src/main/db/repositories/workspace-repository.ts b/apps/code/src/main/db/repositories/workspace-repository.ts index d22079bee..6dfcb391f 100644 --- a/apps/code/src/main/db/repositories/workspace-repository.ts +++ b/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(); } diff --git a/apps/code/src/main/services/handoff/handoff-saga.test.ts b/apps/code/src/main/services/handoff/handoff-saga.test.ts index 8a8067859..eb6760457 100644 --- a/apps/code/src/main/services/handoff/handoff-saga.test.ts +++ b/apps/code/src/main/services/handoff/handoff-saga.test.ts @@ -64,6 +64,7 @@ function createDeps(overrides: Partial = {}): 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(); }); }); diff --git a/apps/code/src/main/services/handoff/handoff-saga.ts b/apps/code/src/main/services/handoff/handoff-saga.ts index 9b00c3b39..bc409c53d 100644 --- a/apps/code/src/main/services/handoff/handoff-saga.ts +++ b/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,15 @@ export class HandoffSaga extends Saga { }); } + let attachmentRevert: (() => void) | null = null; await this.step({ - name: "update_workspace", + name: "attach_workspace_to_folder", execute: async () => { - this.deps.updateWorkspaceMode(taskId, "local"); + const { revert } = this.deps.attachWorkspaceToFolder(taskId, repoPath); + attachmentRevert = revert; }, rollback: async () => { - this.deps.updateWorkspaceMode(taskId, "cloud"); + attachmentRevert?.(); }, }); diff --git a/apps/code/src/main/services/handoff/service.test.ts b/apps/code/src/main/services/handoff/service.test.ts index 56e6a00ed..42b8a2195 100644 --- a/apps/code/src/main/services/handoff/service.test.ts +++ b/apps/code/src/main/services/handoff/service.test.ts @@ -87,6 +87,7 @@ function createService(): HandoffService { createPosthogConfig: mockCreatePosthogConfig, } as never; const workspaceRepo = { updateMode: mockUpdateMode } 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 +99,7 @@ function createService(): HandoffService { cloudTaskService, agentAuthAdapter, workspaceRepo, + repositoryRepo, dialog, appLifecycle, ); diff --git a/apps/code/src/main/services/handoff/service.ts b/apps/code/src/main/services/handoff/service.ts index 08e5fccfc..f7f3d3b18 100644 --- a/apps/code/src/main/services/handoff/service.ts +++ b/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 { 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 { 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 { 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, diff --git a/apps/code/src/renderer/features/sessions/service/localHandoffService.ts b/apps/code/src/renderer/features/sessions/service/localHandoffService.ts index 581c3b32e..d74072370 100644 --- a/apps/code/src/renderer/features/sessions/service/localHandoffService.ts +++ b/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 { +async function resolveRepoPathFromPicker(): Promise { 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; From aa675669f9a0e75c45fe88f8cfeb1cc2e1f5bc60 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Apr 2026 18:47:54 +0200 Subject: [PATCH 2/2] fix(handoff): refactor attachment step to use saga return values Generated-By: PostHog Code Task-Id: e2a44dd9-f5de-433e-8d8c-dce9920f20ba --- apps/code/src/main/services/handoff/handoff-saga.ts | 12 ++++-------- apps/code/src/main/services/handoff/service.test.ts | 6 +++++- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/apps/code/src/main/services/handoff/handoff-saga.ts b/apps/code/src/main/services/handoff/handoff-saga.ts index bc409c53d..05d38d3ae 100644 --- a/apps/code/src/main/services/handoff/handoff-saga.ts +++ b/apps/code/src/main/services/handoff/handoff-saga.ts @@ -124,15 +124,11 @@ export class HandoffSaga extends Saga { }); } - let attachmentRevert: (() => void) | null = null; - await this.step({ + await this.step<{ revert: () => void }>({ name: "attach_workspace_to_folder", - execute: async () => { - const { revert } = this.deps.attachWorkspaceToFolder(taskId, repoPath); - attachmentRevert = revert; - }, - rollback: async () => { - attachmentRevert?.(); + execute: async () => this.deps.attachWorkspaceToFolder(taskId, repoPath), + rollback: async ({ revert }) => { + revert(); }, }); diff --git a/apps/code/src/main/services/handoff/service.test.ts b/apps/code/src/main/services/handoff/service.test.ts index 42b8a2195..94bd62800 100644 --- a/apps/code/src/main/services/handoff/service.test.ts +++ b/apps/code/src/main/services/handoff/service.test.ts @@ -86,7 +86,11 @@ 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 = {