diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index e9f84c43ad..135aaad90f 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -799,6 +799,22 @@ container.bind(LOGS_SERVICE).toDynamicValue((ctx) => { ws.localLogs.readTail.query({ taskRunId, maxBytes }), writeLocalLogs: (taskRunId: string, content: string) => ws.localLogs.write.mutate({ taskRunId, content }), + seedLocalLogs: (taskRunId: string, content: string) => + ws.localLogs.seed.mutate({ taskRunId, content }), + cloneLocalLogs: async ( + sourceTaskRunId: string, + targetTaskRunId: string, + ) => { + const content = await ws.localLogs.read.query({ + taskRunId: sourceTaskRunId, + }); + if (content) { + await ws.localLogs.seed.mutate({ + taskRunId: targetTaskRunId, + content, + }); + } + }, }; }); container.bind(MAIN_ENCRYPTION_SERVICE).to(EncryptionService); diff --git a/apps/code/src/renderer/di/bindings.ts b/apps/code/src/renderer/di/bindings.ts index f8556c970f..e4513c39d2 100644 --- a/apps/code/src/renderer/di/bindings.ts +++ b/apps/code/src/renderer/di/bindings.ts @@ -120,10 +120,12 @@ import { import { TASK_CREATION_EFFECTS, TASK_CREATION_HOST, + TASK_FORK_SERVICE, WORKSPACE_SETUP_SAGA, } from "@posthog/core/task-detail/identifiers"; import type { TaskCreationEffects } from "@posthog/core/task-detail/taskCreationEffects"; import type { ITaskCreationHost } from "@posthog/core/task-detail/taskCreationHost"; +import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; import { TASK_SERVICE, type TaskService, @@ -299,6 +301,7 @@ export interface RendererBindings { [TASK_CREATION_EFFECTS]: TaskCreationEffects; [RENDERER_TASK_SERVICE]: TaskService; [TASK_SERVICE]: TaskService; + [TASK_FORK_SERVICE]: TaskForkService; [WORKSPACE_SETUP_SAGA]: WorkspaceSetupSaga; [SESSION_SERVICE]: SessionService; [LOCAL_HANDOFF_HOST]: LocalHandoffHost; diff --git a/apps/code/src/renderer/di/container.ts b/apps/code/src/renderer/di/container.ts index 5d7d27ddf4..357e4c7732 100644 --- a/apps/code/src/renderer/di/container.ts +++ b/apps/code/src/renderer/di/container.ts @@ -63,9 +63,11 @@ import type { SkillsWorkspaceClient } from "@posthog/core/skills/teamSkillsServi import { TASK_CREATION_EFFECTS, TASK_CREATION_HOST, + TASK_FORK_SERVICE, WORKSPACE_SETUP_SAGA, } from "@posthog/core/task-detail/identifiers"; import type { ITaskCreationHost } from "@posthog/core/task-detail/taskCreationHost"; +import { TaskForkService } from "@posthog/core/task-detail/taskForkService"; import { TASK_SERVICE, TaskService, @@ -302,6 +304,10 @@ container.load(piRuntimeModule); container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects); container.bind(RENDERER_TASK_SERVICE).to(TaskService); container.bind(TASK_SERVICE).toService(RENDERER_TASK_SERVICE); +container + .bind(TASK_FORK_SERVICE) + .to(TaskForkService) + .inSingletonScope(); container .bind(WORKSPACE_SETUP_SAGA) .to(WorkspaceSetupSaga) diff --git a/packages/agent/src/adapters/claude/claude-agent.resume-model.test.ts b/packages/agent/src/adapters/claude/claude-agent.resume-model.test.ts index 3ce2eebe0f..c077a482c7 100644 --- a/packages/agent/src/adapters/claude/claude-agent.resume-model.test.ts +++ b/packages/agent/src/adapters/claude/claude-agent.resume-model.test.ts @@ -38,6 +38,9 @@ function makeQueryHandle(): SdkQueryHandle { } const createdQueries: SdkQueryHandle[] = []; +const forkSession = vi.fn().mockResolvedValue({ + sessionId: "0197a000-0000-7000-8000-0000000000f0", +}); vi.mock("@anthropic-ai/claude-agent-sdk", () => ({ query: vi.fn(() => { @@ -46,6 +49,7 @@ vi.mock("@anthropic-ai/claude-agent-sdk", () => ({ return handle; }), getSessionMessages: vi.fn().mockResolvedValue([]), + forkSession, listSessions: vi.fn().mockResolvedValue([]), createSdkMcpServer: vi.fn(() => ({ type: "sdk", @@ -235,4 +239,19 @@ describe("ClaudeAcpAgent session model on resume", () => { expect(createdQueries[0]?.close).toHaveBeenCalledTimes(1); }); + + it("copies the complete transcript before resuming the fork", async () => { + const agent = makeAgent(); + const sourceSessionId = "0197a000-0000-7000-8000-0000000000aa"; + + const response = await agent.unstable_forkSession({ + sessionId: sourceSessionId, + cwd, + mcpServers: [], + _meta: { taskRunId: "run-fork" }, + }); + + expect(forkSession).toHaveBeenCalledWith(sourceSessionId); + expect(response.sessionId).toBe("0197a000-0000-7000-8000-0000000000f0"); + }); }); diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 983567d8a5..f7280273b9 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -33,6 +33,7 @@ import { import { type CanUseTool, type FastModeState, + forkSession as forkClaudeSession, getSessionInfo, getSessionMessages, listSessions, @@ -360,6 +361,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { async unstable_forkSession( params: ForkSessionRequest, ): Promise { + const forked = await forkClaudeSession(params.sessionId); return this.createSession( { cwd: params.cwd, @@ -367,7 +369,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { additionalDirectories: params.additionalDirectories, _meta: params._meta, }, - { resume: params.sessionId, forkSession: true }, + { resume: forked.sessionId }, ); } diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 472a8eeb74..8fde892a89 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -223,6 +223,7 @@ describe("PostHogAPIClient", () => { model: "gpt-5.4", reasoningLevel: "high", initialPermissionMode: "auto", + resumeFromRunId: "run-source", }), ).resolves.toEqual({ id: "run-123", environment: "cloud" }); @@ -238,6 +239,7 @@ describe("PostHogAPIClient", () => { model: "gpt-5.4", reasoning_effort: "high", initial_permission_mode: "auto", + resume_from_run_id: "run-source", environment: "cloud", }), }, diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 29c687a4fb..a1a32a067b 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -616,6 +616,7 @@ interface CreateTaskRunOptions extends CloudRunOptions { environment?: "local" | "cloud"; mode?: "interactive" | "background"; branch?: string | null; + resumeFromRunId?: string; } interface StartTaskRunOptions { diff --git a/packages/core/src/context-menu/context-menu.test.ts b/packages/core/src/context-menu/context-menu.test.ts index b06a4a79cb..7f7f9df0b4 100644 --- a/packages/core/src/context-menu/context-menu.test.ts +++ b/packages/core/src/context-menu/context-menu.test.ts @@ -119,6 +119,30 @@ describe("ContextMenuService.showTaskContextMenu", () => { expect(labels(idle.lastItems)).not.toContain("Stop task"); }); + it("offers fork and source navigation when available", async () => { + const forkMenu = new FakeContextMenu(); + const forkResult = makeService(forkMenu).showTaskContextMenu({ + ...baseTask, + canFork: true, + hasSourceTask: true, + }); + await forkMenu.shown; + + expect(labels(forkMenu.lastItems)).toContain("Fork task"); + expect(labels(forkMenu.lastItems)).toContain("Open source task"); + findItem(forkMenu.lastItems, "Fork task").click(); + expect(await forkResult).toEqual({ action: { type: "fork" } }); + + const sourceMenu = new FakeContextMenu(); + const sourceResult = makeService(sourceMenu).showTaskContextMenu({ + ...baseTask, + hasSourceTask: true, + }); + await sourceMenu.shown; + findItem(sourceMenu.lastItems, "Open source task").click(); + expect(await sourceResult).toEqual({ action: { type: "open-source" } }); + }); + it("hides Add to Command Center when already in it", async () => { const inCc = new FakeContextMenu(); makeService(inCc).showTaskContextMenu({ diff --git a/packages/core/src/context-menu/context-menu.ts b/packages/core/src/context-menu/context-menu.ts index d4437e9488..c6dec44c53 100644 --- a/packages/core/src/context-menu/context-menu.ts +++ b/packages/core/src/context-menu/context-menu.ts @@ -115,6 +115,8 @@ export class ContextMenuService { isPinned, isSuspended, canStop, + canFork, + hasSourceTask, isInCommandCenter, hasEmptyCommandCenterCell, channels, @@ -142,6 +144,21 @@ export class ContextMenuService { return this.showMenu([ this.item(isPinned ? "Unpin" : "Pin", { type: "pin" }), this.item("Rename", { type: "rename" }), + ...(canFork || hasSourceTask + ? [ + this.separator(), + ...(canFork + ? [this.item("Fork task", { type: "fork" as const })] + : []), + ...(hasSourceTask + ? [ + this.item("Open source task", { + type: "open-source" as const, + }), + ] + : []), + ] + : []), ...(canStop ? [this.separator(), this.item("Stop task", { type: "stop" as const })] : []), diff --git a/packages/core/src/context-menu/schemas.ts b/packages/core/src/context-menu/schemas.ts index b6b664ede4..91627a341e 100644 --- a/packages/core/src/context-menu/schemas.ts +++ b/packages/core/src/context-menu/schemas.ts @@ -7,6 +7,8 @@ export const taskContextMenuInput = z.object({ isPinned: z.boolean().optional(), isSuspended: z.boolean().optional(), canStop: z.boolean().optional(), + canFork: z.boolean().optional(), + hasSourceTask: z.boolean().optional(), isInCommandCenter: z.boolean().optional(), hasEmptyCommandCenterCell: z.boolean().optional(), // Top-level desktop_file_system channels available as "File to…" targets. @@ -47,6 +49,8 @@ const taskAction = z.discriminatedUnion("type", [ z.object({ type: z.literal("pin") }), z.object({ type: z.literal("suspend") }), z.object({ type: z.literal("stop") }), + z.object({ type: z.literal("fork") }), + z.object({ type: z.literal("open-source") }), z.object({ type: z.literal("archive") }), z.object({ type: z.literal("archive-prior") }), z.object({ type: z.literal("delete") }), diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 3846d2164f..98f6915453 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -154,6 +154,8 @@ type TrpcSubscription = { export interface SessionTrpc { agent: { start: TrpcMutation; + fork: TrpcMutation; + flushLogs: TrpcMutation; reconnect: TrpcMutation; cancel: TrpcMutation; prompt: TrpcMutation; @@ -197,6 +199,8 @@ export interface SessionTrpc { readLocalLogsTail?: TrpcQuery; fetchS3Logs: TrpcQuery; writeLocalLogs: TrpcMutation; + seedLocalLogs: TrpcMutation; + cloneLocalLogs: TrpcMutation; }; os: { openExternal: TrpcMutation }; } @@ -1503,6 +1507,115 @@ export class SessionService { } } + async forkLocalTask(params: { + sourceTaskId: string; + task: Task; + repoPath: string; + }): Promise { + const source = this.d.store.getSessionByTaskId(params.sourceTaskId); + if (!source || source.isCloud) { + throw new Error("The source task session is not available"); + } + if (source.isPromptPending) { + throw new Error("Wait for the source task to finish before forking it"); + } + + const authStatus = await this.getAuthCredentialsStatus(); + if (authStatus.kind !== "ready") { + throw new Error("Unable to reach server. Please check your connection."); + } + + const taskRun = await authStatus.auth.client.createTaskRun(params.task.id); + if (!taskRun?.id) { + throw new Error("Failed to create task run. Please try again."); + } + params.task.latest_run = taskRun; + const sourceTaskRunId = source.taskRunId; + + await this.d.trpc.agent.flushLogs.mutate({ + taskRunId: sourceTaskRunId, + }); + await this.d.trpc.logs.cloneLocalLogs.mutate({ + sourceTaskRunId, + targetTaskRunId: taskRun.id, + }); + + params.task.latest_run = await authStatus.auth.client.updateTaskRun( + params.task.id, + taskRun.id, + { + state: { + ...taskRun.state, + forked_from_task_id: params.sourceTaskId, + forked_from_run_id: sourceTaskRunId, + }, + }, + ); + + const { customInstructions, rtkEnabledLocal, spokenNarrationEnabled } = + this.d.settings; + const result = await this.d.trpc.agent.fork.mutate({ + sourceTaskRunId, + taskId: params.task.id, + taskRunId: taskRun.id, + repoPath: params.repoPath, + apiHost: authStatus.auth.apiHost, + projectId: authStatus.auth.projectId, + permissionMode: source.executionMode, + adapter: source.adapter, + customInstructions: customInstructions || undefined, + rtkEnabled: rtkEnabledLocal, + spokenNarration: spokenNarrationEnabled === true, + effort: effortLevelSchema.safeParse(source.reasoningLevel).success + ? (source.reasoningLevel as EffortLevel) + : undefined, + model: source.model ?? this.d.DEFAULT_GATEWAY_MODEL, + }); + const { rawEntries } = await this.fetchSessionLogs(undefined, taskRun.id); + + const session = createBaseSession( + taskRun.id, + params.task.id, + params.task.title, + ); + session.channel = result.channel; + session.status = "connected"; + session.events = + rawEntries.length > 0 + ? convertStoredEntriesToEvents(rawEntries) + : [...source.events]; + session.adapter = source.adapter; + session.model = source.model; + session.executionMode = source.executionMode; + session.reasoningLevel = source.reasoningLevel; + session.configOptions = result.configOptions as + | SessionConfigOption[] + | undefined; + session.steering = readSteering(result); + + if (session.configOptions) { + this.d.setPersistedConfigOptions(taskRun.id, session.configOptions); + } + if (source.adapter) { + this.d.adapterStore.setAdapter(taskRun.id, source.adapter); + } + + this.localRepoPaths.set(params.task.id, params.repoPath); + this.d.store.setSession(session); + this.subscribeToChannel(taskRun.id); + + this.d.track(ANALYTICS_EVENTS.TASK_RUN_STARTED, { + task_id: params.task.id, + execution_type: "local", + initial_mode: source.executionMode, + adapter: source.adapter, + }); + } + + getSessionByTaskId(taskId: string): AgentSession | undefined { + return this.d.store.getSessionByTaskId(taskId); + } + async loadLogsOnly(params: { taskId: string; taskRunId: string; diff --git a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts new file mode 100644 index 0000000000..d7fd726fc5 --- /dev/null +++ b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts @@ -0,0 +1,254 @@ +import type { AgentSession } from "@posthog/shared"; +import type { Task, TaskRun } from "@posthog/shared/domain-types"; +import { describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; + +const sourceSession = { + taskRunId: "source-run", + taskId: "source-task", + taskTitle: "Source task", + channel: "source-channel", + events: [], + startedAt: 1, + status: "connected", + isPromptPending: false, + isCompacting: false, + promptStartedAt: null, + pendingPermissions: new Map(), + pausedDurationMs: 0, + messageQueue: [], + optimisticItems: [], + executionMode: "acceptEdits", + adapter: "claude", + model: "claude-sonnet-4-5", + reasoningLevel: "medium", +} as AgentSession; + +const childRun = { + id: "child-run", + task: "child-task", + team: 1, + branch: "main", + environment: "local", + status: "in_progress", + log_url: "https://example.com/logs/child-run", + error_message: null, + output: null, + state: { existing: true }, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + completed_at: null, +} as TaskRun; + +const childTask = { + id: "child-task", + task_number: 2, + slug: "child-task", + title: "Child task", + description: "Fork the task", + origin_product: "user_created", + repository: "PostHog/code", + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + latest_run: undefined, +} as Task; + +function createHarness(updateTaskRun: ReturnType) { + const setSession = vi.fn(); + const setPersistedConfigOptions = vi.fn(); + const setAdapter = vi.fn(); + const subscribe = vi.fn(() => ({ unsubscribe: vi.fn() })); + const track = vi.fn(); + const flushLogs = vi.fn().mockResolvedValue(undefined); + const cloneLocalLogs = vi.fn().mockResolvedValue(undefined); + const readLocalLogs = vi.fn().mockResolvedValue(""); + const fork = vi.fn().mockResolvedValue({ + channel: "child-channel", + configOptions: [{ id: "model", type: "select" }], + }); + const deps = { + store: { + getSessionByTaskId: vi.fn(() => sourceSession), + setSession, + }, + trpc: { + agent: { + flushLogs: { mutate: flushLogs }, + fork: { mutate: fork }, + onSessionEvent: { subscribe }, + onPermissionRequest: { subscribe }, + onSessionIdleKilled: { subscribe }, + }, + logs: { + cloneLocalLogs: { mutate: cloneLocalLogs }, + readLocalLogs: { query: readLocalLogs }, + }, + }, + fetchAuthState: vi.fn().mockResolvedValue({ + status: "authenticated", + bootstrapComplete: true, + cloudRegion: "us", + currentProjectId: 1, + }), + createAuthenticatedClient: vi.fn(() => ({ + createTaskRun: vi.fn().mockResolvedValue(childRun), + updateTaskRun, + })), + settings: {}, + setPersistedConfigOptions, + adapterStore: { setAdapter }, + DEFAULT_GATEWAY_MODEL: "claude-sonnet-4-5", + track, + log: { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }, + } as unknown as SessionServiceDeps; + + return { + service: new SessionService(deps), + setSession, + setPersistedConfigOptions, + setAdapter, + subscribe, + track, + flushLogs, + cloneLocalLogs, + readLocalLogs, + fork, + }; +} + +describe("SessionService.forkLocalTask", () => { + it("persists lineage before publishing the child session", async () => { + const updateTaskRun = vi.fn().mockResolvedValue({ + ...childRun, + state: { + ...childRun.state, + forked_from_task_id: "source-task", + forked_from_run_id: "source-run", + }, + }); + const { + service, + setSession, + setPersistedConfigOptions, + setAdapter, + flushLogs, + cloneLocalLogs, + fork, + } = createHarness(updateTaskRun); + + await service.forkLocalTask({ + sourceTaskId: "source-task", + task: { ...childTask }, + repoPath: "/repos/code", + }); + + expect(updateTaskRun).toHaveBeenCalledWith("child-task", "child-run", { + state: { + existing: true, + forked_from_task_id: "source-task", + forked_from_run_id: "source-run", + }, + }); + expect(flushLogs).toHaveBeenCalledWith({ taskRunId: "source-run" }); + expect(cloneLocalLogs).toHaveBeenCalledWith({ + sourceTaskRunId: "source-run", + targetTaskRunId: "child-run", + }); + expect(flushLogs.mock.invocationCallOrder[0]).toBeLessThan( + cloneLocalLogs.mock.invocationCallOrder[0], + ); + expect(cloneLocalLogs.mock.invocationCallOrder[0]).toBeLessThan( + updateTaskRun.mock.invocationCallOrder[0], + ); + expect(updateTaskRun.mock.invocationCallOrder[0]).toBeLessThan( + fork.mock.invocationCallOrder[0], + ); + expect(updateTaskRun.mock.invocationCallOrder[0]).toBeLessThan( + setPersistedConfigOptions.mock.invocationCallOrder[0], + ); + expect(updateTaskRun.mock.invocationCallOrder[0]).toBeLessThan( + setAdapter.mock.invocationCallOrder[0], + ); + expect(updateTaskRun.mock.invocationCallOrder[0]).toBeLessThan( + setSession.mock.invocationCallOrder[0], + ); + }); + + it("does not publish or persist a child session when lineage fails", async () => { + const updateTaskRun = vi + .fn() + .mockRejectedValue(new Error("lineage failed")); + const { + service, + setSession, + setPersistedConfigOptions, + setAdapter, + subscribe, + track, + fork, + } = createHarness(updateTaskRun); + + await expect( + service.forkLocalTask({ + sourceTaskId: "source-task", + task: { ...childTask }, + repoPath: "/repos/code", + }), + ).rejects.toThrow("lineage failed"); + + expect(setPersistedConfigOptions).not.toHaveBeenCalled(); + expect(setAdapter).not.toHaveBeenCalled(); + expect(setSession).not.toHaveBeenCalled(); + expect(fork).not.toHaveBeenCalled(); + expect(subscribe).toHaveBeenCalledTimes(1); + expect(track).not.toHaveBeenCalled(); + }); + + it("hydrates the child transcript from the flushed cloned log", async () => { + const updateTaskRun = vi.fn().mockResolvedValue(childRun); + const { service, readLocalLogs, setSession } = createHarness(updateTaskRun); + readLocalLogs.mockResolvedValue( + JSON.stringify({ + type: "notification", + notification: { + jsonrpc: "2.0", + method: "session/update", + params: { + update: { + sessionUpdate: "assistant_message_chunk", + content: { type: "text", text: "The completed answer" }, + }, + }, + }, + }), + ); + + await service.forkLocalTask({ + sourceTaskId: "source-task", + task: { ...childTask }, + repoPath: "/repos/code", + }); + + expect(setSession).toHaveBeenCalledWith( + expect.objectContaining({ + events: [ + expect.objectContaining({ + message: expect.objectContaining({ + method: "session/update", + params: expect.objectContaining({ + update: expect.objectContaining({ + content: { type: "text", text: "The completed answer" }, + }), + }), + }), + }), + ], + }), + ); + }); +}); diff --git a/packages/core/src/task-detail/identifiers.ts b/packages/core/src/task-detail/identifiers.ts index 4369b3d0fb..98772630e6 100644 --- a/packages/core/src/task-detail/identifiers.ts +++ b/packages/core/src/task-detail/identifiers.ts @@ -1,4 +1,7 @@ export const TASK_SERVICE = Symbol.for("posthog.core.taskDetail.taskService"); +export const TASK_FORK_SERVICE = Symbol.for( + "posthog.core.taskDetail.taskForkService", +); export const TASK_CREATION_HOST = Symbol.for( "posthog.core.taskDetail.taskCreationHost", ); diff --git a/packages/core/src/task-detail/task-detail.module.ts b/packages/core/src/task-detail/task-detail.module.ts index 193d444269..7c4b25d840 100644 --- a/packages/core/src/task-detail/task-detail.module.ts +++ b/packages/core/src/task-detail/task-detail.module.ts @@ -1,9 +1,15 @@ import { ContainerModule } from "inversify"; -import { TASK_SERVICE, WORKSPACE_SETUP_SAGA } from "./identifiers"; +import { + TASK_FORK_SERVICE, + TASK_SERVICE, + WORKSPACE_SETUP_SAGA, +} from "./identifiers"; +import { TaskForkService } from "./taskForkService"; import { TaskService } from "./taskService"; import { WorkspaceSetupSaga } from "./workspaceSetupSaga"; export const taskDetailModule = new ContainerModule(({ bind }) => { bind(TASK_SERVICE).to(TaskService).inSingletonScope(); + bind(TASK_FORK_SERVICE).to(TaskForkService).inSingletonScope(); bind(WORKSPACE_SETUP_SAGA).to(WorkspaceSetupSaga).inSingletonScope(); }); diff --git a/packages/core/src/task-detail/taskCreationApiClient.ts b/packages/core/src/task-detail/taskCreationApiClient.ts index b09c5f75d2..816b7a5855 100644 --- a/packages/core/src/task-detail/taskCreationApiClient.ts +++ b/packages/core/src/task-detail/taskCreationApiClient.ts @@ -25,6 +25,7 @@ export interface CreateTaskRunClientOptions { homeQuickAction?: string; importedMcpServers?: CloudMcpServerImport[]; relayedMcpServers?: CloudMcpServerRelayDesignation[]; + resumeFromRunId?: string; } export interface StartTaskRunClientOptions { @@ -46,4 +47,9 @@ export interface TaskCreationApiClient { runId: string, options?: StartTaskRunClientOptions, ): Promise; + updateTaskRun( + taskId: string, + runId: string, + updates: Partial>, + ): Promise; } diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index 92aefe9d1c..e3a28fd039 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/packages/core/src/task-detail/taskCreationHost.ts @@ -19,6 +19,7 @@ export interface CreateWorkspaceArgs { branch?: string; allowRemoteBranchCheckout?: boolean; reuseExistingWorktree?: boolean; + forkFromTaskId?: string; } export type CreatedWorkspaceInfo = WorkspaceInfo; @@ -82,6 +83,7 @@ export interface ITaskCreationHost { getFolders(): Promise; addFolder(args: { folderPath: string }): Promise; addAdditionalDirectory(args: { taskId: string; path: string }): Promise; + getAdditionalDirectories(taskId: string): Promise; removeAdditionalDirectory(args: { taskId: string; path: string; diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 5558ab6580..c9da88e531 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -45,6 +45,7 @@ const host = mockHost as unknown as ITaskCreationHost; const sessionService = { connectToTask: vi.fn(), + forkLocalTask: vi.fn(), disconnectFromTask: vi.fn(), rememberInitialCloudPrompt: vi.fn(), markTaskCreationInFlight: vi.fn(), @@ -93,6 +94,7 @@ function makeSaga( startTaskRun: vi.fn(), sendRunCommand: vi.fn(), updateTask: vi.fn(), + updateTaskRun: vi.fn(), ...posthog, } as never, host, @@ -197,6 +199,107 @@ describe("TaskCreationSaga", () => { ); }); + it("forks a cloud run from the backend snapshot without sending a prompt", async () => { + const createdTask = createTask(); + const createTaskMock = vi.fn().mockResolvedValue(createdTask); + const startedRun = createRun({ id: "fork-run", state: {} }); + const startedTask = createTask({ latest_run: startedRun }); + const updateTaskRunMock = vi.fn().mockResolvedValue( + createRun({ + id: "fork-run", + state: { + forked_from_task_id: "source-task", + forked_from_run_id: "source-run", + }, + }), + ); + const createTaskRunMock = vi.fn().mockResolvedValue(startedRun); + const startTaskRunMock = vi.fn().mockResolvedValue(startedTask); + const saga = makeSaga({ + createTask: createTaskMock, + createTaskRun: createTaskRunMock, + startTaskRun: startTaskRunMock, + updateTaskRun: updateTaskRunMock, + }); + + const result = await saga.run({ + content: "Ship the fix", + repository: "posthog/posthog", + workspaceMode: "cloud", + branch: "feature/source", + adapter: "codex", + sandboxEnvironmentId: "sandbox-1", + customImageId: "image-1", + forkFrom: { + kind: "cloud", + taskId: "source-task", + taskRunId: "source-run", + }, + }); + + expect(result.success).toBe(true); + expect(createTaskMock).toHaveBeenCalledWith( + expect.objectContaining({ + branch: undefined, + sandbox_environment_id: undefined, + custom_image_id: undefined, + }), + ); + expect(createTaskRunMock).toHaveBeenCalledWith( + "task-123", + expect.objectContaining({ resumeFromRunId: "source-run" }), + ); + expect(startTaskRunMock).toHaveBeenCalledWith("task-123", "fork-run", { + pendingUserMessage: undefined, + pendingUserArtifactIds: undefined, + }); + expect(mockHost.getCloudPromptTransport).not.toHaveBeenCalled(); + expect(updateTaskRunMock).toHaveBeenCalledWith("task-123", "fork-run", { + state: { + forked_from_task_id: "source-task", + forked_from_run_id: "source-run", + }, + }); + }); + + it("forks a local session into the cloned workspace", async () => { + mockHost.addFolder.mockResolvedValue({ id: "folder-1", path: "/repo" }); + mockHost.detectRepo.mockResolvedValue(null); + mockHost.createWorkspace.mockResolvedValue({ + taskId: "task-123", + mode: "worktree", + worktree: { + worktreePath: "/worktrees/fork", + worktreeName: "fork", + branchName: null, + baseBranch: "abc123", + createdAt: "2026-07-21T00:00:00Z", + }, + branchName: null, + linkedBranch: null, + }); + const task = createTask(); + const saga = makeSaga({ createTask: vi.fn().mockResolvedValue(task) }); + + const result = await saga.run({ + content: "Ship the fix", + repoPath: "/repo", + workspaceMode: "worktree", + forkFrom: { kind: "local", taskId: "source-task" }, + }); + + expect(result.success).toBe(true); + expect(mockHost.createWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ forkFromTaskId: "source-task" }), + ); + expect(sessionService.forkLocalTask).toHaveBeenCalledWith({ + sourceTaskId: "source-task", + task, + repoPath: "/worktrees/fork", + }); + expect(sessionService.connectToTask).not.toHaveBeenCalled(); + }); + it("folds channel CONTEXT.md into the cloud prompt and stashes it for the optimistic placeholder", async () => { const createdTask = createTask(); const startedTask = createTask({ latest_run: createRun() }); @@ -1212,6 +1315,36 @@ describe("TaskCreationSaga", () => { expect(onTaskReady.mock.calls[0][0].workspace).toBeNull(); }); + it("rolls back a fork when checkpoint provisioning fails", async () => { + const createdTask = createTask(); + const createTaskMock = vi.fn().mockResolvedValue(createdTask); + const deleteTaskMock = vi.fn().mockResolvedValue(undefined); + const onTaskReady = vi.fn(); + mockHost.addFolder.mockResolvedValue({ id: "folder-1", path: "/repo" }); + mockHost.detectRepo.mockResolvedValue(null); + mockHost.createWorkspace.mockRejectedValue(new Error("checkpoint failed")); + + const saga = makeSaga( + { createTask: createTaskMock, deleteTask: deleteTaskMock }, + { onTaskReady }, + ); + + const result = await saga.run({ + content: "Fork the task", + repoPath: "/repo", + workspaceMode: "worktree", + forkFrom: { kind: "local", taskId: "source-task" }, + }); + + expect(result.success).toBe(false); + expect(deleteTaskMock).toHaveBeenCalledWith(createdTask.id); + expect(mockHost.clearProvisioning).toHaveBeenCalledWith(createdTask.id); + expect(onTaskReady).toHaveBeenCalledWith({ + task: createdTask, + workspace: null, + }); + }); + it("still rolls back a worktree task when a later (non-provisioning) step fails", async () => { const createdTask = createTask(); const createTaskMock = vi.fn().mockResolvedValue(createdTask); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index f6da0ae29d..416aae82ef 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -95,7 +95,7 @@ export class TaskCreationSaga extends Saga< const importedClaude = await this.importClaudeSession(input); const warmPayload = - !taskId && input.workspaceMode === "cloud" + !taskId && input.workspaceMode === "cloud" && !input.forkFrom ? await this.prepareWarmActivation(input) : null; @@ -156,6 +156,7 @@ export class TaskCreationSaga extends Saga< branch: branch ?? undefined, allowRemoteBranchCheckout: input.allowRemoteBranchCheckout, reuseExistingWorktree: input.reuseExistingWorktree, + forkFromTaskId: input.forkFrom?.taskId, }); }, rollback: async () => { @@ -196,7 +197,12 @@ export class TaskCreationSaga extends Saga< // back here would run task_creation's deleteTask and destroy that task, // losing the prompt. Instead keep the task with no workspace (the shape // openTask re-provisions from) so the user can retry setup on it. - if (!hasProvisioning) throw error; + if (!hasProvisioning || input.forkFrom) { + if (hasProvisioning) { + this.deps.host.clearProvisioning(task.id); + } + throw error; + } const provisioningError = error instanceof Error ? error.message : String(error); this.log.error("Worktree provisioning failed; keeping task for retry", { @@ -357,6 +363,7 @@ export class TaskCreationSaga extends Saga< const buildTransport = async (): Promise => { if ( + input.forkFrom || !(input.content || input.filePaths?.length) || workspaceMode !== "cloud" ) { @@ -411,6 +418,9 @@ export class TaskCreationSaga extends Saga< homeQuickAction: input.homeQuickActionLabel, importedMcpServers: input.importedMcpServers, relayedMcpServers: input.relayedMcpServers, + ...(input.forkFrom?.kind === "cloud" && { + resumeFromRunId: input.forkFrom.taskRunId, + }), initialPermissionMode: input.executionMode ?? (cloudAdapter === "codex" ? "auto" : "plan"), @@ -452,6 +462,20 @@ export class TaskCreationSaga extends Saga< }, ); + if (input.forkFrom?.kind === "cloud" && startedRun.latest_run) { + startedRun.latest_run = await this.deps.posthogClient.updateTaskRun( + task.id, + startedRun.latest_run.id, + { + state: { + ...startedRun.latest_run.state, + forked_from_task_id: input.forkFrom.taskId, + forked_from_run_id: input.forkFrom.taskRunId, + }, + }, + ); + } + if (transport) { this.deps.track(ANALYTICS_EVENTS.PROMPT_SENT, { task_id: task.id, @@ -486,7 +510,7 @@ export class TaskCreationSaga extends Saga< if (shouldConnect) { const initialPrompt = - !input.taskId && input.content + !input.taskId && input.content && !input.forkFrom ? await this.readOnlyStep("build_prompt_blocks", () => buildPromptBlocks( input.content ?? "", @@ -527,7 +551,17 @@ export class TaskCreationSaga extends Saga< connectParams.adapter = "claude"; } - this.deps.sessionService.connectToTask(connectParams); + if (input.forkFrom?.kind === "local") { + await this.deps.sessionService.forkLocalTask({ + sourceTaskId: input.forkFrom.taskId, + task, + repoPath: agentCwd ?? "", + }); + } else if (input.forkFrom) { + throw new Error("Cloud forks require a cloud workspace"); + } else { + this.deps.sessionService.connectToTask(connectParams); + } return { taskId: task.id }; }, rollback: async ({ taskId }) => { @@ -750,6 +784,8 @@ export class TaskCreationSaga extends Saga< name: "task_creation", execute: async () => { const description = input.taskDescription ?? input.content ?? ""; + const suppressWarmReuse = + !!input.forkFrom || warmPayload?.suppressWarmReuse === true; const result = await this.deps.posthogClient.createTask({ description, repository: repository ?? undefined, @@ -769,7 +805,7 @@ export class TaskCreationSaga extends Saga< // The server associates the task with the report and records the implementation // task_run artefact — no relationship label is sent (associations are unlabelled). branch: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && !suppressWarmReuse ? (input.branch ?? null) : undefined, runtime_adapter: @@ -783,11 +819,11 @@ export class TaskCreationSaga extends Saga< ? (input.reasoningLevel ?? null) : undefined, sandbox_environment_id: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && !suppressWarmReuse ? input.sandboxEnvironmentId : undefined, custom_image_id: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && !suppressWarmReuse ? input.customImageId : undefined, signal_report: input.signalReportId ?? undefined, diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts new file mode 100644 index 0000000000..dc4660a1bc --- /dev/null +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -0,0 +1,251 @@ +import type { Workspace } from "@posthog/shared"; +import type { Task, TaskRun } from "@posthog/shared/domain-types"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { SessionService } from "../sessions/sessionService"; +import type { ITaskCreationHost } from "./taskCreationHost"; +import { TaskForkService } from "./taskForkService"; +import type { TaskService } from "./taskService"; + +const createTask = (overrides: Partial = {}): Task => ({ + id: "source-task", + task_number: 1, + slug: "source-task", + title: "Fork task", + description: "Preserve the existing context", + origin_product: "user_created", + repository: "PostHog/code", + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + latest_run: createRun(), + ...overrides, +}); + +const createRun = (overrides: Partial = {}): TaskRun => ({ + id: "source-run", + task: "source-task", + team: 1, + branch: "main", + environment: "local", + status: "completed", + log_url: "https://example.com/logs/source-run", + error_message: null, + output: null, + state: {}, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + completed_at: "2026-07-21T00:01:00Z", + ...overrides, +}); + +const createWorkspace = (overrides: Partial = {}): Workspace => ({ + taskId: "source-task", + folderId: "folder-1", + folderPath: "/repos/code", + mode: "worktree", + worktreePath: "/worktrees/source-task", + worktreeName: "source-task", + branchName: "posthog-code/source-task", + baseBranch: "main", + linkedBranch: null, + createdAt: "2026-07-21T00:00:00Z", + ...overrides, +}); + +describe("TaskForkService", () => { + const getTaskRun = vi.fn(); + const host = { + getAuthenticatedClient: vi.fn(() => ({ getTaskRun })), + getWorkspace: vi.fn(), + getAdditionalDirectories: vi.fn(), + } as unknown as ITaskCreationHost; + const taskService = { + createTask: vi.fn(), + } as unknown as TaskService; + const sessionService = { + getSessionByTaskId: vi.fn(), + } as unknown as SessionService; + const service = new TaskForkService(host, taskService, sessionService); + + beforeEach(() => { + vi.clearAllMocks(); + getTaskRun.mockResolvedValue(createRun()); + vi.mocked(host.getWorkspace).mockResolvedValue(createWorkspace()); + vi.mocked(host.getAdditionalDirectories).mockResolvedValue([]); + vi.mocked(sessionService.getSessionByTaskId).mockReturnValue(undefined); + vi.mocked(taskService.createTask).mockResolvedValue({ + success: true, + data: {} as never, + }); + }); + + it("rejects a cloud run while its current turn is active", async () => { + vi.mocked(host.getWorkspace).mockResolvedValue( + createWorkspace({ mode: "cloud", folderPath: "" }), + ); + const task = createTask({ + latest_run: createRun({ environment: "cloud", status: "in_progress" }), + }); + + const result = await service.forkTask(task); + + expect(result).toEqual({ + success: false, + error: "Wait for the current cloud turn to finish before forking it", + failedStep: "validation", + }); + }); + + it("forks a live cloud run after its current turn completes", async () => { + vi.mocked(host.getWorkspace).mockResolvedValue( + createWorkspace({ mode: "cloud", folderPath: "" }), + ); + vi.mocked(sessionService.getSessionByTaskId).mockReturnValue({ + taskRunId: "source-run", + isCloud: true, + cloudStatus: "in_progress", + isPromptPending: false, + agentIdleForRunId: "source-run", + } as never); + const task = createTask({ + latest_run: createRun({ environment: "cloud", status: "in_progress" }), + }); + + const result = await service.forkTask(task); + + expect(result.success).toBe(true); + expect(taskService.createTask).toHaveBeenCalledOnce(); + }); + + it("creates a local worktree fork", async () => { + vi.mocked(host.getAdditionalDirectories).mockResolvedValue([ + "/repos/shared", + ]); + const task = createTask({ latest_run: undefined }); + + await service.forkTask(task); + + expect(taskService.createTask).toHaveBeenCalledWith( + expect.objectContaining({ + content: "Preserve the existing context", + repoPath: "/repos/code", + repository: "PostHog/code", + workspaceMode: "worktree", + branch: null, + additionalDirectories: ["/repos/shared"], + forkFrom: { kind: "local", taskId: "source-task" }, + }), + undefined, + ); + expect(getTaskRun).not.toHaveBeenCalled(); + }); + + it("uses the live cloud run instead of a stale local latest run", async () => { + getTaskRun.mockResolvedValue( + createRun({ + id: "live-cloud-run", + environment: "cloud", + }), + ); + const task = createTask({ + latest_run: createRun({ id: "stale-local-run", environment: "local" }), + }); + vi.mocked(sessionService.getSessionByTaskId).mockReturnValue({ + taskRunId: "live-cloud-run", + isCloud: true, + cloudStatus: "completed", + isPromptPending: false, + } as never); + + const result = await service.forkTask(task); + + expect(result.success).toBe(true); + expect(taskService.createTask).toHaveBeenCalledWith( + expect.objectContaining({ + workspaceMode: "cloud", + repoPath: undefined, + forkFrom: { + kind: "cloud", + taskId: "source-task", + taskRunId: "live-cloud-run", + }, + }), + undefined, + ); + expect(host.getAdditionalDirectories).not.toHaveBeenCalled(); + }); + + it("preserves cloud runtime and sandbox options", async () => { + vi.mocked(host.getWorkspace).mockResolvedValue( + createWorkspace({ mode: "cloud", folderPath: "" }), + ); + const task = createTask({ + github_integration: 12, + github_user_integration: "integration-1", + latest_run: createRun({ + environment: "cloud", + branch: "posthog-code/source", + runtime_adapter: "claude", + model: "claude-sonnet-4-5", + reasoning_effort: "medium", + state: { + auto_publish: true, + custom_image_id: "image-1", + initial_permission_mode: "acceptEdits", + pr_authorship_mode: "bot", + rtk_enabled: false, + run_source: "signal_report", + sandbox_environment_id: "environment-1", + }, + }), + }); + + await service.forkTask(task); + + expect(taskService.createTask).toHaveBeenCalledWith( + expect.objectContaining({ + workspaceMode: "cloud", + branch: "posthog-code/source", + githubIntegrationId: 12, + githubUserIntegrationId: "integration-1", + executionMode: "acceptEdits", + adapter: "claude", + model: "claude-sonnet-4-5", + reasoningLevel: "medium", + sandboxEnvironmentId: "environment-1", + customImageId: "image-1", + cloudAutoPublish: true, + cloudRtkEnabled: false, + cloudRunSource: "signal_report", + cloudPrAuthorshipMode: "bot", + forkFrom: { + kind: "cloud", + taskId: "source-task", + taskRunId: "source-run", + }, + }), + undefined, + ); + }); + + it("keeps signal-report forks bot-authored by default", async () => { + vi.mocked(host.getWorkspace).mockResolvedValue( + createWorkspace({ mode: "cloud", folderPath: "" }), + ); + const task = createTask({ + latest_run: createRun({ + environment: "cloud", + state: { run_source: "signal_report" }, + }), + }); + + await service.forkTask(task); + + expect(taskService.createTask).toHaveBeenCalledWith( + expect.objectContaining({ + cloudRunSource: "signal_report", + cloudPrAuthorshipMode: "bot", + }), + undefined, + ); + }); +}); diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts new file mode 100644 index 0000000000..86498fc685 --- /dev/null +++ b/packages/core/src/task-detail/taskForkService.ts @@ -0,0 +1,201 @@ +import type { AgentSession, TaskCreationOutput } from "@posthog/shared"; +import { + executionModeSchema, + isTerminalStatus, + type Task, + type TaskRunStatus, +} from "@posthog/shared/domain-types"; +import { inject, injectable } from "inversify"; +import { + getCloudPrAuthorshipMode, + getCloudRunSource, +} from "../sessions/cloudRunOptions"; +import { + SESSION_SERVICE, + type SessionService, +} from "../sessions/sessionService"; +import { TASK_CREATION_HOST, TASK_SERVICE } from "./identifiers"; +import type { ITaskCreationHost } from "./taskCreationHost"; +import type { CreateTaskResult, TaskService } from "./taskService"; + +export interface ForkTaskOptions { + onTaskReady?: (output: TaskCreationOutput) => void; +} + +type CloudForkSession = Pick< + AgentSession, + | "agentIdleForRunId" + | "cloudStatus" + | "isCloud" + | "isPromptPending" + | "taskRunId" +>; + +export function canForkCloudRun( + taskRunId: string, + persistedStatus: TaskRunStatus | null | undefined, + session: CloudForkSession | null | undefined, +): boolean { + const matchingSession = + session?.isCloud === true && session.taskRunId === taskRunId + ? session + : undefined; + + return ( + isTerminalStatus(persistedStatus) || + isTerminalStatus(matchingSession?.cloudStatus) || + (matchingSession?.agentIdleForRunId === taskRunId && + !matchingSession.isPromptPending) + ); +} + +export function getForkSourceTaskId(task: Task): string | null { + const sourceTaskId = task.latest_run?.state?.forked_from_task_id; + return typeof sourceTaskId === "string" ? sourceTaskId : null; +} + +@injectable() +export class TaskForkService { + constructor( + @inject(TASK_CREATION_HOST) + private readonly host: ITaskCreationHost, + @inject(TASK_SERVICE) + private readonly taskService: TaskService, + @inject(SESSION_SERVICE) + private readonly sessionService: SessionService, + ) {} + + async forkTask( + sourceTask: Task, + options: ForkTaskOptions = {}, + ): Promise { + if (sourceTask.runtime === "pi") { + return this.validationError("Pi tasks cannot be forked yet"); + } + + const sourceWorkspace = await this.host.getWorkspace(sourceTask.id); + const sourceSession = this.sessionService.getSessionByTaskId(sourceTask.id); + const isCloud = + sourceSession?.isCloud !== undefined + ? sourceSession.isCloud + : sourceWorkspace + ? sourceWorkspace.mode === "cloud" + : sourceTask.latest_run?.environment === "cloud"; + + if (!isCloud) { + if ( + !sourceWorkspace || + sourceWorkspace.mode === "cloud" || + sourceWorkspace.isScratch + ) { + return this.validationError( + "Only repository-backed local tasks can be forked", + ); + } + const additionalDirectories = await this.host.getAdditionalDirectories( + sourceTask.id, + ); + + return this.taskService.createTask( + { + content: sourceTask.description || sourceTask.title, + taskDescription: sourceTask.description || sourceTask.title, + repoPath: sourceWorkspace.folderPath, + repository: sourceTask.repository, + workspaceMode: "worktree", + branch: null, + additionalDirectories, + forkFrom: { + kind: "local", + taskId: sourceTask.id, + }, + }, + options.onTaskReady, + ); + } + + const sourceTaskRunId = sourceSession?.isCloud + ? sourceSession.taskRunId + : sourceTask.latest_run?.id; + if (!sourceTaskRunId) { + return this.validationError("This task has no cloud run to fork"); + } + + let sourceRun = sourceTask.latest_run; + if (!sourceRun || sourceRun.id !== sourceTaskRunId) { + const client = await this.host.getAuthenticatedClient(); + if (!client) { + return this.validationError("Not authenticated"); + } + try { + sourceRun = await client.getTaskRun(sourceTask.id, sourceTaskRunId); + } catch (error) { + return { + success: false, + error: + error instanceof Error + ? error.message + : "Failed to fetch the source task run", + failedStep: "fetch_task", + }; + } + } + if (sourceRun.environment !== "cloud") { + return this.validationError("The source run is not a cloud run"); + } + if (!canForkCloudRun(sourceRun.id, sourceRun.status, sourceSession)) { + return this.validationError( + "Wait for the current cloud turn to finish before forking it", + ); + } + + const output = sourceRun.output ?? {}; + const state = sourceRun.state ?? {}; + const cloudBranch = + (typeof output.head_branch === "string" ? output.head_branch : null) ?? + sourceRun.branch ?? + (typeof state.pr_base_branch === "string" ? state.pr_base_branch : null); + + return this.taskService.createTask( + { + content: sourceTask.description || sourceTask.title, + taskDescription: sourceTask.description || sourceTask.title, + repoPath: undefined, + repository: sourceTask.repository, + workspaceMode: "cloud", + branch: cloudBranch, + githubIntegrationId: sourceTask.github_integration ?? undefined, + githubUserIntegrationId: + sourceTask.github_user_integration ?? undefined, + executionMode: executionModeSchema.safeParse( + state.initial_permission_mode, + ).data, + adapter: sourceRun.runtime_adapter ?? undefined, + model: sourceRun.model ?? undefined, + reasoningLevel: sourceRun.reasoning_effort ?? undefined, + sandboxEnvironmentId: + typeof state.sandbox_environment_id === "string" + ? state.sandbox_environment_id + : undefined, + customImageId: + typeof state.custom_image_id === "string" + ? state.custom_image_id + : undefined, + cloudAutoPublish: state.auto_publish === true, + cloudRtkEnabled: state.rtk_enabled === false ? false : undefined, + cloudRunSource: getCloudRunSource(state), + cloudPrAuthorshipMode: getCloudPrAuthorshipMode(state), + forkFrom: { + kind: "cloud", + taskId: sourceTask.id, + taskRunId: sourceRun.id, + }, + }, + options.onTaskReady, + ); + } + + private validationError(error: string): CreateTaskResult { + return { success: false, error, failedStep: "validation" }; + } +} diff --git a/packages/core/src/tasks/contextMenuActions.test.ts b/packages/core/src/tasks/contextMenuActions.test.ts index 9e06ca86fe..c66a9562f1 100644 --- a/packages/core/src/tasks/contextMenuActions.test.ts +++ b/packages/core/src/tasks/contextMenuActions.test.ts @@ -24,6 +24,12 @@ describe("resolveTaskContextMenuIntent", () => { expect(resolveTaskContextMenuIntent({ type: "stop" }, {})).toEqual({ type: "stop", }); + expect(resolveTaskContextMenuIntent({ type: "fork" }, {})).toEqual({ + type: "fork", + }); + expect(resolveTaskContextMenuIntent({ type: "open-source" }, {})).toEqual({ + type: "open-source", + }); expect(resolveTaskContextMenuIntent({ type: "delete" }, {})).toEqual({ type: "delete", }); diff --git a/packages/core/src/tasks/contextMenuActions.ts b/packages/core/src/tasks/contextMenuActions.ts index a9d12056ea..8b8a20144d 100644 --- a/packages/core/src/tasks/contextMenuActions.ts +++ b/packages/core/src/tasks/contextMenuActions.ts @@ -8,6 +8,8 @@ export type TaskContextMenuIntent = | { type: "pin" } | { type: "suspend" } | { type: "stop" } + | { type: "fork" } + | { type: "open-source" } | { type: "restore" } | { type: "archive" } | { type: "archive-prior" } @@ -29,6 +31,10 @@ export function resolveTaskContextMenuIntent( return flags.isSuspended ? { type: "restore" } : { type: "suspend" }; case "stop": return { type: "stop" }; + case "fork": + return { type: "fork" }; + case "open-source": + return { type: "open-source" }; case "archive": return { type: "archive" }; case "archive-prior": diff --git a/packages/git/src/sagas/checkpoint.test.ts b/packages/git/src/sagas/checkpoint.test.ts index e05dfcf5cd..5f4c1b5dab 100644 --- a/packages/git/src/sagas/checkpoint.test.ts +++ b/packages/git/src/sagas/checkpoint.test.ts @@ -66,18 +66,28 @@ async function getIndexFileContent( async function captureCheckpoint( repoPath: string, checkpointId: string, + maxWorktreeFileBytes?: number | null, ): Promise { const capture = new CaptureCheckpointSaga(); - const result = await capture.run({ baseDir: repoPath, checkpointId }); + const result = await capture.run({ + baseDir: repoPath, + checkpointId, + ...(maxWorktreeFileBytes !== undefined && { maxWorktreeFileBytes }), + }); expect(result.success).toBe(true); } async function revertCheckpoint( repoPath: string, checkpointId: string, + restoreCheckpointBranch?: boolean, ): Promise { const revert = new RevertCheckpointSaga(); - const result = await revert.run({ baseDir: repoPath, checkpointId }); + const result = await revert.run({ + baseDir: repoPath, + checkpointId, + restoreCheckpointBranch, + }); expect(result.success).toBe(true); } @@ -250,6 +260,37 @@ describe("checkpoint sagas", () => { }); }); + it("restores detached when the source branch is checked out elsewhere", async () => { + await withRepoAndWorktree(async (repoPath, sourceWorktreePath) => { + const sourceGit = createGitClient(sourceWorktreePath); + await writeFile(path.join(sourceWorktreePath, "a.txt"), "forked\n"); + await captureCheckpoint(sourceWorktreePath, "fork"); + + const childWorktreePath = await mkdtemp( + path.join(tmpdir(), "posthog-code-child-worktree-"), + ); + try { + const repoGit = createGitClient(repoPath); + await repoGit.raw(["worktree", "add", "--detach", childWorktreePath]); + + await revertCheckpoint(childWorktreePath, "fork", false); + + const childGit = createGitClient(childWorktreePath); + expect( + await readFile(path.join(childWorktreePath, "a.txt"), "utf8"), + ).toBe("forked\n"); + expect((await childGit.raw(["branch", "--show-current"])).trim()).toBe( + "", + ); + expect( + (await sourceGit.raw(["branch", "--show-current"])).trim(), + ).not.toBe(""); + } finally { + await rm(childWorktreePath, { recursive: true, force: true }); + } + }); + }); + it("handles submodules without breaking", { timeout: 15000 }, async () => { const subRepo = await mkdtemp( path.join(tmpdir(), "posthog-code-submodule-"), @@ -471,6 +512,31 @@ describe("checkpoint sagas", () => { }); }); + it("preserves large tracked and untracked files when the size cap is disabled", async () => { + await withRepo(async (repoPath) => { + const largeTracked = "tracked-local\n".repeat(100_000); + const largeUntracked = "untracked\n".repeat(120_000); + await writeFile(path.join(repoPath, "a.txt"), largeTracked); + await writeFile( + path.join(repoPath, "large-untracked.txt"), + largeUntracked, + ); + + await captureCheckpoint(repoPath, "lossless-large", null); + + await writeFile(path.join(repoPath, "a.txt"), "after\n"); + await rm(path.join(repoPath, "large-untracked.txt")); + await revertCheckpoint(repoPath, "lossless-large"); + + expect(await readFile(path.join(repoPath, "a.txt"), "utf8")).toBe( + largeTracked, + ); + expect( + await readFile(path.join(repoPath, "large-untracked.txt"), "utf8"), + ).toBe(largeUntracked); + }); + }); + it("does not leak temp index into normal git operations", async () => { await withRepo(async (repoPath) => { const git = createGitClient(repoPath); diff --git a/packages/git/src/sagas/checkpoint.ts b/packages/git/src/sagas/checkpoint.ts index 1b80222a3b..9fa453a6d8 100644 --- a/packages/git/src/sagas/checkpoint.ts +++ b/packages/git/src/sagas/checkpoint.ts @@ -38,6 +38,7 @@ interface CheckpointMetadata { export interface CaptureCheckpointInput extends GitSagaInput { checkpointId?: string; + maxWorktreeFileBytes?: number | null; } export interface CaptureCheckpointOutput extends CheckpointState {} @@ -76,7 +77,12 @@ export class CaptureCheckpointSaga extends GitSaga< ); const worktreeTree = await this.readOnlyStep("write_worktree_tree", () => - createWorktreeTree(this.git, baseDir, headInfo.head), + createWorktreeTree( + this.git, + baseDir, + headInfo.head, + input.maxWorktreeFileBytes, + ), ); const metaTree = await this.readOnlyStep("write_meta_tree", () => @@ -157,6 +163,7 @@ export class CaptureCheckpointSaga extends GitSaga< export interface RevertCheckpointInput extends GitSagaInput { checkpointId: string; + restoreCheckpointBranch?: boolean; } export interface RevertCheckpointOutput { @@ -210,7 +217,7 @@ export class RevertCheckpointSaga extends GitSaga< name: "checkout_head", execute: async () => { if (!head) return; - if (branch) { + if (input.restoreCheckpointBranch !== false && branch) { const branchExists = await refExists( this.git, `refs/heads/${branch}`, @@ -398,6 +405,7 @@ async function createWorktreeTree( git: GitClient, baseDir: string, head: string | null, + maxWorktreeFileBytes: number | null | undefined = MAX_WORKTREE_FILE_BYTES, ): Promise { const { tempGit, tempIndexPath } = await createTempIndexGit( git, @@ -413,7 +421,13 @@ async function createWorktreeTree( } await tempGit.raw(["add", "-A", "--", "."]); - await reconcileLargeBlobs(tempGit, head, MAX_WORKTREE_FILE_BYTES); + if (maxWorktreeFileBytes !== null) { + await reconcileLargeBlobs( + tempGit, + head, + maxWorktreeFileBytes ?? MAX_WORKTREE_FILE_BYTES, + ); + } const treeHash = await tempGit.raw(["write-tree"]); return treeHash.trim(); } finally { diff --git a/packages/host-router/src/routers/agent.router.ts b/packages/host-router/src/routers/agent.router.ts index f7a7a43f52..de926d0db7 100644 --- a/packages/host-router/src/routers/agent.router.ts +++ b/packages/host-router/src/routers/agent.router.ts @@ -9,6 +9,7 @@ import { cancelPermissionInput, cancelPromptInput, cancelSessionInput, + forkSessionInput, getGatewayModelsInput, getGatewayModelsOutput, getPreviewConfigOptionsInput, @@ -40,6 +41,21 @@ export const agentRouter = router({ ctx.container.get(AGENT_SERVICE).startSession(input), ), + fork: publicProcedure + .input(forkSessionInput) + .output(sessionResponseSchema) + .mutation(({ ctx, input }) => + ctx.container.get(AGENT_SERVICE).forkSession(input), + ), + + flushLogs: publicProcedure + .input(subscribeSessionInput) + .mutation(({ ctx, input }) => + ctx.container + .get(AGENT_SERVICE) + .flushSessionLogs(input.taskRunId), + ), + prompt: publicProcedure .input(promptInput) .output(promptOutput) diff --git a/packages/host-router/src/routers/logs.router.ts b/packages/host-router/src/routers/logs.router.ts index 94a42bd378..d203eceb22 100644 --- a/packages/host-router/src/routers/logs.router.ts +++ b/packages/host-router/src/routers/logs.router.ts @@ -2,6 +2,7 @@ import { publicProcedure, router } from "@posthog/host-trpc/trpc"; import type { ILogsService } from "@posthog/workspace-server/services/local-logs/identifiers"; import { LOGS_SERVICE } from "@posthog/workspace-server/services/local-logs/identifiers"; import { + cloneLocalLogsInput, fetchS3LogsInput, fetchS3LogsOutput, readLocalLogsCollapsedInput, @@ -10,6 +11,7 @@ import { readLocalLogsOutput, readLocalLogsTailInput, readLocalLogsTailOutput, + seedLocalLogsInput, writeLocalLogsInput, } from "@posthog/workspace-server/services/local-logs/schemas"; @@ -55,4 +57,20 @@ export const logsRouter = router({ .get(LOGS_SERVICE) .writeLocalLogs(input.taskRunId, input.content), ), + + seedLocalLogs: publicProcedure + .input(seedLocalLogsInput) + .mutation(({ ctx, input }) => + ctx.container + .get(LOGS_SERVICE) + .seedLocalLogs(input.taskRunId, input.content), + ), + + cloneLocalLogs: publicProcedure + .input(cloneLocalLogsInput) + .mutation(({ ctx, input }) => + ctx.container + .get(LOGS_SERVICE) + .cloneLocalLogs(input.sourceTaskRunId, input.targetTaskRunId), + ), }); diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 0c354e146a..8c85c74c6d 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -109,6 +109,17 @@ export interface TaskCreationInput { * session last worked on, linked so the branch-mismatch prompt can fire. */ importedClaudeSession?: { sourceSessionId: string; branch?: string | null }; + /** Create an independent task from a local session or cloud run snapshot. */ + forkFrom?: + | { + kind: "local"; + taskId: string; + } + | { + kind: "cloud"; + taskId: string; + taskRunId: string; + }; } export interface TaskCreationOutput { diff --git a/packages/ui/src/features/sidebar/components/SidebarMenu.tsx b/packages/ui/src/features/sidebar/components/SidebarMenu.tsx index 3be10a1f11..47f98ffb23 100644 --- a/packages/ui/src/features/sidebar/components/SidebarMenu.tsx +++ b/packages/ui/src/features/sidebar/components/SidebarMenu.tsx @@ -1,5 +1,6 @@ import { findGroupFolder } from "@posthog/core/sidebar/groupTasks"; import { isTaskActivelyRunning } from "@posthog/core/sidebar/taskRunning"; +import { getForkSourceTaskId } from "@posthog/core/task-detail/taskForkService"; import { useHostTRPCClient } from "@posthog/host-router/react"; import type { Task } from "@posthog/shared/types"; import { @@ -18,6 +19,7 @@ import { usePinnedTasks } from "@posthog/ui/features/sidebar/usePinnedTasks"; import { useSidebarData } from "@posthog/ui/features/sidebar/useSidebarData"; import { useTaskViewed } from "@posthog/ui/features/sidebar/useTaskViewed"; import { useTaskContextMenu } from "@posthog/ui/features/tasks/useTaskContextMenu"; +import { useTaskFork } from "@posthog/ui/features/tasks/useTaskFork"; import { useRenameTask } from "@posthog/ui/features/tasks/useTaskMutations"; import { useTasks } from "@posthog/ui/features/tasks/useTasks"; import { useWorkspaces } from "@posthog/ui/features/workspace/useWorkspace"; @@ -67,6 +69,7 @@ function SidebarMenuComponent() { const { showContextMenu, editingTaskId, setEditingTaskId } = useTaskContextMenu(); + const { forkTask } = useTaskFork(); const { archiveTask } = useArchiveTask(); const { renameTask } = useRenameTask(); const { togglePin } = usePinnedTasks(); @@ -302,9 +305,11 @@ function SidebarMenuComponent() { } const taskData = allSidebarTasks.find((t) => t.id === taskId); - const task = taskMap.get(taskId) ?? taskData; + const fullTask = taskMap.get(taskId); + const task = fullTask ?? taskData; if (task) { - const runId = taskMap.get(taskId)?.latest_run?.id; + const runId = fullTask?.latest_run?.id; + const sourceTaskId = fullTask ? getForkSourceTaskId(fullTask) : null; const workspace = workspaces[taskId]; const isInCommandCenter = commandCenterCells.some( (id) => id === taskId && taskMap.has(id), @@ -321,6 +326,8 @@ function SidebarMenuComponent() { canStop: taskData?.taskRunEnvironment === "cloud" && isTaskActivelyRunning(taskData), + canFork: !!fullTask && fullTask.runtime !== "pi", + hasSourceTask: !!sourceTaskId, runId, isInCommandCenter, hasEmptyCommandCenterCell, @@ -331,6 +338,12 @@ function SidebarMenuComponent() { taskTitle, runId: stopRunId, }), + onFork: () => { + if (fullTask) void forkTask(fullTask); + }, + onOpenSource: () => { + if (sourceTaskId) navigateToTaskDetail(sourceTaskId); + }, onArchive: handleTaskArchive, onArchivePrior: handleArchivePrior, onAddToCommandCenter: () => { diff --git a/packages/ui/src/features/task-detail/taskCreationHostImpl.ts b/packages/ui/src/features/task-detail/taskCreationHostImpl.ts index e4d99882e7..0d6a4437bd 100644 --- a/packages/ui/src/features/task-detail/taskCreationHostImpl.ts +++ b/packages/ui/src/features/task-detail/taskCreationHostImpl.ts @@ -121,6 +121,10 @@ export class TrpcTaskCreationHost implements ITaskCreationHost { await hostClient().additionalDirectories.addForTask.mutate(args); } + getAdditionalDirectories(taskId: string): Promise { + return hostClient().additionalDirectories.listForTask.query({ taskId }); + } + async removeAdditionalDirectory(args: { taskId: string; path: string; diff --git a/packages/ui/src/features/tasks/retainedTasks.test.ts b/packages/ui/src/features/tasks/retainedTasks.test.ts new file mode 100644 index 0000000000..3fc196e95e --- /dev/null +++ b/packages/ui/src/features/tasks/retainedTasks.test.ts @@ -0,0 +1,34 @@ +import type { Task } from "@posthog/shared/domain-types"; +import { afterEach, describe, expect, it } from "vitest"; +import { + mergeRetainedTasks, + releaseRetainedTask, + retainTask, +} from "./retainedTasks"; + +const task = { id: "forked-task", title: "Forked task" } as Task; + +describe("retainedTasks", () => { + afterEach(() => releaseRetainedTask(task.id)); + + it("keeps a fork visible while the server list is stale", () => { + retainTask(task); + + expect(mergeRetainedTasks([])).toEqual([task]); + }); + + it("reconciles to the server task once it appears", () => { + const serverTask = { ...task, title: "Server title" }; + retainTask(task); + + expect(mergeRetainedTasks([serverTask])).toEqual([serverTask]); + expect(mergeRetainedTasks([])).toEqual([]); + }); + + it("drops a retained task after rollback", () => { + retainTask(task); + releaseRetainedTask(task.id); + + expect(mergeRetainedTasks([])).toEqual([]); + }); +}); diff --git a/packages/ui/src/features/tasks/retainedTasks.ts b/packages/ui/src/features/tasks/retainedTasks.ts new file mode 100644 index 0000000000..cf1aaf33d5 --- /dev/null +++ b/packages/ui/src/features/tasks/retainedTasks.ts @@ -0,0 +1,38 @@ +import type { Task } from "@posthog/shared/domain-types"; + +const RETENTION_MS = 5 * 60_000; +const retainedTasks = new Map(); + +function pruneExpired(now: number): void { + for (const [taskId, retained] of retainedTasks) { + if (now - retained.retainedAt >= RETENTION_MS) { + retainedTasks.delete(taskId); + } + } +} + +export function retainTask(task: Task): void { + const now = Date.now(); + pruneExpired(now); + retainedTasks.set(task.id, { task, retainedAt: now }); +} + +export function releaseRetainedTask(taskId: string): void { + retainedTasks.delete(taskId); +} + +export function mergeRetainedTasks(tasks: Task[]): Task[] { + pruneExpired(Date.now()); + const serverTaskIds = new Set(tasks.map((task) => task.id)); + const missing: Task[] = []; + + for (const [taskId, retained] of retainedTasks) { + if (serverTaskIds.has(taskId)) { + retainedTasks.delete(taskId); + } else { + missing.push(retained.task); + } + } + + return missing.length > 0 ? [...missing, ...tasks] : tasks; +} diff --git a/packages/ui/src/features/tasks/useTaskContextMenu.ts b/packages/ui/src/features/tasks/useTaskContextMenu.ts index 2aa6a32b4e..32e4c5ce2e 100644 --- a/packages/ui/src/features/tasks/useTaskContextMenu.ts +++ b/packages/ui/src/features/tasks/useTaskContextMenu.ts @@ -46,11 +46,15 @@ export function useTaskContextMenu() { isPinned?: boolean; isSuspended?: boolean; canStop?: boolean; + canFork?: boolean; + hasSourceTask?: boolean; runId?: string; isInCommandCenter?: boolean; hasEmptyCommandCenterCell?: boolean; onTogglePin?: () => void; onStop?: (taskId: string, taskTitle: string, runId?: string) => void; + onFork?: () => void; + onOpenSource?: () => void; onArchive?: (taskId: string) => void; onArchivePrior?: (taskId: string) => void; onAddToCommandCenter?: () => void; @@ -65,11 +69,15 @@ export function useTaskContextMenu() { isPinned, isSuspended, canStop, + canFork, + hasSourceTask, runId, isInCommandCenter, hasEmptyCommandCenterCell, onTogglePin, onStop, + onFork, + onOpenSource, onArchive, onArchivePrior, onAddToCommandCenter, @@ -83,6 +91,8 @@ export function useTaskContextMenu() { isPinned, isSuspended, canStop, + canFork, + hasSourceTask, isInCommandCenter, hasEmptyCommandCenterCell, channels: channels.map(({ id, name }) => ({ id, name })), @@ -111,6 +121,12 @@ export function useTaskContextMenu() { onStop?.(task.id, task.title, runId); break; } + case "fork": + onFork?.(); + break; + case "open-source": + onOpenSource?.(); + break; case "archive": if (onArchive) { onArchive(task.id); diff --git a/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx b/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx index 76bde2513e..fc0280a7ac 100644 --- a/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx +++ b/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx @@ -255,3 +255,42 @@ describe("useCreateTask.invalidateTasks", () => { expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: taskKeys.lists() }); }); }); + +describe("useCreateTask.seedTask", () => { + it("retains a newly created task without immediately refetching the list", () => { + const queryClient = new QueryClient(); + const key = taskKeys.list(); + queryClient.setQueryData(key, []); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const localWrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + const { result } = renderHook(() => useCreateTask(), { + wrapper: localWrapper, + }); + + result.current.seedTask(createTask()); + + expect(queryClient.getQueryData(key)).toHaveLength(1); + expect(invalidateSpy).not.toHaveBeenCalled(); + }); + + it("removes a seeded task when its creation rolls back", () => { + const queryClient = new QueryClient(); + const key = taskKeys.list(); + const task = createTask(); + queryClient.setQueryData(key, [task]); + + const localWrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + const { result } = renderHook(() => useCreateTask(), { + wrapper: localWrapper, + }); + + result.current.removeSeededTask(task.id); + + expect(queryClient.getQueryData(key)).toEqual([]); + }); +}); diff --git a/packages/ui/src/features/tasks/useTaskCrudMutations.ts b/packages/ui/src/features/tasks/useTaskCrudMutations.ts index 628dde456f..7b1898b1a3 100644 --- a/packages/ui/src/features/tasks/useTaskCrudMutations.ts +++ b/packages/ui/src/features/tasks/useTaskCrudMutations.ts @@ -17,6 +17,7 @@ import { useAuthenticatedMutation } from "@posthog/ui/hooks/useAuthenticatedMuta import { logger } from "@posthog/ui/shell/logger"; import { useQueryClient } from "@tanstack/react-query"; import { useCallback } from "react"; +import { releaseRetainedTask, retainTask } from "./retainedTasks"; import { taskKeys } from "./taskKeys"; const log = logger.scope("tasks"); @@ -42,30 +43,40 @@ export async function releaseDeletedTaskResources( export function useCreateTask() { const queryClient = useQueryClient(); - const invalidateTasks = (newTask?: Task) => { - if (newTask) { - // Only seed list caches that aren't scoped to a specific origin_product. - // An origin-scoped list (e.g. the slack-origin list behind useSlackTasks) - // is read by the sidebar to brand a task's icon by id membership, so - // seeding a freshly created, non-slack task into it would make that task - // briefly render as a Slack task until the list refetches. Origin-less - // lists, by contrast, should mirror every new task. - queryClient.setQueriesData( - { - queryKey: taskKeys.lists(), - predicate: (query) => { - const isOriginScopedList = Boolean( - taskKeys.filtersOf(query.queryKey)?.originProduct, - ); - return !isOriginScopedList; - }, + const seedTask = (newTask: Task) => { + retainTask(newTask); + // Only seed list caches that aren't scoped to a specific origin_product. + // An origin-scoped list (e.g. the slack-origin list behind useSlackTasks) + // is read by the sidebar to brand a task's icon by id membership, so + // seeding a freshly created, non-slack task into it would make that task + // briefly render as a Slack task until the list refetches. Origin-less + // lists, by contrast, should mirror every new task. + queryClient.setQueriesData( + { + queryKey: taskKeys.lists(), + predicate: (query) => { + const isOriginScopedList = Boolean( + taskKeys.filtersOf(query.queryKey)?.originProduct, + ); + return !isOriginScopedList; }, - (old) => insertTaskDedup(old, newTask), - ); - } + }, + (old) => insertTaskDedup(old, newTask), + ); + }; + + const invalidateTasks = (newTask?: Task) => { + if (newTask) seedTask(newTask); queryClient.invalidateQueries({ queryKey: taskKeys.lists() }); }; + const removeSeededTask = (taskId: string) => { + releaseRetainedTask(taskId); + queryClient.setQueriesData({ queryKey: taskKeys.lists() }, (old) => + removeTaskFromList(old, taskId), + ); + }; + const mutation = useAuthenticatedMutation( ( client, @@ -87,7 +98,7 @@ export function useCreateTask() { }) as unknown as Promise, ); - return { ...mutation, invalidateTasks }; + return { ...mutation, invalidateTasks, removeSeededTask, seedTask }; } interface DeleteTaskOptions { diff --git a/packages/ui/src/features/tasks/useTaskFork.test.tsx b/packages/ui/src/features/tasks/useTaskFork.test.tsx new file mode 100644 index 0000000000..34e5b7a000 --- /dev/null +++ b/packages/ui/src/features/tasks/useTaskFork.test.tsx @@ -0,0 +1,108 @@ +import type { Task } from "@posthog/shared"; +import { act, renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + forkTask: vi.fn(), + openTask: vi.fn(), + removeSeededTask: vi.fn(), + seedTask: vi.fn(), + setQueryData: vi.fn(), + setProvisioningFailed: vi.fn(), + toast: vi.fn(), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: () => ({ forkTask: mocks.forkTask }), +})); +vi.mock("@posthog/host-router/react", () => ({ + useHostTRPC: () => ({ + workspace: { getAll: { queryKey: () => ["workspace", "getAll"] } }, + }), +})); +vi.mock("@tanstack/react-query", () => ({ + useQueryClient: () => ({ setQueryData: mocks.setQueryData }), +})); +vi.mock("@posthog/ui/features/tasks/useTaskCrudMutations", () => ({ + useCreateTask: () => ({ + removeSeededTask: mocks.removeSeededTask, + seedTask: mocks.seedTask, + }), +})); +vi.mock("@posthog/ui/features/provisioning/store", () => ({ + useProvisioningStore: { + getState: () => ({ setFailed: mocks.setProvisioningFailed }), + }, +})); +vi.mock("@posthog/ui/features/notifications/errorDetails", () => ({ + toastError: vi.fn(), +})); +vi.mock("@posthog/ui/router/useOpenTask", () => ({ + openTask: mocks.openTask, +})); +vi.mock("@posthog/ui/primitives/toast", () => ({ + toast: { error: mocks.toast }, +})); + +import { useTaskFork } from "./useTaskFork"; + +const sourceTask = { id: "source", title: "Source" } as Task; +const childTask = { id: "child", title: "Child" } as Task; + +describe("useTaskFork", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("retains the child from task creation through final navigation", async () => { + mocks.forkTask.mockImplementation( + async ( + _task: Task, + options: { + onTaskReady: (output: { + task: Task; + workspace: { taskId: string; mode: string } | null; + }) => void; + }, + ) => { + options.onTaskReady({ task: childTask, workspace: null }); + return { + success: true, + data: { + task: childTask, + workspace: { taskId: "child", mode: "worktree" }, + }, + }; + }, + ); + const { result } = renderHook(() => useTaskFork()); + + await act(() => result.current.forkTask(sourceTask)); + + expect(mocks.seedTask).toHaveBeenNthCalledWith(1, childTask); + expect(mocks.seedTask).toHaveBeenNthCalledWith(2, childTask); + expect(mocks.removeSeededTask).not.toHaveBeenCalled(); + expect(mocks.openTask).toHaveBeenCalledWith(childTask); + expect(mocks.openTask.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.seedTask.mock.invocationCallOrder[1], + ); + }); + + it("removes the retained child when fork creation rolls back", async () => { + mocks.forkTask.mockImplementation( + async ( + _task: Task, + options: { onTaskReady: (output: { task: Task }) => void }, + ) => { + options.onTaskReady({ task: childTask }); + return { success: false, error: "fork failed" }; + }, + ); + const { result } = renderHook(() => useTaskFork()); + + await act(() => result.current.forkTask(sourceTask)); + + expect(mocks.removeSeededTask).toHaveBeenCalledWith("child"); + expect(mocks.openTask).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/ui/src/features/tasks/useTaskFork.ts b/packages/ui/src/features/tasks/useTaskFork.ts new file mode 100644 index 0000000000..1ef267b952 --- /dev/null +++ b/packages/ui/src/features/tasks/useTaskFork.ts @@ -0,0 +1,75 @@ +import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; +import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; +import { getErrorTitle } from "@posthog/core/task-detail/taskInput"; +import { useService } from "@posthog/di/react"; +import { useHostTRPC } from "@posthog/host-router/react"; +import type { Task, Workspace } from "@posthog/shared"; +import { toastError } from "@posthog/ui/features/notifications/errorDetails"; +import { useProvisioningStore } from "@posthog/ui/features/provisioning/store"; +import { useCreateTask } from "@posthog/ui/features/tasks/useTaskCrudMutations"; +import { toast } from "@posthog/ui/primitives/toast"; +import { openTask } from "@posthog/ui/router/useOpenTask"; +import { useQueryClient } from "@tanstack/react-query"; +import { useCallback } from "react"; + +export function useTaskFork() { + const taskForkService = useService(TASK_FORK_SERVICE); + const trpc = useHostTRPC(); + const queryClient = useQueryClient(); + const { removeSeededTask, seedTask } = useCreateTask(); + + const seedFork = useCallback( + (task: Task, workspace: Workspace | null) => { + seedTask(task); + if (workspace) { + queryClient.setQueryData>( + trpc.workspace.getAll.queryKey(), + (workspaces) => ({ + ...workspaces, + [task.id]: workspace, + }), + ); + } + }, + [queryClient, seedTask, trpc], + ); + + const forkTask = useCallback( + async (sourceTask: Task): Promise => { + let seededTaskId: string | null = null; + try { + const result = await taskForkService.forkTask(sourceTask, { + onTaskReady: ({ task, workspace }) => { + seededTaskId = task.id; + seedFork(task, workspace); + }, + }); + if (!result.success) { + if (seededTaskId) removeSeededTask(seededTaskId); + toast.error("Could not fork task", { description: result.error }); + return; + } + + seedFork(result.data.task, result.data.workspace); + if (result.data.provisioningError) { + useProvisioningStore + .getState() + .setFailed(result.data.task.id, result.data.provisioningError); + toastError( + getErrorTitle("workspace_creation"), + result.data.provisioningError, + ); + } + await openTask(result.data.task); + } catch (error) { + if (seededTaskId) removeSeededTask(seededTaskId); + toast.error("Could not fork task", { + description: error instanceof Error ? error.message : String(error), + }); + } + }, + [removeSeededTask, seedFork, taskForkService], + ); + + return { forkTask }; +} diff --git a/packages/ui/src/features/tasks/useTasks.ts b/packages/ui/src/features/tasks/useTasks.ts index d902484060..41e10abcfc 100644 --- a/packages/ui/src/features/tasks/useTasks.ts +++ b/packages/ui/src/features/tasks/useTasks.ts @@ -3,6 +3,7 @@ import type { Task } from "@posthog/shared/domain-types"; import { keepPreviousData } from "@tanstack/react-query"; import { useAuthenticatedQuery } from "../../hooks/useAuthenticatedQuery"; import { useMeQuery } from "../auth/useMeQuery"; +import { mergeRetainedTasks } from "./retainedTasks"; import { taskKeys } from "./taskKeys"; // Full-task polls are heavy (~630KB per response at 100 tasks — descriptions @@ -33,12 +34,14 @@ export function useTasks( return useAuthenticatedQuery( taskKeys.list({ repository: filters?.repository, createdBy, internal }), - (client) => - client.getTasks({ - repository: filters?.repository, - createdBy, - internal, - }) as unknown as Promise, + async (client) => + mergeRetainedTasks( + (await client.getTasks({ + repository: filters?.repository, + createdBy, + internal, + })) as unknown as Task[], + ), { enabled: (options?.enabled ?? true) && !!currentUser?.id, refetchInterval: TASK_LIST_POLL_INTERVAL_MS, diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index 7d479b71df..e1cb89a820 100644 --- a/packages/workspace-server/src/services/agent/agent.test.ts +++ b/packages/workspace-server/src/services/agent/agent.test.ts @@ -16,6 +16,13 @@ const mockNewSession = vi.hoisted(() => }), ); +const mockForkSession = vi.hoisted(() => + vi.fn().mockResolvedValue({ + sessionId: "forked-session-id", + configOptions: [], + }), +); + const mockAcpClient = vi.hoisted(() => ({ current: undefined as | { @@ -39,6 +46,7 @@ const mockClientSideConnection = vi.hoisted(() => mockAcpClient.current = clientFactory({}); this.initialize = vi.fn().mockResolvedValue({}); this.newSession = mockNewSession; + this.unstable_forkSession = mockForkSession; this.loadSession = vi.fn().mockResolvedValue({ configOptions: [] }); this.resumeSession = vi.fn().mockResolvedValue({ configOptions: [] }); this.setSessionConfigOption = vi.fn( @@ -276,6 +284,61 @@ describe("AgentService", () => { }); describe("MCP servers", () => { + it("forks the source agent session into the target run", async () => { + await service.startSession({ + ...baseSessionParams, + adapter: "claude", + }); + + await service.forkSession({ + ...baseSessionParams, + taskId: "task-2", + taskRunId: "run-2", + sourceTaskRunId: "run-1", + adapter: "claude", + }); + + expect(mockForkSession).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: "test-session-id", + cwd: "/mock/repo", + _meta: expect.objectContaining({ + taskRunId: "run-2", + environment: "local", + }), + }), + ); + expect(mockNewSession).toHaveBeenCalledTimes(1); + }); + + it("retries a fork as a fork after an authentication error", async () => { + await service.startSession({ + ...baseSessionParams, + adapter: "claude", + }); + mockForkSession.mockRejectedValueOnce( + new Error("Authentication required"), + ); + + await service.forkSession({ + ...baseSessionParams, + taskId: "task-2", + taskRunId: "run-2", + sourceTaskRunId: "run-1", + adapter: "claude", + }); + + expect(mockForkSession).toHaveBeenCalledTimes(2); + expect(mockForkSession).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + sessionId: "test-session-id", + _meta: expect.objectContaining({ taskRunId: "run-2" }), + }), + ); + expect(mockNewSession).toHaveBeenCalledTimes(1); + }); + it("marks desktop sessions as local even though they have a taskRunId", async () => { await service.startSession({ ...baseSessionParams, diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index 9af21e716d..646c20c36a 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -88,7 +88,11 @@ import type { ProcessTrackingService } from "../process-tracking/process-trackin import { loadSessionEnvOverrides } from "../session-env/loader"; import { isScratchPath } from "../workspace/scratch"; import type { AgentAuthAdapter, McpToolInstallations } from "./auth-adapter"; -import { cleanupCodexHome, prepareCodexHome } from "./codex-home"; +import { + cleanupCodexHome, + copyCodexSessionState, + prepareCodexHome, +} from "./codex-home"; import { discoverExternalPlugins } from "./discover-plugins"; import { AGENT_AUTH_ADAPTER, @@ -109,6 +113,7 @@ import { type AgentServiceEvents, type Credentials, type EffortLevel, + type ForkSessionInput, type InterruptReason, type PromptOutput, type ReconnectSessionInput, @@ -331,6 +336,11 @@ interface ManagedSession { prAttachChain: Promise; } +interface ForkSessionSource { + sessionId: string; + taskRunId: string; +} + /** Get the agent session ID from a managed session, throwing if not set. */ function getAgentSessionId(session: ManagedSession): string { const { sessionId } = session.config; @@ -698,6 +708,24 @@ If a repository IS genuinely required, attach one in this priority order: return this.toSessionResponse(session); } + async forkSession(params: ForkSessionInput): Promise { + this.validateSessionParams(params); + const source = this.sessions.get(params.sourceTaskRunId); + if (!source) { + throw new Error("The source agent session is not connected"); + } + if (source.promptPending) { + throw new Error("Wait for the source task to finish before forking it"); + } + + const config = this.toSessionConfig(params); + const session = await this.getOrCreateSession(config, false, false, { + sessionId: getAgentSessionId(source), + taskRunId: source.taskRunId, + }); + return this.toSessionResponse(session); + } + async reconnectSession( params: ReconnectSessionInput, ): Promise { @@ -717,16 +745,19 @@ If a repository IS genuinely required, attach one in this priority order: config: SessionConfig, isReconnect: false, isRetry?: boolean, + forkSource?: ForkSessionSource, ): Promise; private async getOrCreateSession( config: SessionConfig, isReconnect: true, isRetry?: boolean, + forkSource?: ForkSessionSource, ): Promise; private async getOrCreateSession( config: SessionConfig, isReconnect: boolean, isRetry = false, + forkSource?: ForkSessionSource, ): Promise { const { taskId, @@ -835,6 +866,13 @@ If a repository IS genuinely required, attach one in this priority order: bundledSkillsDir, log: this.log, }); + if (forkSource) { + await copyCodexSessionState({ + appDataPath: this.storagePaths.appDataPath, + sourceTaskRunId: forkSource.taskRunId, + targetTaskRunId: taskRunId, + }); + } } catch (err) { // A skills-prep failure must not kill the session; Codex falls back // to its default home and the user's own ~/.agents/skills. @@ -973,11 +1011,50 @@ If a repository IS genuinely required, attach one in this priority order: let configOptions: SessionConfigOption[] | undefined; let agentSessionId: string | undefined; + if (forkSource) { + const liveForkSource = this.sessions.get(forkSource.taskRunId); + if (!liveForkSource) { + throw new Error("The source agent session is not connected"); + } + if (liveForkSource.promptPending) { + throw new Error("Wait for the source session to become idle"); + } + const forkResponse = await connection.unstable_forkSession({ + sessionId: getAgentSessionId(liveForkSource), + cwd: repoPath, + mcpServers: sessionMcpServers, + _meta: { + ...(logUrl && { + persistence: { taskId, runId: taskRunId, logUrl }, + }), + taskRunId, + environment: "local", + systemPrompt, + ...(channelMode && { channelMode }), + ...(config.spokenNarration !== undefined && { + spokenNarration: config.spokenNarration, + }), + mcpToolApprovals: toolApprovals, + ...(permissionMode && { permissionMode }), + ...(model != null && { model }), + ...(jsonSchema && { jsonSchema }), + claudeCode: { options: claudeCodeOptions }, + }, + }); + configOptions = forkResponse.configOptions ?? undefined; + agentSessionId = forkResponse.sessionId; + } + // Imported Claude Code CLI session: the transcript JSONL was copied // into CLAUDE_CONFIG_DIR at import time, so load it directly and let // the adapter replay its history to the client. On failure, fall // through to a fresh session so the task still starts. - if (!isReconnect && config.importedSessionId && adapter !== "codex") { + if ( + !forkSource && + !isReconnect && + config.importedSessionId && + adapter !== "codex" + ) { const importedSessionId = config.importedSessionId; try { const loadResponse = await connection.loadSession({ @@ -1150,6 +1227,12 @@ If a repository IS genuinely required, attach one in this priority order: taskRunId, }); } + await cleanupCodexHome(this.storagePaths.appDataPath, taskRunId).catch( + () => + this.log.debug("Codex home cleanup failed during error handling", { + taskRunId, + }), + ); if (!isRetry && isAuthError(err)) { this.log.warn( @@ -1157,9 +1240,9 @@ If a repository IS genuinely required, attach one in this priority order: { taskRunId }, ); if (isReconnect) { - return this.getOrCreateSession(config, true, true); + return this.getOrCreateSession(config, true, true, forkSource); } - return this.getOrCreateSession(config, false, true); + return this.getOrCreateSession(config, false, true, forkSource); } // When the in-process ACP layer masks a thrown error as a generic // "Internal error", the real text survives in `data.details`. Surface it @@ -1392,6 +1475,14 @@ If a repository IS genuinely required, attach one in this priority order: return this.sessions.get(taskRunId); } + async flushSessionLogs(taskRunId: string): Promise { + const session = this.sessions.get(taskRunId); + if (!session) { + throw new Error("The agent session is not connected"); + } + await session.agent.flushAllLogs(); + } + getDebugSnapshot(): { sessions: Array<{ taskRunId: string; diff --git a/packages/workspace-server/src/services/agent/codex-home.test.ts b/packages/workspace-server/src/services/agent/codex-home.test.ts index 985c64a401..4acef2c93c 100644 --- a/packages/workspace-server/src/services/agent/codex-home.test.ts +++ b/packages/workspace-server/src/services/agent/codex-home.test.ts @@ -14,6 +14,7 @@ vi.mock("node:os", async (importOriginal) => { import { cleanupCodexHome, + copyCodexSessionState, getCodexHomeDir, prepareCodexHome, } from "./codex-home"; @@ -188,3 +189,43 @@ describe("prepareCodexHome", () => { expect(existsSync(path.join(outside, "precious", "SKILL.md"))).toBe(true); }); }); + +describe("copyCodexSessionState", () => { + it("copies sessions and shell snapshots from the source run only", async () => { + const sourceHome = getCodexHomeDir(appDataPath, "source-run"); + const otherHome = getCodexHomeDir(appDataPath, "other-run"); + const targetHome = getCodexHomeDir(appDataPath, "target-run"); + await mkdir(path.join(sourceHome, "sessions"), { recursive: true }); + await mkdir(path.join(sourceHome, "shell_snapshots"), { recursive: true }); + await mkdir(path.join(otherHome, "sessions"), { recursive: true }); + await mkdir(targetHome, { recursive: true }); + await writeFile( + path.join(sourceHome, "sessions", "session.jsonl"), + "source", + ); + await writeFile( + path.join(sourceHome, "shell_snapshots", "snapshot.sh"), + "source snapshot", + ); + await writeFile(path.join(otherHome, "sessions", "other.jsonl"), "other"); + + await copyCodexSessionState({ + appDataPath, + sourceTaskRunId: "source-run", + targetTaskRunId: "target-run", + }); + + expect( + readFileSync(path.join(targetHome, "sessions", "session.jsonl"), "utf-8"), + ).toBe("source"); + expect( + readFileSync( + path.join(targetHome, "shell_snapshots", "snapshot.sh"), + "utf-8", + ), + ).toBe("source snapshot"); + expect(existsSync(path.join(targetHome, "sessions", "other.jsonl"))).toBe( + false, + ); + }); +}); diff --git a/packages/workspace-server/src/services/agent/codex-home.ts b/packages/workspace-server/src/services/agent/codex-home.ts index 520281520c..9fd28ffbbe 100644 --- a/packages/workspace-server/src/services/agent/codex-home.ts +++ b/packages/workspace-server/src/services/agent/codex-home.ts @@ -1,6 +1,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; +import { clonePath } from "@posthog/git/utils"; import { findSkillDirs, getUserSkillsDir, @@ -94,3 +95,23 @@ export async function prepareCodexHome(options: { return codexHome; } + +export async function copyCodexSessionState(options: { + appDataPath: string; + sourceTaskRunId: string; + targetTaskRunId: string; +}): Promise { + const sourceHome = getCodexHomeDir( + options.appDataPath, + options.sourceTaskRunId, + ); + const targetHome = getCodexHomeDir( + options.appDataPath, + options.targetTaskRunId, + ); + + for (const directory of ["sessions", "shell_snapshots"]) { + const source = path.join(sourceHome, directory); + await clonePath(source, path.join(targetHome, directory)); + } +} diff --git a/packages/workspace-server/src/services/agent/schemas.ts b/packages/workspace-server/src/services/agent/schemas.ts index 5d8eeab508..14c2cffb48 100644 --- a/packages/workspace-server/src/services/agent/schemas.ts +++ b/packages/workspace-server/src/services/agent/schemas.ts @@ -102,6 +102,12 @@ export const startSessionInput = z.object({ export type StartSessionInput = z.infer; +export const forkSessionInput = startSessionInput.extend({ + sourceTaskRunId: z.string(), +}); + +export type ForkSessionInput = z.infer; + export const modelOptionSchema = z.object({ modelId: z.string(), name: z.string(), diff --git a/packages/workspace-server/src/services/local-logs/identifiers.ts b/packages/workspace-server/src/services/local-logs/identifiers.ts index 517b1bb4c6..e4d3ae9aaf 100644 --- a/packages/workspace-server/src/services/local-logs/identifiers.ts +++ b/packages/workspace-server/src/services/local-logs/identifiers.ts @@ -20,4 +20,9 @@ export interface ILogsService { maxBytes: number, ): Promise<{ content: string; truncated: boolean } | null>; writeLocalLogs(taskRunId: string, content: string): Promise; + seedLocalLogs(taskRunId: string, content: string): Promise; + cloneLocalLogs( + sourceTaskRunId: string, + targetTaskRunId: string, + ): Promise; } diff --git a/packages/workspace-server/src/services/local-logs/schemas.ts b/packages/workspace-server/src/services/local-logs/schemas.ts index dc6f36eace..67e7f68ba6 100644 --- a/packages/workspace-server/src/services/local-logs/schemas.ts +++ b/packages/workspace-server/src/services/local-logs/schemas.ts @@ -34,6 +34,11 @@ export const seedLocalLogsInput = z.object({ content: z.string(), }); +export const cloneLocalLogsInput = z.object({ + sourceTaskRunId: z.string().min(1), + targetTaskRunId: z.string().min(1), +}); + export const countLocalLogEntriesInput = z.object({ taskRunId: z.string().min(1), }); diff --git a/packages/workspace-server/src/services/local-logs/service.test.ts b/packages/workspace-server/src/services/local-logs/service.test.ts index c6d93d8493..31c7f64675 100644 --- a/packages/workspace-server/src/services/local-logs/service.test.ts +++ b/packages/workspace-server/src/services/local-logs/service.test.ts @@ -255,6 +255,20 @@ describe("LocalLogsService", () => { }); }); + describe("cloneLocalLogs", () => { + it("seeds the target from the source log", async () => { + const service = new LocalLogsService(); + vi.spyOn(service, "readLocalLogs").mockResolvedValue("source log\n"); + const seedLocalLogs = vi + .spyOn(service, "seedLocalLogs") + .mockResolvedValue(undefined); + + await service.cloneLocalLogs("source-run", "target-run"); + + expect(seedLocalLogs).toHaveBeenCalledWith("target-run", "source log\n"); + }); + }); + describe("countLocalLogEntries", () => { it("counts non-blank lines", async () => { mockReadFile.mockResolvedValue("a\n\nb\n c \n\n"); diff --git a/packages/workspace-server/src/services/local-logs/service.ts b/packages/workspace-server/src/services/local-logs/service.ts index b139d950a2..728d509f0a 100644 --- a/packages/workspace-server/src/services/local-logs/service.ts +++ b/packages/workspace-server/src/services/local-logs/service.ts @@ -231,6 +231,16 @@ export class LocalLogsService implements ILogsService { ); } + async cloneLocalLogs( + sourceTaskRunId: string, + targetTaskRunId: string, + ): Promise { + const content = await this.readLocalLogs(sourceTaskRunId); + if (content) { + await this.seedLocalLogs(targetTaskRunId, content); + } + } + async countLocalLogEntries(taskRunId: string): Promise { const logPath = this.getLocalLogPath(taskRunId); try { diff --git a/packages/workspace-server/src/services/workspace/schemas.ts b/packages/workspace-server/src/services/workspace/schemas.ts index 0de533573e..c488e3485f 100644 --- a/packages/workspace-server/src/services/workspace/schemas.ts +++ b/packages/workspace-server/src/services/workspace/schemas.ts @@ -29,6 +29,7 @@ export const createWorkspaceInput = z // When set, an existing worktree already checked out on the branch is reused // for the task instead of creating a new one. Gated behind a confirmation. reuseExistingWorktree: z.boolean().optional(), + forkFromTaskId: z.string().optional(), }) .refine( (data) => diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 11ff64591e..359bcf5d6a 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -67,6 +67,23 @@ const mockWorktreeManager = { createWorktreeForExistingBranch: vi.fn(), createWorktreeForRemoteBranch: vi.fn(), }; +const { + mockCaptureWorktreeCheckpoint, + mockDeleteWorktreeCheckpoint, + mockRestoreWorktreeFromCheckpoint, +} = vi.hoisted(() => ({ + mockCaptureWorktreeCheckpoint: vi.fn(), + mockDeleteWorktreeCheckpoint: vi.fn(), + mockRestoreWorktreeFromCheckpoint: vi.fn(), +})); +const mockGitClient = { + revparse: vi.fn(), + raw: vi.fn(), +}; + +vi.mock("@posthog/git/client", () => ({ + createGitClient: () => mockGitClient, +})); vi.mock("@posthog/git/worktree", () => ({ WorktreeManager: class { @@ -78,6 +95,12 @@ vi.mock("@posthog/git/worktree", () => ({ }, })); +vi.mock("../worktree-checkpoint/worktree-checkpoint", () => ({ + captureWorktreeCheckpoint: mockCaptureWorktreeCheckpoint, + deleteWorktreeCheckpoint: mockDeleteWorktreeCheckpoint, + restoreWorktreeFromCheckpoint: mockRestoreWorktreeFromCheckpoint, +})); + function createMocks() { const databaseService = { isInitialized: vi.fn(() => true), @@ -186,8 +209,13 @@ describe("WorkspaceService", () => { let service: WorkspaceService; beforeEach(() => { + vi.clearAllMocks(); mocks = createMocks(); service = makeService(mocks); + mockGitClient.revparse.mockResolvedValue("abc123\n"); + mockGitClient.raw.mockResolvedValue(""); + mockCaptureWorktreeCheckpoint.mockResolvedValue(undefined); + mockDeleteWorktreeCheckpoint.mockResolvedValue(undefined); }); describe("reconcileCloudWorkspaces", () => { @@ -738,6 +766,54 @@ describe("WorkspaceService", () => { }); }); + describe("createWorkspace (fork)", () => { + it("captures and restores the source workspace checkpoint", async () => { + seedWorktreeTask(mocks, { + taskId: "source-task", + repoPath: "/repo", + name: "source", + worktreePath: "/worktrees/source", + }); + mockRestoreWorktreeFromCheckpoint.mockResolvedValue({ + worktreePath: "/worktrees/child", + worktreeName: "child", + branchName: null, + baseBranch: "abc123", + createdAt: "2026-07-21T00:00:00Z", + }); + + await service.createWorkspace({ + taskId: "child-task", + mainRepoPath: "/repo", + folderId: "folder-1", + folderPath: "/repo", + mode: "worktree", + forkFromTaskId: "source-task", + }); + + expect(mockCaptureWorktreeCheckpoint).toHaveBeenCalledWith( + "/repo", + "/worktrees/source", + "fork-child-task", + null, + ); + expect(mockRestoreWorktreeFromCheckpoint).toHaveBeenCalledWith( + expect.objectContaining({ + mainRepoPath: "/repo", + worktreeBasePath: "/tmp/worktrees", + branchName: null, + checkpointId: "fork-child-task", + restoreCheckpointBranch: false, + }), + ); + expect(mockDeleteWorktreeCheckpoint).toHaveBeenCalledWith( + "/repo", + "fork-child-task", + ); + expect(mockWorktreeManager.createWorktree).not.toHaveBeenCalled(); + }); + }); + describe("pending creations", () => { afterEach(() => { vi.mocked(getCurrentBranch).mockReset(); diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index e366a53e5d..75c764b9cd 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -46,6 +46,11 @@ import type { ProcessTrackingService } from "../process-tracking/process-trackin import { getBranchFromPath, hasAnyFiles } from "../repo-fs-query/repo-fs-query"; import { SUSPENSION_SERVICE } from "../suspension/identifiers"; import type { SuspensionService } from "../suspension/suspension"; +import { + captureWorktreeCheckpoint, + deleteWorktreeCheckpoint, + restoreWorktreeFromCheckpoint, +} from "../worktree-checkpoint/worktree-checkpoint"; import { deleteWorktree as deleteGitWorktree, listLinkedWorktrees, @@ -636,6 +641,7 @@ export class WorkspaceService extends TypedEventEmitter useExistingBranch, allowRemoteBranchCheckout, reuseExistingWorktree, + forkFromTaskId, } = options; const existingWorkspace = await this.getWorkspaceInfo(taskId); @@ -748,7 +754,50 @@ export class WorkspaceService extends TypedEventEmitter // re-check here so a lost race (the worktree got claimed between preflight // and now) fails the saga step rather than sharing one worktree across two // tasks. - if (existingWorktree) { + if (forkFromTaskId) { + const sourceAssociation = this.findTaskAssociation(forkFromTaskId); + if (!sourceAssociation || sourceAssociation.mode === "cloud") { + throw new Error("The source task does not have a local workspace"); + } + const sourcePath = + sourceAssociation.mode === "worktree" + ? sourceAssociation.path + : sourceAssociation.folderId + ? this.getFolderPath(sourceAssociation.folderId) + : null; + if (!sourcePath) { + throw new Error("The source task workspace is unavailable"); + } + const checkpointId = `fork-${taskId}`; + await captureWorktreeCheckpoint( + mainRepoPath, + sourcePath, + checkpointId, + null, + ); + try { + worktree = await restoreWorktreeFromCheckpoint({ + mainRepoPath, + worktreeBasePath, + preferredName: undefined, + branchName: null, + checkpointId, + restoreCheckpointBranch: false, + logger: this.log, + onOutput, + }); + } finally { + await deleteWorktreeCheckpoint(mainRepoPath, checkpointId).catch( + (error) => { + this.log.warn("Failed to delete fork checkpoint", { + taskId, + checkpointId, + error, + }); + }, + ); + } + } else if (existingWorktree) { const [occupant] = this.getWorktreeTasks(existingWorktree.worktreePath); if (occupant) { throw new Error( diff --git a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.test.ts b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.test.ts index 25a117290a..106921a8a7 100644 --- a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.test.ts +++ b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.test.ts @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; const mockManager = vi.hoisted(() => ({ createWorktreeForExistingBranch: vi.fn(), createDetachedWorktreeAtCommit: vi.fn(), + deleteWorktree: vi.fn(), })); const mockRevertRun = vi.hoisted(() => vi.fn()); const mockCaptureRun = vi.hoisted(() => vi.fn()); @@ -14,6 +15,7 @@ vi.mock("@posthog/git/worktree", () => ({ createWorktreeForExistingBranch = mockManager.createWorktreeForExistingBranch; createDetachedWorktreeAtCommit = mockManager.createDetachedWorktreeAtCommit; + deleteWorktree = mockManager.deleteWorktree; }, })); @@ -35,6 +37,7 @@ vi.mock("@posthog/git/client", () => ({ import { captureWorktreeCheckpoint, + deleteWorktreeCheckpoint, restoreWorktreeFromCheckpoint, } from "./worktree-checkpoint"; @@ -52,6 +55,7 @@ const baseParams = { beforeEach(() => { mockManager.createWorktreeForExistingBranch.mockResolvedValue(BRANCH_WT); mockManager.createDetachedWorktreeAtCommit.mockResolvedValue(DETACHED_WT); + mockManager.deleteWorktree.mockResolvedValue(undefined); mockRevertRun.mockResolvedValue({ success: true }); mockCaptureRun.mockResolvedValue({ success: true }); mockDeleteCheckpoint.mockResolvedValue(undefined); @@ -92,6 +96,20 @@ describe("restoreWorktreeFromCheckpoint", () => { expect(mockRevertRun).toHaveBeenCalledWith({ baseDir: "/wt/branch", checkpointId: "cp-1", + restoreCheckpointBranch: undefined, + }); + }); + + it("can restore the checkpoint without claiming its recorded branch", async () => { + await restoreWorktreeFromCheckpoint({ + ...baseParams, + restoreCheckpointBranch: false, + }); + + expect(mockRevertRun).toHaveBeenCalledWith({ + baseDir: "/wt/branch", + checkpointId: "cp-1", + restoreCheckpointBranch: false, }); }); @@ -101,6 +119,7 @@ describe("restoreWorktreeFromCheckpoint", () => { await expect(restoreWorktreeFromCheckpoint(baseParams)).rejects.toThrow( /failed to apply checkpoint: bad patch/, ); + expect(mockManager.deleteWorktree).toHaveBeenCalledWith("/wt/branch"); }); it("recreates the branch after revert when recreateBranch is set", async () => { @@ -150,3 +169,14 @@ describe("captureWorktreeCheckpoint", () => { ).rejects.toThrow(/Failed to capture checkpoint: dirty index/); }); }); + +describe("deleteWorktreeCheckpoint", () => { + it("deletes the checkpoint from the repository", async () => { + await deleteWorktreeCheckpoint("/repo", "cp-1"); + + expect(mockDeleteCheckpoint).toHaveBeenCalledWith( + expect.anything(), + "cp-1", + ); + }); +}); diff --git a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts index 4a1ce61e43..f92457ec02 100644 --- a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts +++ b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts @@ -15,7 +15,9 @@ export interface RestoreWorktreeFromCheckpointParams { branchName: string | null; checkpointId: string; recreateBranch?: boolean; + restoreCheckpointBranch?: boolean; logger?: SagaLogger; + onOutput?: (data: string) => void; } /** @@ -34,31 +36,62 @@ export async function restoreWorktreeFromCheckpoint( let newWorktree: WorktreeInfo; if (params.branchName && !params.recreateBranch) { - newWorktree = await manager.createWorktreeForExistingBranch( - params.branchName, - params.preferredName, - ); + newWorktree = params.onOutput + ? await manager.createWorktreeForExistingBranch( + params.branchName, + params.preferredName, + { onOutput: params.onOutput }, + ) + : await manager.createWorktreeForExistingBranch( + params.branchName, + params.preferredName, + ); } else { - newWorktree = await manager.createDetachedWorktreeAtCommit( - "HEAD", - params.preferredName, - ); + newWorktree = params.onOutput + ? await manager.createDetachedWorktreeAtCommit( + "HEAD", + params.preferredName, + { onOutput: params.onOutput }, + ) + : await manager.createDetachedWorktreeAtCommit( + "HEAD", + params.preferredName, + ); } - const revertSaga = new RevertCheckpointSaga(); - const result = await revertSaga.run({ - baseDir: newWorktree.worktreePath, - checkpointId: params.checkpointId, - }); - if (!result.success) { - throw new Error( - `Worktree restored but failed to apply checkpoint: ${result.error}`, - ); - } + try { + const revertSaga = new RevertCheckpointSaga(); + const result = await revertSaga.run({ + baseDir: newWorktree.worktreePath, + checkpointId: params.checkpointId, + restoreCheckpointBranch: params.restoreCheckpointBranch, + }); + if (!result.success) { + throw new Error( + `Worktree restored but failed to apply checkpoint: ${result.error}`, + ); + } - if (params.recreateBranch && params.branchName) { - const git = createGitClient(newWorktree.worktreePath); - await git.checkoutLocalBranch(params.branchName); + if (params.recreateBranch && params.branchName) { + const git = createGitClient(newWorktree.worktreePath); + await git.checkoutLocalBranch(params.branchName); + } + } catch (error) { + await manager + .deleteWorktree(newWorktree.worktreePath) + .catch((cleanupError) => { + params.logger?.warn( + "Failed to clean up worktree after restore failure", + { + worktreePath: newWorktree.worktreePath, + error: + cleanupError instanceof Error + ? cleanupError.message + : String(cleanupError), + }, + ); + }); + throw error; } return newWorktree; @@ -73,6 +106,7 @@ export async function captureWorktreeCheckpoint( folderPath: string, worktreePath: string, checkpointId: string, + maxWorktreeFileBytes?: number | null, ): Promise { const git = createGitClient(folderPath); try { @@ -80,8 +114,19 @@ export async function captureWorktreeCheckpoint( } catch {} const saga = new CaptureCheckpointSaga(); - const result = await saga.run({ baseDir: worktreePath, checkpointId }); + const result = await saga.run({ + baseDir: worktreePath, + checkpointId, + ...(maxWorktreeFileBytes !== undefined && { maxWorktreeFileBytes }), + }); if (!result.success) { throw new Error(`Failed to capture checkpoint: ${result.error}`); } } + +export async function deleteWorktreeCheckpoint( + folderPath: string, + checkpointId: string, +): Promise { + await deleteCheckpoint(createGitClient(folderPath), checkpointId); +}