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..05d38d3ae 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,11 @@ export class HandoffSaga extends Saga { }); } - 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(); }, }); diff --git a/apps/code/src/main/services/handoff/service.test.ts b/apps/code/src/main/services/handoff/service.test.ts index 56e6a00ed..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,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, ); 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;