From 9701fc16e9f8235a753bb42053a04cb293caa07f Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 12:15:15 +0300 Subject: [PATCH 01/10] Add task forking with preserved context Create independent local and cloud tasks from an existing run while preserving conversation and workspace state. Closes #330 Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- .../api-client/src/posthog-client.test.ts | 2 + packages/api-client/src/posthog-client.ts | 1 + packages/core/src/sessions/sessionService.ts | 109 ++++++++++++ packages/core/src/task-detail/identifiers.ts | 3 + .../src/task-detail/task-detail.module.ts | 8 +- .../src/task-detail/taskCreationApiClient.ts | 6 + .../core/src/task-detail/taskCreationHost.ts | 1 + .../src/task-detail/taskCreationSaga.test.ts | 89 ++++++++++ .../core/src/task-detail/taskCreationSaga.ts | 34 +++- .../src/task-detail/taskForkService.test.ts | 162 ++++++++++++++++++ .../core/src/task-detail/taskForkService.ts | 97 +++++++++++ .../host-router/src/routers/agent.router.ts | 16 ++ .../host-router/src/routers/logs.router.ts | 9 + packages/shared/src/task-creation-domain.ts | 9 + .../components/ForkTaskButton.test.tsx | 89 ++++++++++ .../task-detail/components/ForkTaskButton.tsx | 92 ++++++++++ .../components/TaskHeaderActions.test.tsx | 4 + .../components/TaskHeaderActions.tsx | 3 + .../src/services/agent/agent.test.ts | 35 ++++ .../src/services/agent/agent.ts | 82 ++++++++- .../src/services/agent/codex-home.ts | 24 +++ .../src/services/agent/schemas.ts | 6 + .../src/services/local-logs/identifiers.ts | 1 + .../src/services/workspace/schemas.ts | 1 + .../src/services/workspace/workspace.test.ts | 116 +++++++++++++ .../src/services/workspace/workspace.ts | 77 ++++++++- 26 files changed, 1069 insertions(+), 7 deletions(-) create mode 100644 packages/core/src/task-detail/taskForkService.test.ts create mode 100644 packages/core/src/task-detail/taskForkService.ts create mode 100644 packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx create mode 100644 packages/ui/src/features/task-detail/components/ForkTaskButton.tsx 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/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 3846d2164f..8c9da2e8b4 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,7 @@ export interface SessionTrpc { readLocalLogsTail?: TrpcQuery; fetchS3Logs: TrpcQuery; writeLocalLogs: TrpcMutation; + seedLocalLogs: TrpcMutation; }; os: { openExternal: TrpcMutation }; } @@ -1503,6 +1506,112 @@ export class SessionService { } } + async forkLocalTask(params: { + sourceTaskId: string; + sourceTaskRunId: string; + task: Task; + repoPath: string; + }): Promise { + const source = this.d.store.getSessionByTaskId(params.sourceTaskId); + if (!source || source.taskRunId !== params.sourceTaskRunId) { + 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; + + await this.d.trpc.agent.flushLogs.mutate({ + taskRunId: params.sourceTaskRunId, + }); + const sourceLogs = await this.d.trpc.logs.readLocalLogs.query({ + taskRunId: params.sourceTaskRunId, + }); + if (sourceLogs) { + await this.d.trpc.logs.seedLocalLogs.mutate({ + taskRunId: taskRun.id, + content: sourceLogs, + }); + } + + const { customInstructions, rtkEnabledLocal, spokenNarrationEnabled } = + this.d.settings; + const result = await this.d.trpc.agent.fork.mutate({ + sourceTaskRunId: params.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 session = createBaseSession( + taskRun.id, + params.task.id, + params.task.title, + ); + session.channel = result.channel; + session.status = "connected"; + session.events = [...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); + + 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: params.sourceTaskRunId, + }, + }, + ); + + this.d.track(ANALYTICS_EVENTS.TASK_RUN_STARTED, { + task_id: params.task.id, + execution_type: "local", + initial_mode: source.executionMode, + adapter: source.adapter, + }); + } + async loadLogsOnly(params: { taskId: string; taskRunId: string; 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..86435725f5 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; diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 5558ab6580..f0e29e6006 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,93 @@ describe("TaskCreationSaga", () => { ); }); + it("forks a cloud run from the backend snapshot without sending a prompt", async () => { + const createdTask = createTask(); + 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: vi.fn().mockResolvedValue(createdTask), + createTaskRun: createTaskRunMock, + startTaskRun: startTaskRunMock, + updateTaskRun: updateTaskRunMock, + }); + + const result = await saga.run({ + content: "Ship the fix", + repository: "posthog/posthog", + workspaceMode: "cloud", + adapter: "codex", + forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + }); + + expect(result.success).toBe(true); + 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: { taskId: "source-task", taskRunId: "source-run" }, + }); + + expect(result.success).toBe(true); + expect(mockHost.createWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ forkFromTaskId: "source-task" }), + ); + expect(sessionService.forkLocalTask).toHaveBeenCalledWith({ + sourceTaskId: "source-task", + sourceTaskRunId: "source-run", + 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() }); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index f6da0ae29d..c6847947a5 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 () => { @@ -357,6 +358,7 @@ export class TaskCreationSaga extends Saga< const buildTransport = async (): Promise => { if ( + input.forkFrom || !(input.content || input.filePaths?.length) || workspaceMode !== "cloud" ) { @@ -411,6 +413,9 @@ export class TaskCreationSaga extends Saga< homeQuickAction: input.homeQuickActionLabel, importedMcpServers: input.importedMcpServers, relayedMcpServers: input.relayedMcpServers, + ...(input.forkFrom && { + resumeFromRunId: input.forkFrom.taskRunId, + }), initialPermissionMode: input.executionMode ?? (cloudAdapter === "codex" ? "auto" : "plan"), @@ -452,6 +457,20 @@ export class TaskCreationSaga extends Saga< }, ); + if (input.forkFrom && 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 +505,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 +546,16 @@ export class TaskCreationSaga extends Saga< connectParams.adapter = "claude"; } - this.deps.sessionService.connectToTask(connectParams); + if (input.forkFrom) { + await this.deps.sessionService.forkLocalTask({ + sourceTaskId: input.forkFrom.taskId, + sourceTaskRunId: input.forkFrom.taskRunId, + task, + repoPath: agentCwd ?? "", + }); + } else { + this.deps.sessionService.connectToTask(connectParams); + } return { taskId: task.id }; }, rollback: async ({ taskId }) => { 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..95aa956eb1 --- /dev/null +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -0,0 +1,162 @@ +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 { 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 host = { + getWorkspace: vi.fn(), + } as unknown as ITaskCreationHost; + const taskService = { + createTask: vi.fn(), + } as unknown as TaskService; + const service = new TaskForkService(host, taskService); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(host.getWorkspace).mockResolvedValue(createWorkspace()); + vi.mocked(taskService.createTask).mockResolvedValue({ + success: true, + data: {} as never, + }); + }); + + it("rejects an active cloud run", async () => { + 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 cloud run to finish before forking it", + failedStep: "validation", + }); + }); + + it("creates a local worktree fork", async () => { + const task = createTask({ + latest_run: createRun({ + runtime_adapter: "codex", + model: "gpt-5.4", + reasoning_effort: "high", + state: { initial_permission_mode: "full-access" }, + }), + }); + + 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, + adapter: "codex", + model: "gpt-5.4", + reasoningLevel: "high", + executionMode: "full-access", + forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + }), + undefined, + ); + }); + + 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: { taskId: "source-task", taskRunId: "source-run" }, + }), + 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..b378ab9e0e --- /dev/null +++ b/packages/core/src/task-detail/taskForkService.ts @@ -0,0 +1,97 @@ +import type { TaskCreationOutput } from "@posthog/shared"; +import { + executionModeSchema, + isTerminalStatus, + type Task, +} from "@posthog/shared/domain-types"; +import { inject, injectable } from "inversify"; +import { TASK_CREATION_HOST, TASK_SERVICE } from "./identifiers"; +import type { ITaskCreationHost } from "./taskCreationHost"; +import type { CreateTaskResult, TaskService } from "./taskService"; + +@injectable() +export class TaskForkService { + constructor( + @inject(TASK_CREATION_HOST) + private readonly host: ITaskCreationHost, + @inject(TASK_SERVICE) + private readonly taskService: TaskService, + ) {} + + async forkTask( + sourceTask: Task, + onTaskReady?: (output: TaskCreationOutput) => void, + ): Promise { + const sourceRun = sourceTask.latest_run; + if (!sourceRun) { + return this.validationError("The source task has no run to fork"); + } + if (sourceTask.runtime === "pi") { + return this.validationError("Pi tasks cannot be forked yet"); + } + + const sourceWorkspace = await this.host.getWorkspace(sourceTask.id); + const isCloud = + sourceWorkspace?.mode === "cloud" || sourceRun.environment === "cloud"; + if (isCloud && !isTerminalStatus(sourceRun.status)) { + return this.validationError( + "Wait for the cloud run to finish before forking it", + ); + } + if (!isCloud && (!sourceWorkspace || sourceWorkspace.isScratch)) { + return this.validationError( + "Only repository-backed local tasks can be forked", + ); + } + + 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: isCloud ? undefined : sourceWorkspace?.folderPath, + repository: sourceTask.repository, + workspaceMode: isCloud ? "cloud" : "worktree", + branch: isCloud ? cloudBranch : null, + 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: + state.run_source === "signal_report" ? "signal_report" : "manual", + cloudPrAuthorshipMode: + state.pr_authorship_mode === "bot" ? "bot" : "user", + forkFrom: { + taskId: sourceTask.id, + taskRunId: sourceRun.id, + }, + }, + onTaskReady, + ); + } + + private validationError(error: string): CreateTaskResult { + return { success: false, error, failedStep: "validation" }; + } +} 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..1a4bab5324 100644 --- a/packages/host-router/src/routers/logs.router.ts +++ b/packages/host-router/src/routers/logs.router.ts @@ -10,6 +10,7 @@ import { readLocalLogsOutput, readLocalLogsTailInput, readLocalLogsTailOutput, + seedLocalLogsInput, writeLocalLogsInput, } from "@posthog/workspace-server/services/local-logs/schemas"; @@ -55,4 +56,12 @@ 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), + ), }); diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 0c354e146a..785aca706c 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -109,6 +109,15 @@ 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 the source run's current conversation and + * workspace state. Local runs use the native agent fork API; cloud runs use + * the backend's persisted run snapshot. + */ + forkFrom?: { + taskId: string; + taskRunId: string; + }; } export interface TaskCreationOutput { diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx new file mode 100644 index 0000000000..5a054f5925 --- /dev/null +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -0,0 +1,89 @@ +import type { Task } from "@posthog/shared/domain-types"; +import { Theme } from "@radix-ui/themes"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + forkTask: vi.fn(), + navigateToTaskDetail: vi.fn(), + useWorkspace: vi.fn(), + useSessionForTask: vi.fn(), +})); + +vi.mock("@posthog/di/react", () => ({ + useService: () => ({ forkTask: mocks.forkTask }), +})); +vi.mock("@posthog/ui/features/workspace/useWorkspace", () => ({ + useWorkspace: mocks.useWorkspace, +})); +vi.mock("@posthog/ui/features/sessions/useSession", () => ({ + useSessionForTask: mocks.useSessionForTask, +})); +vi.mock("@posthog/ui/router/navigationBridge", () => ({ + navigateToTaskDetail: mocks.navigateToTaskDetail, +})); +vi.mock("@posthog/ui/primitives/Tooltip", () => ({ + Tooltip: ({ children }: { children: React.ReactNode }) => children, +})); +vi.mock("@posthog/ui/primitives/toast", () => ({ + toast: { error: vi.fn() }, +})); + +import { ForkTaskButton, ForkedFromTaskButton } from "./ForkTaskButton"; + +const cloudTask = { + id: "task-1", + title: "Source task", + latest_run: { + id: "run-1", + environment: "cloud", + status: "completed", + state: {}, + }, +} as Task; + +describe("ForkTaskButton", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.useWorkspace.mockReturnValue({ mode: "cloud" }); + mocks.useSessionForTask.mockReturnValue(null); + }); + + it("forks a completed cloud task and opens the child", async () => { + mocks.forkTask.mockImplementation( + async (_task: Task, onReady: (output: { task: Task }) => void) => { + onReady({ task: { id: "task-2" } as Task }); + return { success: true, data: {} }; + }, + ); + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => expect(mocks.forkTask).toHaveBeenCalledWith(cloudTask, expect.any(Function))); + expect(mocks.navigateToTaskDetail).toHaveBeenCalledWith("task-2"); + }); + + it("links a forked task back to its source", () => { + const child = { + ...cloudTask, + latest_run: { + ...cloudTask.latest_run, + state: { forked_from_task_id: "task-parent" }, + }, + } as Task; + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Open source task" })); + + expect(mocks.navigateToTaskDetail).toHaveBeenCalledWith("task-parent"); + }); +}); diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx new file mode 100644 index 0000000000..e48b12390f --- /dev/null +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -0,0 +1,92 @@ +import { GitBranch, Spinner } from "@phosphor-icons/react"; +import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; +import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; +import { useService } from "@posthog/di/react"; +import { Button } from "@posthog/quill"; +import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; +import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; +import { Tooltip } from "@posthog/ui/primitives/Tooltip"; +import { toast } from "@posthog/ui/primitives/toast"; +import { navigateToTaskDetail } from "@posthog/ui/router/navigationBridge"; +import { useState } from "react"; + +export function ForkTaskButton({ task }: { task: Task }) { + const taskForkService = useService(TASK_FORK_SERVICE); + const workspace = useWorkspace(task.id); + const session = useSessionForTask(task.id); + const [isForking, setIsForking] = useState(false); + const run = task.latest_run; + const isCloud = workspace?.mode === "cloud" || run?.environment === "cloud"; + const canForkCloud = !!run && isTerminalStatus(run.status); + const canForkLocal = + !!run && + !!workspace && + !workspace.isScratch && + session?.status === "connected" && + !session.isPromptPending; + const canFork = isCloud ? canForkCloud : canForkLocal; + const tooltip = canFork + ? "Fork task" + : isCloud + ? "Wait for the cloud run to finish" + : "Wait for the local task to finish"; + + const handleFork = async () => { + setIsForking(true); + try { + const result = await taskForkService.forkTask(task, ({ task: child }) => { + navigateToTaskDetail(child.id); + }); + if (!result.success) { + toast.error("Could not fork task", { description: result.error }); + } + } catch (error) { + toast.error("Could not fork task", { + description: error instanceof Error ? error.message : String(error), + }); + } finally { + setIsForking(false); + } + }; + + return ( + + + + + + ); +} + +export function ForkedFromTaskButton({ task }: { task: Task }) { + const parentTaskId = task.latest_run?.state?.forked_from_task_id; + if (typeof parentTaskId !== "string") return null; + + return ( + + + + + + ); +} diff --git a/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx b/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx index f80d03734d..464d6d0665 100644 --- a/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx +++ b/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx @@ -79,6 +79,10 @@ vi.mock("@posthog/ui/features/diff-stats/DiffStatsBadge", () => ({ vi.mock("@posthog/ui/primitives/Tooltip", () => ({ Tooltip: ({ children }: { children: ReactNode }) => children, })); +vi.mock("./ForkTaskButton", () => ({ + ForkTaskButton: () => null, + ForkedFromTaskButton: () => null, +})); import { TaskHeaderActions } from "./TaskHeaderActions"; diff --git a/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx b/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx index d69a198bd9..c3966ac717 100644 --- a/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx +++ b/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx @@ -26,6 +26,7 @@ import { import { Tooltip } from "@posthog/ui/primitives/Tooltip"; import { Flex } from "@radix-ui/themes"; import { useState } from "react"; +import { ForkedFromTaskButton, ForkTaskButton } from "./ForkTaskButton"; const CLOUD_HANDOFF_FLAG = "phc-cloud-handoff"; @@ -158,6 +159,8 @@ export function TaskHeaderActions({ task }: { task: Task }) { )} + + {workspaceLoaded && ( <> diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index 7d479b71df..cf5dceaa4b 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,33 @@ 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("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..5f15ad771b 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,43 @@ If a repository IS genuinely required, attach one in this priority order: let configOptions: SessionConfigOption[] | undefined; let agentSessionId: string | undefined; + if (forkSource) { + const forkResponse = await connection.unstable_forkSession({ + sessionId: forkSource.sessionId, + 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({ @@ -1392,6 +1462,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.ts b/packages/workspace-server/src/services/agent/codex-home.ts index 520281520c..85615fee23 100644 --- a/packages/workspace-server/src/services/agent/codex-home.ts +++ b/packages/workspace-server/src/services/agent/codex-home.ts @@ -94,3 +94,27 @@ 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); + if (!fs.existsSync(source)) continue; + await fs.promises.cp(source, path.join(targetHome, directory), { + recursive: true, + force: true, + }); + } +} 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..e114b6b952 100644 --- a/packages/workspace-server/src/services/local-logs/identifiers.ts +++ b/packages/workspace-server/src/services/local-logs/identifiers.ts @@ -20,4 +20,5 @@ export interface ILogsService { maxBytes: number, ): Promise<{ content: string; truncated: boolean } | null>; writeLocalLogs(taskRunId: string, content: string): Promise; + seedLocalLogs(taskRunId: string, content: string): Promise; } 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..08331f5f16 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -67,6 +67,22 @@ const mockWorktreeManager = { createWorktreeForExistingBranch: vi.fn(), createWorktreeForRemoteBranch: vi.fn(), }; +const mockGitClient = { + revparse: vi.fn(), + diff: vi.fn(), + raw: vi.fn(), +}; +const mockApplyPatch = vi.fn(); + +vi.mock("@posthog/git/client", () => ({ + createGitClient: () => mockGitClient, +})); + +vi.mock("@posthog/git/sagas/patch", () => ({ + ApplyPatchSaga: class { + run = mockApplyPatch; + }, +})); vi.mock("@posthog/git/worktree", () => ({ WorktreeManager: class { @@ -186,8 +202,13 @@ describe("WorkspaceService", () => { let service: WorkspaceService; beforeEach(() => { + vi.clearAllMocks(); mocks = createMocks(); service = makeService(mocks); + mockGitClient.revparse.mockResolvedValue("abc123\n"); + mockGitClient.diff.mockResolvedValue(""); + mockGitClient.raw.mockResolvedValue(""); + mockApplyPatch.mockResolvedValue({ success: true, data: undefined }); }); describe("reconcileCloudWorkspaces", () => { @@ -738,6 +759,101 @@ describe("WorkspaceService", () => { }); }); + describe("createWorkspace (fork)", () => { + const tempDirs: string[] = []; + + afterEach(() => { + for (const dir of tempDirs.splice(0)) { + fs.rmSync(dir, { recursive: true, force: true }); + } + }); + + function mkTemp(prefix: string): string { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix)); + tempDirs.push(dir); + return dir; + } + + it("creates the child worktree at the source HEAD and applies its changes", async () => { + seedWorktreeTask(mocks, { + taskId: "source-task", + repoPath: "/repo", + name: "source", + worktreePath: "/worktrees/source", + }); + mockWorktreeManager.createWorktree.mockResolvedValue({ + worktreePath: "/worktrees/child", + worktreeName: "child", + branchName: null, + baseBranch: "abc123", + createdAt: "2026-07-21T00:00:00Z", + }); + mockGitClient.diff + .mockResolvedValueOnce("staged patch") + .mockResolvedValueOnce("unstaged patch"); + + await service.createWorkspace({ + taskId: "child-task", + mainRepoPath: "/repo", + folderId: "folder-1", + folderPath: "/repo", + mode: "worktree", + forkFromTaskId: "source-task", + }); + + expect(mockWorktreeManager.createWorktree).toHaveBeenCalledWith( + expect.objectContaining({ baseBranch: "abc123" }), + ); + expect(mockApplyPatch).toHaveBeenNthCalledWith(1, { + baseDir: "/worktrees/child", + patch: "staged patch", + cached: true, + }); + expect(mockApplyPatch).toHaveBeenNthCalledWith(2, { + baseDir: "/worktrees/child", + patch: "unstaged patch", + cached: false, + }); + }); + + it("preserves untracked symlinks", async () => { + const sourcePath = mkTemp("posthog-fork-source-"); + const targetPath = mkTemp("posthog-fork-target-"); + fs.writeFileSync(path.join(sourcePath, "target.txt"), "target"); + fs.symlinkSync("target.txt", path.join(sourcePath, "link.txt")); + seedWorktreeTask(mocks, { + taskId: "source-task", + repoPath: "/repo", + name: "source", + worktreePath: sourcePath, + }); + mockWorktreeManager.createWorktree.mockResolvedValue({ + worktreePath: targetPath, + worktreeName: "child", + branchName: null, + baseBranch: "abc123", + createdAt: "2026-07-21T00:00:00Z", + }); + mockGitClient.raw.mockResolvedValue("link.txt\0"); + + await service.createWorkspace({ + taskId: "child-task", + mainRepoPath: "/repo", + folderId: "folder-1", + folderPath: "/repo", + mode: "worktree", + forkFromTaskId: "source-task", + }); + + expect( + fs.lstatSync(path.join(targetPath, "link.txt")).isSymbolicLink(), + ).toBe(true); + expect(fs.readlinkSync(path.join(targetPath, "link.txt"))).toBe( + "target.txt", + ); + }); + }); + 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..779338120a 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -16,6 +16,7 @@ import { } from "@posthog/git/queries"; import { CreateOrSwitchBranchSaga } from "@posthog/git/sagas/branch"; import { DetachHeadSaga } from "@posthog/git/sagas/head"; +import { ApplyPatchSaga } from "@posthog/git/sagas/patch"; import { WorktreeManager } from "@posthog/git/worktree"; import { ANALYTICS_SERVICE, @@ -636,6 +637,7 @@ export class WorkspaceService extends TypedEventEmitter useExistingBranch, allowRemoteBranchCheckout, reuseExistingWorktree, + forkFromTaskId, } = options; const existingWorkspace = await this.getWorkspaceInfo(taskId); @@ -748,7 +750,27 @@ 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 sourceHead = await createGitClient(sourcePath).revparse(["HEAD"]); + worktree = await worktreeManager.createWorktree({ + baseBranch: sourceHead.trim(), + onOutput, + }); + await this.copyWorkspaceChanges(sourcePath, worktree.worktreePath); + } else if (existingWorktree) { const [occupant] = this.getWorktreeTasks(existingWorktree.worktreePath); if (occupant) { throw new Error( @@ -871,6 +893,59 @@ export class WorkspaceService extends TypedEventEmitter }; } + private async copyWorkspaceChanges( + sourcePath: string, + targetPath: string, + ): Promise { + const sourceGit = createGitClient(sourcePath); + const [stagedPatch, unstagedPatch, untrackedFiles] = await Promise.all([ + sourceGit.diff(["--binary", "--cached", "HEAD"]), + sourceGit.diff(["--binary"]), + sourceGit.raw(["ls-files", "--others", "--exclude-standard", "-z"]), + ]); + + for (const [patch, cached] of [ + [stagedPatch, true], + [unstagedPatch, false], + ] as const) { + if (!patch) continue; + const result = await new ApplyPatchSaga().run({ + baseDir: targetPath, + patch, + cached, + }); + if (!result.success) { + throw new Error( + `Failed to copy forked workspace changes: ${result.error}`, + ); + } + } + + await Promise.all( + untrackedFiles + .split("\0") + .filter(Boolean) + .map(async (filePath) => { + const sourceFile = path.join(sourcePath, filePath); + const targetFile = path.join(targetPath, filePath); + await fs.promises.mkdir(path.dirname(targetFile), { + recursive: true, + }); + const stat = await fs.promises.lstat(sourceFile); + if (stat.isSymbolicLink()) { + await fs.promises.symlink( + await fs.promises.readlink(sourceFile), + targetFile, + ); + } else if (stat.isFile()) { + await fs.promises.copyFile(sourceFile, targetFile); + } else { + throw new Error(`Cannot copy untracked path: ${filePath}`); + } + }), + ); + } + async deleteWorkspace(taskId: string, mainRepoPath: string): Promise { this.log.info(`Deleting workspace for task ${taskId}`); From be049769d5f68b7027709a0b60911c790e18dd6f Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 12:49:20 +0300 Subject: [PATCH 02/10] Fix task fork state preservation Preserve fork semantics across retries and partial failures, reuse checkpoint snapshots for local workspaces, and align fork UI state with live sessions. Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- packages/core/src/sessions/sessionService.ts | 36 ++-- .../sessionServiceForkLocalTask.test.ts | 183 ++++++++++++++++++ .../src/task-detail/taskForkService.test.ts | 22 +++ .../core/src/task-detail/taskForkService.ts | 10 +- .../host-router/src/routers/logs.router.ts | 9 + .../components/ForkTaskButton.test.tsx | 128 +++++++++++- .../task-detail/components/ForkTaskButton.tsx | 51 +++-- .../src/services/agent/agent.test.ts | 28 +++ .../src/services/agent/agent.ts | 4 +- .../src/services/agent/codex-home.test.ts | 41 ++++ .../src/services/local-logs/identifiers.ts | 4 + .../src/services/local-logs/schemas.ts | 5 + .../src/services/local-logs/service.test.ts | 14 ++ .../src/services/local-logs/service.ts | 10 + .../src/services/workspace/workspace.test.ts | 112 ++++------- .../src/services/workspace/workspace.ts | 88 +++------ .../worktree-checkpoint.test.ts | 16 ++ .../worktree-checkpoint.ts | 80 ++++++-- 18 files changed, 634 insertions(+), 207 deletions(-) create mode 100644 packages/core/src/sessions/sessionServiceForkLocalTask.test.ts diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 8c9da2e8b4..7aeebcfbdf 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -200,6 +200,7 @@ export interface SessionTrpc { fetchS3Logs: TrpcQuery; writeLocalLogs: TrpcMutation; seedLocalLogs: TrpcMutation; + cloneLocalLogs: TrpcMutation; }; os: { openExternal: TrpcMutation }; } @@ -1534,15 +1535,10 @@ export class SessionService { await this.d.trpc.agent.flushLogs.mutate({ taskRunId: params.sourceTaskRunId, }); - const sourceLogs = await this.d.trpc.logs.readLocalLogs.query({ - taskRunId: params.sourceTaskRunId, + await this.d.trpc.logs.cloneLocalLogs.mutate({ + sourceTaskRunId: params.sourceTaskRunId, + targetTaskRunId: taskRun.id, }); - if (sourceLogs) { - await this.d.trpc.logs.seedLocalLogs.mutate({ - taskRunId: taskRun.id, - content: sourceLogs, - }); - } const { customInstructions, rtkEnabledLocal, spokenNarrationEnabled } = this.d.settings; @@ -1564,6 +1560,18 @@ export class SessionService { model: source.model ?? this.d.DEFAULT_GATEWAY_MODEL, }); + 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: params.sourceTaskRunId, + }, + }, + ); + const session = createBaseSession( taskRun.id, params.task.id, @@ -1592,18 +1600,6 @@ export class SessionService { this.d.store.setSession(session); this.subscribeToChannel(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: params.sourceTaskRunId, - }, - }, - ); - this.d.track(ANALYTICS_EVENTS.TASK_RUN_STARTED, { task_id: params.task.id, execution_type: "local", diff --git a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts new file mode 100644 index 0000000000..6f8e21ecf5 --- /dev/null +++ b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts @@ -0,0 +1,183 @@ +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 deps = { + store: { + getSessionByTaskId: vi.fn(() => sourceSession), + setSession, + }, + trpc: { + agent: { + flushLogs: { mutate: vi.fn().mockResolvedValue(undefined) }, + fork: { + mutate: vi.fn().mockResolvedValue({ + channel: "child-channel", + configOptions: [{ id: "model", type: "select" }], + }), + }, + onSessionEvent: { subscribe }, + onPermissionRequest: { subscribe }, + onSessionIdleKilled: { subscribe }, + }, + logs: { + cloneLocalLogs: { mutate: vi.fn().mockResolvedValue(undefined) }, + }, + }, + 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, + }; +} + +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 } = + createHarness(updateTaskRun); + + await service.forkLocalTask({ + sourceTaskId: "source-task", + sourceTaskRunId: "source-run", + 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(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, + } = createHarness(updateTaskRun); + + await expect( + service.forkLocalTask({ + sourceTaskId: "source-task", + sourceTaskRunId: "source-run", + task: { ...childTask }, + repoPath: "/repos/code", + }), + ).rejects.toThrow("lineage failed"); + + expect(setPersistedConfigOptions).not.toHaveBeenCalled(); + expect(setAdapter).not.toHaveBeenCalled(); + expect(setSession).not.toHaveBeenCalled(); + expect(subscribe).toHaveBeenCalledTimes(1); + expect(track).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index 95aa956eb1..e12c2a6306 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -159,4 +159,26 @@ describe("TaskForkService", () => { 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 index b378ab9e0e..a8b1ec2689 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -5,6 +5,10 @@ import { type Task, } from "@posthog/shared/domain-types"; import { inject, injectable } from "inversify"; +import { + getCloudPrAuthorshipMode, + getCloudRunSource, +} from "../sessions/cloudRunOptions"; import { TASK_CREATION_HOST, TASK_SERVICE } from "./identifiers"; import type { ITaskCreationHost } from "./taskCreationHost"; import type { CreateTaskResult, TaskService } from "./taskService"; @@ -78,10 +82,8 @@ export class TaskForkService { : undefined, cloudAutoPublish: state.auto_publish === true, cloudRtkEnabled: state.rtk_enabled === false ? false : undefined, - cloudRunSource: - state.run_source === "signal_report" ? "signal_report" : "manual", - cloudPrAuthorshipMode: - state.pr_authorship_mode === "bot" ? "bot" : "user", + cloudRunSource: getCloudRunSource(state), + cloudPrAuthorshipMode: getCloudPrAuthorshipMode(state), forkFrom: { taskId: sourceTask.id, taskRunId: sourceRun.id, diff --git a/packages/host-router/src/routers/logs.router.ts b/packages/host-router/src/routers/logs.router.ts index 1a4bab5324..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, @@ -64,4 +65,12 @@ export const logsRouter = router({ .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/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 5a054f5925..8930a7f5eb 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -6,8 +6,9 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ forkTask: vi.fn(), navigateToTaskDetail: vi.fn(), + openTask: vi.fn(), + session: null as Record | null, useWorkspace: vi.fn(), - useSessionForTask: vi.fn(), })); vi.mock("@posthog/di/react", () => ({ @@ -17,19 +18,31 @@ vi.mock("@posthog/ui/features/workspace/useWorkspace", () => ({ useWorkspace: mocks.useWorkspace, })); vi.mock("@posthog/ui/features/sessions/useSession", () => ({ - useSessionForTask: mocks.useSessionForTask, + useSessionSelector: ( + _taskId: string, + select: (session: Record | undefined) => unknown, + ) => select(mocks.session ?? undefined), })); vi.mock("@posthog/ui/router/navigationBridge", () => ({ navigateToTaskDetail: mocks.navigateToTaskDetail, })); +vi.mock("@posthog/ui/router/useOpenTask", () => ({ + openTask: mocks.openTask, +})); vi.mock("@posthog/ui/primitives/Tooltip", () => ({ - Tooltip: ({ children }: { children: React.ReactNode }) => children, + Tooltip: ({ + children, + content, + }: { + children: React.ReactNode; + content: string; + }) =>
{children}
, })); vi.mock("@posthog/ui/primitives/toast", () => ({ toast: { error: vi.fn() }, })); -import { ForkTaskButton, ForkedFromTaskButton } from "./ForkTaskButton"; +import { ForkedFromTaskButton, ForkTaskButton } from "./ForkTaskButton"; const cloudTask = { id: "task-1", @@ -46,7 +59,7 @@ describe("ForkTaskButton", () => { beforeEach(() => { vi.clearAllMocks(); mocks.useWorkspace.mockReturnValue({ mode: "cloud" }); - mocks.useSessionForTask.mockReturnValue(null); + mocks.session = null; }); it("forks a completed cloud task and opens the child", async () => { @@ -64,8 +77,109 @@ describe("ForkTaskButton", () => { ); fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - await waitFor(() => expect(mocks.forkTask).toHaveBeenCalledWith(cloudTask, expect.any(Function))); - expect(mocks.navigateToTaskDetail).toHaveBeenCalledWith("task-2"); + await waitFor(() => + expect(mocks.forkTask).toHaveBeenCalledWith( + cloudTask, + expect.any(Function), + ), + ); + expect(mocks.openTask).toHaveBeenCalledWith( + expect.objectContaining({ id: "task-2" }), + ); + expect(mocks.navigateToTaskDetail).not.toHaveBeenCalledWith("task-2"); + }); + + it("uses the live cloud status before the persisted run status", () => { + mocks.session = { cloudStatus: "completed" }; + + render( + + + , + ); + + expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); + }); + + it("disables Pi tasks with the runtime-specific reason", () => { + render( + + + , + ); + + const button = screen.getByRole("button", { name: "Fork task" }); + expect(button).toHaveAttribute("aria-disabled", "true"); + expect(button.closest("[data-tooltip]")).toHaveAttribute( + "data-tooltip", + "Pi tasks cannot be forked yet", + ); + }); + + it("disables scratch tasks as ineligible for local forking", () => { + mocks.useWorkspace.mockReturnValue({ mode: "local", isScratch: true }); + mocks.session = { status: "connected", isPromptPending: false }; + + render( + + + , + ); + + const button = screen.getByRole("button", { name: "Fork task" }); + expect(button).toHaveAttribute("aria-disabled", "true"); + expect(button.closest("[data-tooltip]")).toHaveAttribute( + "data-tooltip", + "Only repository-backed local tasks can be forked", + ); + }); + + it.each([ + { + session: { status: "disconnected", isPromptPending: false }, + reason: "Reconnect the local task before forking it", + }, + { + session: { status: "connected", isPromptPending: true }, + reason: "Wait for the local task to finish before forking it", + }, + ])("explains unavailable local state: $reason", ({ session, reason }) => { + mocks.useWorkspace.mockReturnValue({ mode: "local", isScratch: false }); + mocks.session = session; + + render( + + + , + ); + + const button = screen.getByRole("button", { name: "Fork task" }); + expect(button).toHaveAttribute("aria-disabled", "true"); + expect(button.closest("[data-tooltip]")).toHaveAttribute( + "data-tooltip", + reason, + ); }); it("links a forked task back to its source", () => { diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index e48b12390f..7d6ab58a47 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -4,39 +4,58 @@ import type { TaskForkService } from "@posthog/core/task-detail/taskForkService" import { useService } from "@posthog/di/react"; import { Button } from "@posthog/quill"; import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; -import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; +import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; import { Tooltip } from "@posthog/ui/primitives/Tooltip"; import { toast } from "@posthog/ui/primitives/toast"; import { navigateToTaskDetail } from "@posthog/ui/router/navigationBridge"; +import { openTask } from "@posthog/ui/router/useOpenTask"; import { useState } from "react"; +import { shallow } from "zustand/shallow"; export function ForkTaskButton({ task }: { task: Task }) { const taskForkService = useService(TASK_FORK_SERVICE); const workspace = useWorkspace(task.id); - const session = useSessionForTask(task.id); + const { cloudStatus, isPromptPending, sessionStatus } = useSessionSelector( + task.id, + (session) => ({ + cloudStatus: session?.cloudStatus, + isPromptPending: session?.isPromptPending ?? false, + sessionStatus: session?.status, + }), + shallow, + ); const [isForking, setIsForking] = useState(false); const run = task.latest_run; const isCloud = workspace?.mode === "cloud" || run?.environment === "cloud"; - const canForkCloud = !!run && isTerminalStatus(run.status); - const canForkLocal = - !!run && - !!workspace && - !workspace.isScratch && - session?.status === "connected" && - !session.isPromptPending; - const canFork = isCloud ? canForkCloud : canForkLocal; - const tooltip = canFork - ? "Fork task" - : isCloud - ? "Wait for the cloud run to finish" - : "Wait for the local task to finish"; + const currentCloudStatus = cloudStatus ?? run?.status; + + let disabledReason: string | null = null; + if (task.runtime === "pi") { + disabledReason = "Pi tasks cannot be forked yet"; + } else if (!run) { + disabledReason = "This task has no run to fork"; + } else if (isCloud && !isTerminalStatus(currentCloudStatus)) { + disabledReason = "Wait for the cloud run to finish before forking it"; + } else if (!isCloud && (!workspace || workspace.isScratch)) { + disabledReason = "Only repository-backed local tasks can be forked"; + } else if (!isCloud && sessionStatus !== "connected") { + disabledReason = + sessionStatus === "disconnected" || sessionStatus === "error" + ? "Reconnect the local task before forking it" + : "Wait for the local task to connect before forking it"; + } else if (!isCloud && isPromptPending) { + disabledReason = "Wait for the local task to finish before forking it"; + } + + const canFork = disabledReason === null; + const tooltip = disabledReason ?? "Fork task"; const handleFork = async () => { setIsForking(true); try { const result = await taskForkService.forkTask(task, ({ task: child }) => { - navigateToTaskDetail(child.id); + void openTask(child); }); if (!result.success) { toast.error("Could not fork task", { description: result.error }); diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index cf5dceaa4b..e1cb89a820 100644 --- a/packages/workspace-server/src/services/agent/agent.test.ts +++ b/packages/workspace-server/src/services/agent/agent.test.ts @@ -311,6 +311,34 @@ describe("AgentService", () => { 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 5f15ad771b..e66a2ea92d 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -1227,9 +1227,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 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/local-logs/identifiers.ts b/packages/workspace-server/src/services/local-logs/identifiers.ts index e114b6b952..e4d3ae9aaf 100644 --- a/packages/workspace-server/src/services/local-logs/identifiers.ts +++ b/packages/workspace-server/src/services/local-logs/identifiers.ts @@ -21,4 +21,8 @@ export interface ILogsService { ): 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/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 08331f5f16..106199f3ab 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -67,23 +67,24 @@ 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(), - diff: vi.fn(), raw: vi.fn(), }; -const mockApplyPatch = vi.fn(); vi.mock("@posthog/git/client", () => ({ createGitClient: () => mockGitClient, })); -vi.mock("@posthog/git/sagas/patch", () => ({ - ApplyPatchSaga: class { - run = mockApplyPatch; - }, -})); - vi.mock("@posthog/git/worktree", () => ({ WorktreeManager: class { createWorktree = mockWorktreeManager.createWorktree; @@ -94,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), @@ -206,9 +213,9 @@ describe("WorkspaceService", () => { mocks = createMocks(); service = makeService(mocks); mockGitClient.revparse.mockResolvedValue("abc123\n"); - mockGitClient.diff.mockResolvedValue(""); mockGitClient.raw.mockResolvedValue(""); - mockApplyPatch.mockResolvedValue({ success: true, data: undefined }); + mockCaptureWorktreeCheckpoint.mockResolvedValue(undefined); + mockDeleteWorktreeCheckpoint.mockResolvedValue(undefined); }); describe("reconcileCloudWorkspaces", () => { @@ -760,37 +767,20 @@ describe("WorkspaceService", () => { }); describe("createWorkspace (fork)", () => { - const tempDirs: string[] = []; - - afterEach(() => { - for (const dir of tempDirs.splice(0)) { - fs.rmSync(dir, { recursive: true, force: true }); - } - }); - - function mkTemp(prefix: string): string { - const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix)); - tempDirs.push(dir); - return dir; - } - - it("creates the child worktree at the source HEAD and applies its changes", async () => { + it("captures and restores the source workspace checkpoint", async () => { seedWorktreeTask(mocks, { taskId: "source-task", repoPath: "/repo", name: "source", worktreePath: "/worktrees/source", }); - mockWorktreeManager.createWorktree.mockResolvedValue({ + mockRestoreWorktreeFromCheckpoint.mockResolvedValue({ worktreePath: "/worktrees/child", worktreeName: "child", branchName: null, baseBranch: "abc123", createdAt: "2026-07-21T00:00:00Z", }); - mockGitClient.diff - .mockResolvedValueOnce("staged patch") - .mockResolvedValueOnce("unstaged patch"); await service.createWorkspace({ taskId: "child-task", @@ -801,56 +791,24 @@ describe("WorkspaceService", () => { forkFromTaskId: "source-task", }); - expect(mockWorktreeManager.createWorktree).toHaveBeenCalledWith( - expect.objectContaining({ baseBranch: "abc123" }), + expect(mockCaptureWorktreeCheckpoint).toHaveBeenCalledWith( + "/repo", + "/worktrees/source", + "fork-child-task", ); - expect(mockApplyPatch).toHaveBeenNthCalledWith(1, { - baseDir: "/worktrees/child", - patch: "staged patch", - cached: true, - }); - expect(mockApplyPatch).toHaveBeenNthCalledWith(2, { - baseDir: "/worktrees/child", - patch: "unstaged patch", - cached: false, - }); - }); - - it("preserves untracked symlinks", async () => { - const sourcePath = mkTemp("posthog-fork-source-"); - const targetPath = mkTemp("posthog-fork-target-"); - fs.writeFileSync(path.join(sourcePath, "target.txt"), "target"); - fs.symlinkSync("target.txt", path.join(sourcePath, "link.txt")); - seedWorktreeTask(mocks, { - taskId: "source-task", - repoPath: "/repo", - name: "source", - worktreePath: sourcePath, - }); - mockWorktreeManager.createWorktree.mockResolvedValue({ - worktreePath: targetPath, - worktreeName: "child", - branchName: null, - baseBranch: "abc123", - createdAt: "2026-07-21T00:00:00Z", - }); - mockGitClient.raw.mockResolvedValue("link.txt\0"); - - await service.createWorkspace({ - taskId: "child-task", - mainRepoPath: "/repo", - folderId: "folder-1", - folderPath: "/repo", - mode: "worktree", - forkFromTaskId: "source-task", - }); - - expect( - fs.lstatSync(path.join(targetPath, "link.txt")).isSymbolicLink(), - ).toBe(true); - expect(fs.readlinkSync(path.join(targetPath, "link.txt"))).toBe( - "target.txt", + expect(mockRestoreWorktreeFromCheckpoint).toHaveBeenCalledWith( + expect.objectContaining({ + mainRepoPath: "/repo", + worktreeBasePath: "/tmp/worktrees", + branchName: null, + checkpointId: "fork-child-task", + }), + ); + expect(mockDeleteWorktreeCheckpoint).toHaveBeenCalledWith( + "/repo", + "fork-child-task", ); + expect(mockWorktreeManager.createWorktree).not.toHaveBeenCalled(); }); }); diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index 779338120a..ebcfdd9319 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -16,7 +16,6 @@ import { } from "@posthog/git/queries"; import { CreateOrSwitchBranchSaga } from "@posthog/git/sagas/branch"; import { DetachHeadSaga } from "@posthog/git/sagas/head"; -import { ApplyPatchSaga } from "@posthog/git/sagas/patch"; import { WorktreeManager } from "@posthog/git/worktree"; import { ANALYTICS_SERVICE, @@ -47,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, @@ -764,12 +768,29 @@ export class WorkspaceService extends TypedEventEmitter if (!sourcePath) { throw new Error("The source task workspace is unavailable"); } - const sourceHead = await createGitClient(sourcePath).revparse(["HEAD"]); - worktree = await worktreeManager.createWorktree({ - baseBranch: sourceHead.trim(), - onOutput, - }); - await this.copyWorkspaceChanges(sourcePath, worktree.worktreePath); + const checkpointId = `fork-${taskId}`; + await captureWorktreeCheckpoint(mainRepoPath, sourcePath, checkpointId); + try { + worktree = await restoreWorktreeFromCheckpoint({ + mainRepoPath, + worktreeBasePath, + preferredName: undefined, + branchName: null, + checkpointId, + 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) { @@ -893,59 +914,6 @@ export class WorkspaceService extends TypedEventEmitter }; } - private async copyWorkspaceChanges( - sourcePath: string, - targetPath: string, - ): Promise { - const sourceGit = createGitClient(sourcePath); - const [stagedPatch, unstagedPatch, untrackedFiles] = await Promise.all([ - sourceGit.diff(["--binary", "--cached", "HEAD"]), - sourceGit.diff(["--binary"]), - sourceGit.raw(["ls-files", "--others", "--exclude-standard", "-z"]), - ]); - - for (const [patch, cached] of [ - [stagedPatch, true], - [unstagedPatch, false], - ] as const) { - if (!patch) continue; - const result = await new ApplyPatchSaga().run({ - baseDir: targetPath, - patch, - cached, - }); - if (!result.success) { - throw new Error( - `Failed to copy forked workspace changes: ${result.error}`, - ); - } - } - - await Promise.all( - untrackedFiles - .split("\0") - .filter(Boolean) - .map(async (filePath) => { - const sourceFile = path.join(sourcePath, filePath); - const targetFile = path.join(targetPath, filePath); - await fs.promises.mkdir(path.dirname(targetFile), { - recursive: true, - }); - const stat = await fs.promises.lstat(sourceFile); - if (stat.isSymbolicLink()) { - await fs.promises.symlink( - await fs.promises.readlink(sourceFile), - targetFile, - ); - } else if (stat.isFile()) { - await fs.promises.copyFile(sourceFile, targetFile); - } else { - throw new Error(`Cannot copy untracked path: ${filePath}`); - } - }), - ); - } - async deleteWorkspace(taskId: string, mainRepoPath: string): Promise { this.log.info(`Deleting workspace for task ${taskId}`); 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..e497716ce4 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); @@ -101,6 +105,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 +155,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..e83522f382 100644 --- a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts +++ b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts @@ -16,6 +16,7 @@ export interface RestoreWorktreeFromCheckpointParams { checkpointId: string; recreateBranch?: boolean; logger?: SagaLogger; + onOutput?: (data: string) => void; } /** @@ -34,31 +35,61 @@ 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, + }); + 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; @@ -85,3 +116,10 @@ export async function captureWorktreeCheckpoint( throw new Error(`Failed to capture checkpoint: ${result.error}`); } } + +export async function deleteWorktreeCheckpoint( + folderPath: string, + checkpointId: string, +): Promise { + await deleteCheckpoint(createGitClient(folderPath), checkpointId); +} From 1619f18724d583cb0bba4efcfae4a63074a54b3b Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 13:12:01 +0300 Subject: [PATCH 03/10] Fix task fork recovery and status handling Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- packages/core/src/sessions/sessionService.ts | 24 +++---- .../sessionServiceForkLocalTask.test.ts | 47 +++++++++++--- .../src/task-detail/taskForkService.test.ts | 16 +++++ .../core/src/task-detail/taskForkService.ts | 15 ++++- packages/git/src/sagas/checkpoint.test.ts | 38 ++++++++++- packages/git/src/sagas/checkpoint.ts | 3 +- .../components/ForkTaskButton.test.tsx | 63 ++++++++++++++++++- .../task-detail/components/ForkTaskButton.tsx | 21 ++++++- .../src/services/workspace/workspace.test.ts | 1 + .../src/services/workspace/workspace.ts | 1 + .../worktree-checkpoint.test.ts | 14 +++++ .../worktree-checkpoint.ts | 2 + 12 files changed, 213 insertions(+), 32 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 7aeebcfbdf..5b42a38c97 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1540,6 +1540,18 @@ export class SessionService { 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: params.sourceTaskRunId, + }, + }, + ); + const { customInstructions, rtkEnabledLocal, spokenNarrationEnabled } = this.d.settings; const result = await this.d.trpc.agent.fork.mutate({ @@ -1560,18 +1572,6 @@ export class SessionService { model: source.model ?? this.d.DEFAULT_GATEWAY_MODEL, }); - 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: params.sourceTaskRunId, - }, - }, - ); - const session = createBaseSession( taskRun.id, params.task.id, diff --git a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts index 6f8e21ecf5..1ccfa0895f 100644 --- a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts +++ b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts @@ -59,6 +59,12 @@ function createHarness(updateTaskRun: ReturnType) { 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 fork = vi.fn().mockResolvedValue({ + channel: "child-channel", + configOptions: [{ id: "model", type: "select" }], + }); const deps = { store: { getSessionByTaskId: vi.fn(() => sourceSession), @@ -66,19 +72,14 @@ function createHarness(updateTaskRun: ReturnType) { }, trpc: { agent: { - flushLogs: { mutate: vi.fn().mockResolvedValue(undefined) }, - fork: { - mutate: vi.fn().mockResolvedValue({ - channel: "child-channel", - configOptions: [{ id: "model", type: "select" }], - }), - }, + flushLogs: { mutate: flushLogs }, + fork: { mutate: fork }, onSessionEvent: { subscribe }, onPermissionRequest: { subscribe }, onSessionIdleKilled: { subscribe }, }, logs: { - cloneLocalLogs: { mutate: vi.fn().mockResolvedValue(undefined) }, + cloneLocalLogs: { mutate: cloneLocalLogs }, }, }, fetchAuthState: vi.fn().mockResolvedValue({ @@ -111,6 +112,9 @@ function createHarness(updateTaskRun: ReturnType) { setAdapter, subscribe, track, + flushLogs, + cloneLocalLogs, + fork, }; } @@ -124,8 +128,15 @@ describe("SessionService.forkLocalTask", () => { forked_from_run_id: "source-run", }, }); - const { service, setSession, setPersistedConfigOptions, setAdapter } = - createHarness(updateTaskRun); + const { + service, + setSession, + setPersistedConfigOptions, + setAdapter, + flushLogs, + cloneLocalLogs, + fork, + } = createHarness(updateTaskRun); await service.forkLocalTask({ sourceTaskId: "source-task", @@ -141,6 +152,20 @@ describe("SessionService.forkLocalTask", () => { 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], ); @@ -163,6 +188,7 @@ describe("SessionService.forkLocalTask", () => { setAdapter, subscribe, track, + fork, } = createHarness(updateTaskRun); await expect( @@ -177,6 +203,7 @@ describe("SessionService.forkLocalTask", () => { expect(setPersistedConfigOptions).not.toHaveBeenCalled(); expect(setAdapter).not.toHaveBeenCalled(); expect(setSession).not.toHaveBeenCalled(); + expect(fork).not.toHaveBeenCalled(); expect(subscribe).toHaveBeenCalledTimes(1); expect(track).not.toHaveBeenCalled(); }); diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index e12c2a6306..c8b8e91a09 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -82,6 +82,22 @@ describe("TaskForkService", () => { }); }); + it("uses the effective cloud status when the persisted run is stale", 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, { + sourceRunStatus: "completed", + }); + + expect(result.success).toBe(true); + expect(taskService.createTask).toHaveBeenCalledOnce(); + }); + it("creates a local worktree fork", async () => { const task = createTask({ latest_run: createRun({ diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts index a8b1ec2689..f88664003a 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -3,6 +3,7 @@ import { executionModeSchema, isTerminalStatus, type Task, + type TaskRun, } from "@posthog/shared/domain-types"; import { inject, injectable } from "inversify"; import { @@ -13,6 +14,11 @@ import { TASK_CREATION_HOST, TASK_SERVICE } from "./identifiers"; import type { ITaskCreationHost } from "./taskCreationHost"; import type { CreateTaskResult, TaskService } from "./taskService"; +export interface ForkTaskOptions { + sourceRunStatus?: TaskRun["status"]; + onTaskReady?: (output: TaskCreationOutput) => void; +} + @injectable() export class TaskForkService { constructor( @@ -24,7 +30,7 @@ export class TaskForkService { async forkTask( sourceTask: Task, - onTaskReady?: (output: TaskCreationOutput) => void, + options: ForkTaskOptions = {}, ): Promise { const sourceRun = sourceTask.latest_run; if (!sourceRun) { @@ -37,7 +43,10 @@ export class TaskForkService { const sourceWorkspace = await this.host.getWorkspace(sourceTask.id); const isCloud = sourceWorkspace?.mode === "cloud" || sourceRun.environment === "cloud"; - if (isCloud && !isTerminalStatus(sourceRun.status)) { + if ( + isCloud && + !isTerminalStatus(options.sourceRunStatus ?? sourceRun.status) + ) { return this.validationError( "Wait for the cloud run to finish before forking it", ); @@ -89,7 +98,7 @@ export class TaskForkService { taskRunId: sourceRun.id, }, }, - onTaskReady, + options.onTaskReady, ); } diff --git a/packages/git/src/sagas/checkpoint.test.ts b/packages/git/src/sagas/checkpoint.test.ts index e05dfcf5cd..6bcdd3ee15 100644 --- a/packages/git/src/sagas/checkpoint.test.ts +++ b/packages/git/src/sagas/checkpoint.test.ts @@ -75,9 +75,14 @@ async function captureCheckpoint( async function revertCheckpoint( repoPath: string, checkpointId: string, + restoreBranch?: boolean, ): Promise { const revert = new RevertCheckpointSaga(); - const result = await revert.run({ baseDir: repoPath, checkpointId }); + const result = await revert.run({ + baseDir: repoPath, + checkpointId, + restoreBranch, + }); expect(result.success).toBe(true); } @@ -250,6 +255,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-"), diff --git a/packages/git/src/sagas/checkpoint.ts b/packages/git/src/sagas/checkpoint.ts index 1b80222a3b..8ac3d6a53c 100644 --- a/packages/git/src/sagas/checkpoint.ts +++ b/packages/git/src/sagas/checkpoint.ts @@ -157,6 +157,7 @@ export class CaptureCheckpointSaga extends GitSaga< export interface RevertCheckpointInput extends GitSagaInput { checkpointId: string; + restoreBranch?: boolean; } export interface RevertCheckpointOutput { @@ -210,7 +211,7 @@ export class RevertCheckpointSaga extends GitSaga< name: "checkout_head", execute: async () => { if (!head) return; - if (branch) { + if (input.restoreBranch !== false && branch) { const branchExists = await refExists( this.git, `refs/heads/${branch}`, diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 8930a7f5eb..f7bc2f6e4f 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -5,9 +5,12 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ forkTask: vi.fn(), + invalidateTasks: vi.fn(), navigateToTaskDetail: vi.fn(), openTask: vi.fn(), + setProvisioningFailed: vi.fn(), session: null as Record | null, + toastError: vi.fn(), useWorkspace: vi.fn(), })); @@ -29,6 +32,17 @@ vi.mock("@posthog/ui/router/navigationBridge", () => ({ vi.mock("@posthog/ui/router/useOpenTask", () => ({ openTask: mocks.openTask, })); +vi.mock("@posthog/ui/features/tasks/useTaskCrudMutations", () => ({ + useCreateTask: () => ({ invalidateTasks: mocks.invalidateTasks }), +})); +vi.mock("@posthog/ui/features/provisioning/store", () => ({ + useProvisioningStore: { + getState: () => ({ setFailed: mocks.setProvisioningFailed }), + }, +})); +vi.mock("@posthog/ui/features/notifications/errorDetails", () => ({ + toastError: mocks.toastError, +})); vi.mock("@posthog/ui/primitives/Tooltip", () => ({ Tooltip: ({ children, @@ -64,8 +78,11 @@ describe("ForkTaskButton", () => { it("forks a completed cloud task and opens the child", async () => { mocks.forkTask.mockImplementation( - async (_task: Task, onReady: (output: { task: Task }) => void) => { - onReady({ task: { id: "task-2" } as Task }); + async ( + _task: Task, + options: { onTaskReady: (output: { task: Task }) => void }, + ) => { + options.onTaskReady({ task: { id: "task-2" } as Task }); return { success: true, data: {} }; }, ); @@ -80,15 +97,55 @@ describe("ForkTaskButton", () => { await waitFor(() => expect(mocks.forkTask).toHaveBeenCalledWith( cloudTask, - expect.any(Function), + expect.objectContaining({ + sourceRunStatus: "completed", + onTaskReady: expect.any(Function), + }), ), ); + expect(mocks.invalidateTasks).toHaveBeenCalledWith( + expect.objectContaining({ id: "task-2" }), + ); expect(mocks.openTask).toHaveBeenCalledWith( expect.objectContaining({ id: "task-2" }), ); expect(mocks.navigateToTaskDetail).not.toHaveBeenCalledWith("task-2"); }); + it("records and reports a provisioning failure on the child task", async () => { + const child = { id: "task-2" } as Task; + mocks.forkTask.mockImplementation( + async ( + _task: Task, + options: { onTaskReady: (output: { task: Task }) => void }, + ) => { + options.onTaskReady({ task: child }); + return { + success: true, + data: { task: child, provisioningError: "git checkout failed" }, + }; + }, + ); + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => + expect(mocks.setProvisioningFailed).toHaveBeenCalledWith( + "task-2", + "git checkout failed", + ), + ); + expect(mocks.toastError).toHaveBeenCalledWith( + "Failed to create workspace", + "git checkout failed", + ); + }); + it("uses the live cloud status before the persisted run status", () => { mocks.session = { cloudStatus: "completed" }; diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 7d6ab58a47..68051047ce 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -1,10 +1,14 @@ import { GitBranch, Spinner } from "@phosphor-icons/react"; import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; +import { getErrorTitle } from "@posthog/core/task-detail/taskInput"; import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; import { useService } from "@posthog/di/react"; import { Button } from "@posthog/quill"; import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; +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 { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; import { Tooltip } from "@posthog/ui/primitives/Tooltip"; import { toast } from "@posthog/ui/primitives/toast"; @@ -15,6 +19,7 @@ import { shallow } from "zustand/shallow"; export function ForkTaskButton({ task }: { task: Task }) { const taskForkService = useService(TASK_FORK_SERVICE); + const { invalidateTasks } = useCreateTask(); const workspace = useWorkspace(task.id); const { cloudStatus, isPromptPending, sessionStatus } = useSessionSelector( task.id, @@ -54,11 +59,23 @@ export function ForkTaskButton({ task }: { task: Task }) { const handleFork = async () => { setIsForking(true); try { - const result = await taskForkService.forkTask(task, ({ task: child }) => { - void openTask(child); + const result = await taskForkService.forkTask(task, { + sourceRunStatus: currentCloudStatus, + onTaskReady: ({ task: child }) => { + invalidateTasks(child); + void openTask(child); + }, }); if (!result.success) { toast.error("Could not fork task", { description: result.error }); + } else if (result.data.provisioningError) { + useProvisioningStore + .getState() + .setFailed(result.data.task.id, result.data.provisioningError); + toastError( + getErrorTitle("workspace_creation"), + result.data.provisioningError, + ); } } catch (error) { toast.error("Could not fork task", { diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 106199f3ab..26f65cff62 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -802,6 +802,7 @@ describe("WorkspaceService", () => { worktreeBasePath: "/tmp/worktrees", branchName: null, checkpointId: "fork-child-task", + restoreBranch: false, }), ); expect(mockDeleteWorktreeCheckpoint).toHaveBeenCalledWith( diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index ebcfdd9319..f2d948055a 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -777,6 +777,7 @@ export class WorkspaceService extends TypedEventEmitter preferredName: undefined, branchName: null, checkpointId, + restoreBranch: false, logger: this.log, onOutput, }); 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 e497716ce4..1267750c49 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 @@ -96,6 +96,20 @@ describe("restoreWorktreeFromCheckpoint", () => { expect(mockRevertRun).toHaveBeenCalledWith({ baseDir: "/wt/branch", checkpointId: "cp-1", + restoreBranch: undefined, + }); + }); + + it("can restore the checkpoint without claiming its recorded branch", async () => { + await restoreWorktreeFromCheckpoint({ + ...baseParams, + restoreBranch: false, + }); + + expect(mockRevertRun).toHaveBeenCalledWith({ + baseDir: "/wt/branch", + checkpointId: "cp-1", + restoreBranch: false, }); }); 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 e83522f382..fbc10e6907 100644 --- a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts +++ b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts @@ -15,6 +15,7 @@ export interface RestoreWorktreeFromCheckpointParams { branchName: string | null; checkpointId: string; recreateBranch?: boolean; + restoreBranch?: boolean; logger?: SagaLogger; onOutput?: (data: string) => void; } @@ -63,6 +64,7 @@ export async function restoreWorktreeFromCheckpoint( const result = await revertSaga.run({ baseDir: newWorktree.worktreePath, checkpointId: params.checkpointId, + restoreBranch: params.restoreBranch, }); if (!result.success) { throw new Error( From 7bd3d5d3087204fe224ab8ed33aa536a5dd7e39d Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 13:34:54 +0300 Subject: [PATCH 04/10] Fix fork finalization and context preservation Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- .../core/src/task-detail/taskCreationHost.ts | 1 + .../src/task-detail/taskCreationSaga.test.ts | 27 +++++++ .../core/src/task-detail/taskCreationSaga.ts | 9 ++- .../src/task-detail/taskForkService.test.ts | 6 ++ .../core/src/task-detail/taskForkService.ts | 4 + packages/git/src/sagas/checkpoint.test.ts | 4 +- packages/git/src/sagas/checkpoint.ts | 4 +- .../components/ForkTaskButton.test.tsx | 78 +++++++++++++------ .../task-detail/components/ForkTaskButton.tsx | 11 +-- .../task-detail/taskCreationHostImpl.ts | 4 + .../src/services/workspace/workspace.test.ts | 2 +- .../src/services/workspace/workspace.ts | 2 +- .../worktree-checkpoint.test.ts | 6 +- .../worktree-checkpoint.ts | 4 +- 14 files changed, 122 insertions(+), 40 deletions(-) diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index 86435725f5..e3a28fd039 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/packages/core/src/task-detail/taskCreationHost.ts @@ -83,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 f0e29e6006..42123ac262 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -1301,6 +1301,33 @@ 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: { taskId: "source-task", taskRunId: "source-run" }, + }); + + expect(result.success).toBe(false); + expect(deleteTaskMock).toHaveBeenCalledWith(createdTask.id); + expect(mockHost.clearProvisioning).toHaveBeenCalledWith(createdTask.id); + expect(onTaskReady).not.toHaveBeenCalled(); + }); + 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 c6847947a5..7576f003af 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -131,7 +131,7 @@ export class TaskCreationSaga extends Saga< if (hasProvisioning) { this.deps.host.setProvisioningActive(task.id); - if (this.deps.onTaskReady) { + if (!input.forkFrom && this.deps.onTaskReady) { this.deps.onTaskReady({ task, workspace }); } } @@ -197,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", { diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index c8b8e91a09..3ce2da0f27 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -53,6 +53,7 @@ const createWorkspace = (overrides: Partial = {}): Workspace => ({ describe("TaskForkService", () => { const host = { getWorkspace: vi.fn(), + getAdditionalDirectories: vi.fn(), } as unknown as ITaskCreationHost; const taskService = { createTask: vi.fn(), @@ -62,6 +63,7 @@ describe("TaskForkService", () => { beforeEach(() => { vi.clearAllMocks(); vi.mocked(host.getWorkspace).mockResolvedValue(createWorkspace()); + vi.mocked(host.getAdditionalDirectories).mockResolvedValue([]); vi.mocked(taskService.createTask).mockResolvedValue({ success: true, data: {} as never, @@ -99,6 +101,9 @@ describe("TaskForkService", () => { }); it("creates a local worktree fork", async () => { + vi.mocked(host.getAdditionalDirectories).mockResolvedValue([ + "/repos/shared", + ]); const task = createTask({ latest_run: createRun({ runtime_adapter: "codex", @@ -121,6 +126,7 @@ describe("TaskForkService", () => { model: "gpt-5.4", reasoningLevel: "high", executionMode: "full-access", + additionalDirectories: ["/repos/shared"], forkFrom: { taskId: "source-task", taskRunId: "source-run" }, }), undefined, diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts index f88664003a..968cc038f7 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -56,6 +56,9 @@ export class TaskForkService { "Only repository-backed local tasks can be forked", ); } + const additionalDirectories = isCloud + ? undefined + : await this.host.getAdditionalDirectories(sourceTask.id); const output = sourceRun.output ?? {}; const state = sourceRun.state ?? {}; @@ -81,6 +84,7 @@ export class TaskForkService { adapter: sourceRun.runtime_adapter ?? undefined, model: sourceRun.model ?? undefined, reasoningLevel: sourceRun.reasoning_effort ?? undefined, + additionalDirectories, sandboxEnvironmentId: typeof state.sandbox_environment_id === "string" ? state.sandbox_environment_id diff --git a/packages/git/src/sagas/checkpoint.test.ts b/packages/git/src/sagas/checkpoint.test.ts index 6bcdd3ee15..c370b5a75d 100644 --- a/packages/git/src/sagas/checkpoint.test.ts +++ b/packages/git/src/sagas/checkpoint.test.ts @@ -75,13 +75,13 @@ async function captureCheckpoint( async function revertCheckpoint( repoPath: string, checkpointId: string, - restoreBranch?: boolean, + restoreCheckpointBranch?: boolean, ): Promise { const revert = new RevertCheckpointSaga(); const result = await revert.run({ baseDir: repoPath, checkpointId, - restoreBranch, + restoreCheckpointBranch, }); expect(result.success).toBe(true); } diff --git a/packages/git/src/sagas/checkpoint.ts b/packages/git/src/sagas/checkpoint.ts index 8ac3d6a53c..3148cd16f9 100644 --- a/packages/git/src/sagas/checkpoint.ts +++ b/packages/git/src/sagas/checkpoint.ts @@ -157,7 +157,7 @@ export class CaptureCheckpointSaga extends GitSaga< export interface RevertCheckpointInput extends GitSagaInput { checkpointId: string; - restoreBranch?: boolean; + restoreCheckpointBranch?: boolean; } export interface RevertCheckpointOutput { @@ -211,7 +211,7 @@ export class RevertCheckpointSaga extends GitSaga< name: "checkout_head", execute: async () => { if (!head) return; - if (input.restoreBranch !== false && branch) { + if (input.restoreCheckpointBranch !== false && branch) { const branchExists = await refExists( this.git, `refs/heads/${branch}`, diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index f7bc2f6e4f..1dda265bae 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -10,6 +10,7 @@ const mocks = vi.hoisted(() => ({ openTask: vi.fn(), setProvisioningFailed: vi.fn(), session: null as Record | null, + toast: vi.fn(), toastError: vi.fn(), useWorkspace: vi.fn(), })); @@ -53,7 +54,7 @@ vi.mock("@posthog/ui/primitives/Tooltip", () => ({ }) =>
{children}
, })); vi.mock("@posthog/ui/primitives/toast", () => ({ - toast: { error: vi.fn() }, + toast: { error: mocks.toast }, })); import { ForkedFromTaskButton, ForkTaskButton } from "./ForkTaskButton"; @@ -77,15 +78,11 @@ describe("ForkTaskButton", () => { }); it("forks a completed cloud task and opens the child", async () => { - mocks.forkTask.mockImplementation( - async ( - _task: Task, - options: { onTaskReady: (output: { task: Task }) => void }, - ) => { - options.onTaskReady({ task: { id: "task-2" } as Task }); - return { success: true, data: {} }; - }, - ); + const child = { id: "task-2" } as Task; + mocks.forkTask.mockResolvedValue({ + success: true, + data: { task: child, workspace: null }, + }); render( @@ -99,7 +96,6 @@ describe("ForkTaskButton", () => { cloudTask, expect.objectContaining({ sourceRunStatus: "completed", - onTaskReady: expect.any(Function), }), ), ); @@ -114,18 +110,14 @@ describe("ForkTaskButton", () => { it("records and reports a provisioning failure on the child task", async () => { const child = { id: "task-2" } as Task; - mocks.forkTask.mockImplementation( - async ( - _task: Task, - options: { onTaskReady: (output: { task: Task }) => void }, - ) => { - options.onTaskReady({ task: child }); - return { - success: true, - data: { task: child, provisioningError: "git checkout failed" }, - }; + mocks.forkTask.mockResolvedValue({ + success: true, + data: { + task: child, + workspace: null, + provisioningError: "git checkout failed", }, - ); + }); render( @@ -146,6 +138,48 @@ describe("ForkTaskButton", () => { ); }); + it("reports a failed fork and leaves the button ready to retry", async () => { + mocks.forkTask.mockResolvedValue({ + success: false, + error: "fork failed", + failedStep: "agent_session", + }); + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => + expect(mocks.toast).toHaveBeenCalledWith("Could not fork task", { + description: "fork failed", + }), + ); + expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); + expect(mocks.openTask).not.toHaveBeenCalled(); + }); + + it("reports a rejected fork and leaves the button ready to retry", async () => { + mocks.forkTask.mockRejectedValue(new Error("network failed")); + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => + expect(mocks.toast).toHaveBeenCalledWith("Could not fork task", { + description: "network failed", + }), + ); + expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); + expect(mocks.openTask).not.toHaveBeenCalled(); + }); + it("uses the live cloud status before the persisted run status", () => { mocks.session = { cloudStatus: "completed" }; diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 68051047ce..b663491fda 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -61,14 +61,15 @@ export function ForkTaskButton({ task }: { task: Task }) { try { const result = await taskForkService.forkTask(task, { sourceRunStatus: currentCloudStatus, - onTaskReady: ({ task: child }) => { - invalidateTasks(child); - void openTask(child); - }, }); if (!result.success) { toast.error("Could not fork task", { description: result.error }); - } else if (result.data.provisioningError) { + return; + } + + invalidateTasks(result.data.task); + void openTask(result.data.task); + if (result.data.provisioningError) { useProvisioningStore .getState() .setFailed(result.data.task.id, result.data.provisioningError); 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/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 26f65cff62..4e7431ee1f 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -802,7 +802,7 @@ describe("WorkspaceService", () => { worktreeBasePath: "/tmp/worktrees", branchName: null, checkpointId: "fork-child-task", - restoreBranch: false, + restoreCheckpointBranch: false, }), ); expect(mockDeleteWorktreeCheckpoint).toHaveBeenCalledWith( diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index f2d948055a..8074cd7242 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -777,7 +777,7 @@ export class WorkspaceService extends TypedEventEmitter preferredName: undefined, branchName: null, checkpointId, - restoreBranch: false, + restoreCheckpointBranch: false, logger: this.log, onOutput, }); 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 1267750c49..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 @@ -96,20 +96,20 @@ describe("restoreWorktreeFromCheckpoint", () => { expect(mockRevertRun).toHaveBeenCalledWith({ baseDir: "/wt/branch", checkpointId: "cp-1", - restoreBranch: undefined, + restoreCheckpointBranch: undefined, }); }); it("can restore the checkpoint without claiming its recorded branch", async () => { await restoreWorktreeFromCheckpoint({ ...baseParams, - restoreBranch: false, + restoreCheckpointBranch: false, }); expect(mockRevertRun).toHaveBeenCalledWith({ baseDir: "/wt/branch", checkpointId: "cp-1", - restoreBranch: false, + restoreCheckpointBranch: false, }); }); 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 fbc10e6907..86be8bca3b 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,7 @@ export interface RestoreWorktreeFromCheckpointParams { branchName: string | null; checkpointId: string; recreateBranch?: boolean; - restoreBranch?: boolean; + restoreCheckpointBranch?: boolean; logger?: SagaLogger; onOutput?: (data: string) => void; } @@ -64,7 +64,7 @@ export async function restoreWorktreeFromCheckpoint( const result = await revertSaga.run({ baseDir: newWorktree.worktreePath, checkpointId: params.checkpointId, - restoreBranch: params.restoreBranch, + restoreCheckpointBranch: params.restoreCheckpointBranch, }); if (!result.success) { throw new Error( From 8fcc9f9252f4a09f0017d8ab4e09ee586b999f57 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 13:46:22 +0300 Subject: [PATCH 05/10] Prevent stale and lossy task forks Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- .../src/task-detail/taskCreationSaga.test.ts | 13 +++++++- .../core/src/task-detail/taskCreationSaga.ts | 8 +++-- packages/git/src/sagas/checkpoint.test.ts | 32 ++++++++++++++++++- packages/git/src/sagas/checkpoint.ts | 17 ++++++++-- .../components/ForkTaskButton.test.tsx | 24 +++++++++++++- .../task-detail/components/ForkTaskButton.tsx | 25 +++++++++------ .../src/services/agent/agent.ts | 15 ++++++++- .../src/services/agent/codex-home.ts | 7 ++-- .../src/services/workspace/workspace.test.ts | 1 + .../src/services/workspace/workspace.ts | 7 +++- .../worktree-checkpoint.ts | 7 +++- 11 files changed, 130 insertions(+), 26 deletions(-) diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 42123ac262..f0349c1b4c 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -201,6 +201,7 @@ 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( @@ -215,7 +216,7 @@ describe("TaskCreationSaga", () => { const createTaskRunMock = vi.fn().mockResolvedValue(startedRun); const startTaskRunMock = vi.fn().mockResolvedValue(startedTask); const saga = makeSaga({ - createTask: vi.fn().mockResolvedValue(createdTask), + createTask: createTaskMock, createTaskRun: createTaskRunMock, startTaskRun: startTaskRunMock, updateTaskRun: updateTaskRunMock, @@ -225,11 +226,21 @@ describe("TaskCreationSaga", () => { content: "Ship the fix", repository: "posthog/posthog", workspaceMode: "cloud", + branch: "feature/source", adapter: "codex", + sandboxEnvironmentId: "sandbox-1", + customImageId: "image-1", forkFrom: { 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" }), diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 7576f003af..8018fa1771 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -783,6 +783,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, @@ -802,7 +804,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: @@ -816,11 +818,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/git/src/sagas/checkpoint.test.ts b/packages/git/src/sagas/checkpoint.test.ts index c370b5a75d..5f4c1b5dab 100644 --- a/packages/git/src/sagas/checkpoint.test.ts +++ b/packages/git/src/sagas/checkpoint.test.ts @@ -66,9 +66,14 @@ 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); } @@ -507,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 3148cd16f9..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", () => @@ -399,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, @@ -414,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/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 1dda265bae..6af92ee600 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -181,7 +181,7 @@ describe("ForkTaskButton", () => { }); it("uses the live cloud status before the persisted run status", () => { - mocks.session = { cloudStatus: "completed" }; + mocks.session = { taskRunId: "run-1", cloudStatus: "completed" }; render( @@ -199,6 +199,28 @@ describe("ForkTaskButton", () => { expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); }); + it("ignores a terminal cloud status from an older run", () => { + mocks.session = { taskRunId: "old-run", cloudStatus: "completed" }; + + render( + + + , + ); + + expect(screen.getByRole("button", { name: "Fork task" })).toHaveAttribute( + "aria-disabled", + "true", + ); + }); + it("disables Pi tasks with the runtime-specific reason", () => { render( diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index b663491fda..8ae6e048fb 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -21,19 +21,24 @@ export function ForkTaskButton({ task }: { task: Task }) { const taskForkService = useService(TASK_FORK_SERVICE); const { invalidateTasks } = useCreateTask(); const workspace = useWorkspace(task.id); - const { cloudStatus, isPromptPending, sessionStatus } = useSessionSelector( - task.id, - (session) => ({ - cloudStatus: session?.cloudStatus, - isPromptPending: session?.isPromptPending ?? false, - sessionStatus: session?.status, - }), - shallow, - ); + const { cloudStatus, isPromptPending, sessionStatus, sessionTaskRunId } = + useSessionSelector( + task.id, + (session) => ({ + cloudStatus: session?.cloudStatus, + isPromptPending: session?.isPromptPending ?? false, + sessionStatus: session?.status, + sessionTaskRunId: session?.taskRunId, + }), + shallow, + ); const [isForking, setIsForking] = useState(false); const run = task.latest_run; const isCloud = workspace?.mode === "cloud" || run?.environment === "cloud"; - const currentCloudStatus = cloudStatus ?? run?.status; + const currentCloudStatus = + run && sessionTaskRunId === run.id + ? (cloudStatus ?? run.status) + : run?.status; let disabledReason: string | null = null; if (task.runtime === "pi") { diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index e66a2ea92d..646c20c36a 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -1012,8 +1012,15 @@ If a repository IS genuinely required, attach one in this priority order: 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: forkSource.sessionId, + sessionId: getAgentSessionId(liveForkSource), cwd: repoPath, mcpServers: sessionMcpServers, _meta: { @@ -1220,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( diff --git a/packages/workspace-server/src/services/agent/codex-home.ts b/packages/workspace-server/src/services/agent/codex-home.ts index 85615fee23..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, @@ -111,10 +112,6 @@ export async function copyCodexSessionState(options: { for (const directory of ["sessions", "shell_snapshots"]) { const source = path.join(sourceHome, directory); - if (!fs.existsSync(source)) continue; - await fs.promises.cp(source, path.join(targetHome, directory), { - recursive: true, - force: true, - }); + await clonePath(source, path.join(targetHome, directory)); } } diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 4e7431ee1f..359bcf5d6a 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -795,6 +795,7 @@ describe("WorkspaceService", () => { "/repo", "/worktrees/source", "fork-child-task", + null, ); expect(mockRestoreWorktreeFromCheckpoint).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index 8074cd7242..75c764b9cd 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -769,7 +769,12 @@ export class WorkspaceService extends TypedEventEmitter throw new Error("The source task workspace is unavailable"); } const checkpointId = `fork-${taskId}`; - await captureWorktreeCheckpoint(mainRepoPath, sourcePath, checkpointId); + await captureWorktreeCheckpoint( + mainRepoPath, + sourcePath, + checkpointId, + null, + ); try { worktree = await restoreWorktreeFromCheckpoint({ mainRepoPath, 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 86be8bca3b..f92457ec02 100644 --- a/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts +++ b/packages/workspace-server/src/services/worktree-checkpoint/worktree-checkpoint.ts @@ -106,6 +106,7 @@ export async function captureWorktreeCheckpoint( folderPath: string, worktreePath: string, checkpointId: string, + maxWorktreeFileBytes?: number | null, ): Promise { const git = createGitClient(folderPath); try { @@ -113,7 +114,11 @@ 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}`); } From 7bc83488ed2f5cb439ca5a735461b0c3eb02ad04 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 13:58:02 +0300 Subject: [PATCH 06/10] Wire task fork services across desktop Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- apps/code/src/main/di/container.ts | 16 ++++++++++++++++ apps/code/src/renderer/di/bindings.ts | 3 +++ apps/code/src/renderer/di/container.ts | 6 ++++++ .../components/ForkTaskButton.test.tsx | 3 +++ .../task-detail/components/ForkTaskButton.tsx | 7 +++++-- 5 files changed, 33 insertions(+), 2 deletions(-) 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/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 6af92ee600..612282c8ff 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -132,6 +132,9 @@ describe("ForkTaskButton", () => { "git checkout failed", ), ); + expect( + mocks.setProvisioningFailed.mock.invocationCallOrder[0], + ).toBeLessThan(mocks.openTask.mock.invocationCallOrder[0]); expect(mocks.toastError).toHaveBeenCalledWith( "Failed to create workspace", "git checkout failed", diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 8ae6e048fb..22ddeaea01 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -72,12 +72,15 @@ export function ForkTaskButton({ task }: { task: Task }) { return; } - invalidateTasks(result.data.task); - void openTask(result.data.task); if (result.data.provisioningError) { useProvisioningStore .getState() .setFailed(result.data.task.id, result.data.provisioningError); + } + + invalidateTasks(result.data.task); + void openTask(result.data.task); + if (result.data.provisioningError) { toastError( getErrorTitle("workspace_creation"), result.data.provisioningError, From 87deb442f46fca7e70521aad766d2fd7691bc87b Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 14:08:26 +0300 Subject: [PATCH 07/10] Fix forked workspace cache hydration Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- .../components/ForkTaskButton.test.tsx | 39 +++++++++++++++++++ .../task-detail/components/ForkTaskButton.tsx | 14 +++++++ 2 files changed, 53 insertions(+) diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 612282c8ff..ad432dd60d 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -8,6 +8,7 @@ const mocks = vi.hoisted(() => ({ invalidateTasks: vi.fn(), navigateToTaskDetail: vi.fn(), openTask: vi.fn(), + setQueryData: vi.fn(), setProvisioningFailed: vi.fn(), session: null as Record | null, toast: vi.fn(), @@ -18,6 +19,14 @@ const mocks = vi.hoisted(() => ({ 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/workspace/useWorkspace", () => ({ useWorkspace: mocks.useWorkspace, })); @@ -108,6 +117,36 @@ describe("ForkTaskButton", () => { expect(mocks.navigateToTaskDetail).not.toHaveBeenCalledWith("task-2"); }); + it("caches the forked workspace before opening a local child", async () => { + const child = { id: "task-2" } as Task; + const workspace = { taskId: "task-2", mode: "worktree" }; + mocks.useWorkspace.mockReturnValue({ mode: "worktree" }); + mocks.session = { status: "connected", taskRunId: "run-1" }; + mocks.forkTask.mockResolvedValue({ + success: true, + data: { task: child, workspace }, + }); + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => expect(mocks.openTask).toHaveBeenCalledWith(child)); + const updateWorkspaces = mocks.setQueryData.mock.calls[0]?.[1] as ( + workspaces: Record, + ) => Record; + expect(updateWorkspaces({ existing: {} })).toEqual({ + existing: {}, + "task-2": workspace, + }); + expect(mocks.setQueryData.mock.invocationCallOrder[0]).toBeLessThan( + mocks.openTask.mock.invocationCallOrder[0], + ); + }); + it("records and reports a provisioning failure on the child task", async () => { const child = { id: "task-2" } as Task; mocks.forkTask.mockResolvedValue({ diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 22ddeaea01..93eb491abf 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -3,7 +3,9 @@ import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; import { getErrorTitle } from "@posthog/core/task-detail/taskInput"; import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; import { useService } from "@posthog/di/react"; +import { useHostTRPC } from "@posthog/host-router/react"; import { Button } from "@posthog/quill"; +import type { Workspace } from "@posthog/shared"; import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; import { toastError } from "@posthog/ui/features/notifications/errorDetails"; @@ -14,11 +16,14 @@ import { Tooltip } from "@posthog/ui/primitives/Tooltip"; import { toast } from "@posthog/ui/primitives/toast"; import { navigateToTaskDetail } from "@posthog/ui/router/navigationBridge"; import { openTask } from "@posthog/ui/router/useOpenTask"; +import { useQueryClient } from "@tanstack/react-query"; import { useState } from "react"; import { shallow } from "zustand/shallow"; export function ForkTaskButton({ task }: { task: Task }) { const taskForkService = useService(TASK_FORK_SERVICE); + const trpc = useHostTRPC(); + const queryClient = useQueryClient(); const { invalidateTasks } = useCreateTask(); const workspace = useWorkspace(task.id); const { cloudStatus, isPromptPending, sessionStatus, sessionTaskRunId } = @@ -78,6 +83,15 @@ export function ForkTaskButton({ task }: { task: Task }) { .setFailed(result.data.task.id, result.data.provisioningError); } + if (result.data.workspace) { + queryClient.setQueryData>( + trpc.workspace.getAll.queryKey(), + (workspaces) => ({ + ...workspaces, + [result.data.task.id]: result.data.workspace as Workspace, + }), + ); + } invalidateTasks(result.data.task); void openTask(result.data.task); if (result.data.provisioningError) { From cbee73be2950ce7a3b7bade8d38bc2416a4bad95 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 15:45:34 +0300 Subject: [PATCH 08/10] Fix local and cloud task fork sources Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- packages/core/src/sessions/sessionService.ts | 12 +- .../sessionServiceForkLocalTask.test.ts | 2 - .../src/task-detail/taskCreationSaga.test.ts | 11 +- .../core/src/task-detail/taskCreationSaga.ts | 9 +- .../src/task-detail/taskForkService.test.ts | 93 +++++++++++--- .../core/src/task-detail/taskForkService.ts | 92 ++++++++++---- packages/shared/src/task-creation-domain.ts | 20 +-- .../components/ForkTaskButton.test.tsx | 118 ++++++++++++++++-- .../task-detail/components/ForkTaskButton.tsx | 72 ++++++++--- 9 files changed, 337 insertions(+), 92 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 5b42a38c97..26f78e5638 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1509,12 +1509,11 @@ export class SessionService { async forkLocalTask(params: { sourceTaskId: string; - sourceTaskRunId: string; task: Task; repoPath: string; }): Promise { const source = this.d.store.getSessionByTaskId(params.sourceTaskId); - if (!source || source.taskRunId !== params.sourceTaskRunId) { + if (!source || source.isCloud) { throw new Error("The source task session is not available"); } if (source.isPromptPending) { @@ -1531,12 +1530,13 @@ export class SessionService { 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: params.sourceTaskRunId, + taskRunId: sourceTaskRunId, }); await this.d.trpc.logs.cloneLocalLogs.mutate({ - sourceTaskRunId: params.sourceTaskRunId, + sourceTaskRunId, targetTaskRunId: taskRun.id, }); @@ -1547,7 +1547,7 @@ export class SessionService { state: { ...taskRun.state, forked_from_task_id: params.sourceTaskId, - forked_from_run_id: params.sourceTaskRunId, + forked_from_run_id: sourceTaskRunId, }, }, ); @@ -1555,7 +1555,7 @@ export class SessionService { const { customInstructions, rtkEnabledLocal, spokenNarrationEnabled } = this.d.settings; const result = await this.d.trpc.agent.fork.mutate({ - sourceTaskRunId: params.sourceTaskRunId, + sourceTaskRunId, taskId: params.task.id, taskRunId: taskRun.id, repoPath: params.repoPath, diff --git a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts index 1ccfa0895f..dc98242c17 100644 --- a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts +++ b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts @@ -140,7 +140,6 @@ describe("SessionService.forkLocalTask", () => { await service.forkLocalTask({ sourceTaskId: "source-task", - sourceTaskRunId: "source-run", task: { ...childTask }, repoPath: "/repos/code", }); @@ -194,7 +193,6 @@ describe("SessionService.forkLocalTask", () => { await expect( service.forkLocalTask({ sourceTaskId: "source-task", - sourceTaskRunId: "source-run", task: { ...childTask }, repoPath: "/repos/code", }), diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index f0349c1b4c..99d530dc23 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -230,7 +230,11 @@ describe("TaskCreationSaga", () => { adapter: "codex", sandboxEnvironmentId: "sandbox-1", customImageId: "image-1", - forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + forkFrom: { + kind: "cloud", + taskId: "source-task", + taskRunId: "source-run", + }, }); expect(result.success).toBe(true); @@ -281,7 +285,7 @@ describe("TaskCreationSaga", () => { content: "Ship the fix", repoPath: "/repo", workspaceMode: "worktree", - forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + forkFrom: { kind: "local", taskId: "source-task" }, }); expect(result.success).toBe(true); @@ -290,7 +294,6 @@ describe("TaskCreationSaga", () => { ); expect(sessionService.forkLocalTask).toHaveBeenCalledWith({ sourceTaskId: "source-task", - sourceTaskRunId: "source-run", task, repoPath: "/worktrees/fork", }); @@ -1330,7 +1333,7 @@ describe("TaskCreationSaga", () => { content: "Fork the task", repoPath: "/repo", workspaceMode: "worktree", - forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + forkFrom: { kind: "local", taskId: "source-task" }, }); expect(result.success).toBe(false); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 8018fa1771..cf71a21f8c 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -418,7 +418,7 @@ export class TaskCreationSaga extends Saga< homeQuickAction: input.homeQuickActionLabel, importedMcpServers: input.importedMcpServers, relayedMcpServers: input.relayedMcpServers, - ...(input.forkFrom && { + ...(input.forkFrom?.kind === "cloud" && { resumeFromRunId: input.forkFrom.taskRunId, }), initialPermissionMode: @@ -462,7 +462,7 @@ export class TaskCreationSaga extends Saga< }, ); - if (input.forkFrom && startedRun.latest_run) { + if (input.forkFrom?.kind === "cloud" && startedRun.latest_run) { startedRun.latest_run = await this.deps.posthogClient.updateTaskRun( task.id, startedRun.latest_run.id, @@ -551,13 +551,14 @@ export class TaskCreationSaga extends Saga< connectParams.adapter = "claude"; } - if (input.forkFrom) { + if (input.forkFrom?.kind === "local") { await this.deps.sessionService.forkLocalTask({ sourceTaskId: input.forkFrom.taskId, - sourceTaskRunId: input.forkFrom.taskRunId, task, repoPath: agentCwd ?? "", }); + } else if (input.forkFrom) { + throw new Error("Cloud forks require a cloud workspace"); } else { this.deps.sessionService.connectToTask(connectParams); } diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index 3ce2da0f27..913ef8ca4d 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -51,7 +51,9 @@ const createWorkspace = (overrides: Partial = {}): Workspace => ({ }); describe("TaskForkService", () => { + const getTaskRun = vi.fn(); const host = { + getAuthenticatedClient: vi.fn(() => ({ getTaskRun })), getWorkspace: vi.fn(), getAdditionalDirectories: vi.fn(), } as unknown as ITaskCreationHost; @@ -62,6 +64,7 @@ describe("TaskForkService", () => { beforeEach(() => { vi.clearAllMocks(); + getTaskRun.mockResolvedValue(createRun()); vi.mocked(host.getWorkspace).mockResolvedValue(createWorkspace()); vi.mocked(host.getAdditionalDirectories).mockResolvedValue([]); vi.mocked(taskService.createTask).mockResolvedValue({ @@ -75,7 +78,13 @@ describe("TaskForkService", () => { latest_run: createRun({ environment: "cloud", status: "in_progress" }), }); - const result = await service.forkTask(task); + const result = await service.forkTask(task, { + source: { + kind: "cloud", + taskRunId: "source-run", + status: "in_progress", + }, + }); expect(result).toEqual({ success: false, @@ -93,7 +102,11 @@ describe("TaskForkService", () => { }); const result = await service.forkTask(task, { - sourceRunStatus: "completed", + source: { + kind: "cloud", + taskRunId: "source-run", + status: "completed", + }, }); expect(result.success).toBe(true); @@ -104,16 +117,9 @@ describe("TaskForkService", () => { vi.mocked(host.getAdditionalDirectories).mockResolvedValue([ "/repos/shared", ]); - const task = createTask({ - latest_run: createRun({ - runtime_adapter: "codex", - model: "gpt-5.4", - reasoning_effort: "high", - state: { initial_permission_mode: "full-access" }, - }), - }); + const task = createTask({ latest_run: undefined }); - await service.forkTask(task); + await service.forkTask(task, { source: { kind: "local" } }); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ @@ -122,15 +128,48 @@ describe("TaskForkService", () => { repository: "PostHog/code", workspaceMode: "worktree", branch: null, - adapter: "codex", - model: "gpt-5.4", - reasoningLevel: "high", - executionMode: "full-access", additionalDirectories: ["/repos/shared"], - forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + 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", + status: "completed", + }), + ); + const task = createTask({ + latest_run: createRun({ id: "stale-local-run", environment: "local" }), + }); + + const result = await service.forkTask(task, { + source: { + kind: "cloud", + taskRunId: "live-cloud-run", + status: "completed", + }, + }); + + 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 () => { @@ -158,7 +197,13 @@ describe("TaskForkService", () => { }), }); - await service.forkTask(task); + await service.forkTask(task, { + source: { + kind: "cloud", + taskRunId: "source-run", + status: "completed", + }, + }); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ @@ -176,7 +221,11 @@ describe("TaskForkService", () => { cloudRtkEnabled: false, cloudRunSource: "signal_report", cloudPrAuthorshipMode: "bot", - forkFrom: { taskId: "source-task", taskRunId: "source-run" }, + forkFrom: { + kind: "cloud", + taskId: "source-task", + taskRunId: "source-run", + }, }), undefined, ); @@ -193,7 +242,13 @@ describe("TaskForkService", () => { }), }); - await service.forkTask(task); + await service.forkTask(task, { + source: { + kind: "cloud", + taskRunId: "source-run", + status: "completed", + }, + }); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts index 968cc038f7..f1ebb52ecc 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -15,7 +15,13 @@ import type { ITaskCreationHost } from "./taskCreationHost"; import type { CreateTaskResult, TaskService } from "./taskService"; export interface ForkTaskOptions { - sourceRunStatus?: TaskRun["status"]; + source: + | { kind: "local" } + | { + kind: "cloud"; + taskRunId: string; + status: TaskRun["status"]; + }; onTaskReady?: (output: TaskCreationOutput) => void; } @@ -30,35 +36,75 @@ export class TaskForkService { async forkTask( sourceTask: Task, - options: ForkTaskOptions = {}, + options: ForkTaskOptions, ): Promise { - const sourceRun = sourceTask.latest_run; - if (!sourceRun) { - return this.validationError("The source task has no run to fork"); - } if (sourceTask.runtime === "pi") { return this.validationError("Pi tasks cannot be forked yet"); } const sourceWorkspace = await this.host.getWorkspace(sourceTask.id); - const isCloud = - sourceWorkspace?.mode === "cloud" || sourceRun.environment === "cloud"; - if ( - isCloud && - !isTerminalStatus(options.sourceRunStatus ?? sourceRun.status) - ) { - return this.validationError( - "Wait for the cloud run to finish before forking it", + if (options.source.kind === "local") { + 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, ); } - if (!isCloud && (!sourceWorkspace || sourceWorkspace.isScratch)) { + + let sourceRun = sourceTask.latest_run; + if (!sourceRun || sourceRun.id !== options.source.taskRunId) { + const client = await this.host.getAuthenticatedClient(); + if (!client) { + return this.validationError("Not authenticated"); + } + try { + sourceRun = await client.getTaskRun( + sourceTask.id, + options.source.taskRunId, + ); + } 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 (!isTerminalStatus(options.source.status)) { return this.validationError( - "Only repository-backed local tasks can be forked", + "Wait for the cloud run to finish before forking it", ); } - const additionalDirectories = isCloud - ? undefined - : await this.host.getAdditionalDirectories(sourceTask.id); const output = sourceRun.output ?? {}; const state = sourceRun.state ?? {}; @@ -71,10 +117,10 @@ export class TaskForkService { { content: sourceTask.description || sourceTask.title, taskDescription: sourceTask.description || sourceTask.title, - repoPath: isCloud ? undefined : sourceWorkspace?.folderPath, + repoPath: undefined, repository: sourceTask.repository, - workspaceMode: isCloud ? "cloud" : "worktree", - branch: isCloud ? cloudBranch : null, + workspaceMode: "cloud", + branch: cloudBranch, githubIntegrationId: sourceTask.github_integration ?? undefined, githubUserIntegrationId: sourceTask.github_user_integration ?? undefined, @@ -84,7 +130,6 @@ export class TaskForkService { adapter: sourceRun.runtime_adapter ?? undefined, model: sourceRun.model ?? undefined, reasoningLevel: sourceRun.reasoning_effort ?? undefined, - additionalDirectories, sandboxEnvironmentId: typeof state.sandbox_environment_id === "string" ? state.sandbox_environment_id @@ -98,6 +143,7 @@ export class TaskForkService { cloudRunSource: getCloudRunSource(state), cloudPrAuthorshipMode: getCloudPrAuthorshipMode(state), forkFrom: { + kind: "cloud", taskId: sourceTask.id, taskRunId: sourceRun.id, }, diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 785aca706c..8c85c74c6d 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -109,15 +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 the source run's current conversation and - * workspace state. Local runs use the native agent fork API; cloud runs use - * the backend's persisted run snapshot. - */ - forkFrom?: { - taskId: string; - taskRunId: string; - }; + /** 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/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index ad432dd60d..1081f8e43e 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -104,7 +104,11 @@ describe("ForkTaskButton", () => { expect(mocks.forkTask).toHaveBeenCalledWith( cloudTask, expect.objectContaining({ - sourceRunStatus: "completed", + source: { + kind: "cloud", + taskRunId: "run-1", + status: "completed", + }, }), ), ); @@ -121,7 +125,11 @@ describe("ForkTaskButton", () => { const child = { id: "task-2" } as Task; const workspace = { taskId: "task-2", mode: "worktree" }; mocks.useWorkspace.mockReturnValue({ mode: "worktree" }); - mocks.session = { status: "connected", taskRunId: "run-1" }; + mocks.session = { + status: "connected", + taskRunId: "run-1", + isCloud: false, + }; mocks.forkTask.mockResolvedValue({ success: true, data: { task: child, workspace }, @@ -223,7 +231,12 @@ describe("ForkTaskButton", () => { }); it("uses the live cloud status before the persisted run status", () => { - mocks.session = { taskRunId: "run-1", cloudStatus: "completed" }; + mocks.session = { + status: "connected", + taskRunId: "run-1", + isCloud: true, + cloudStatus: "completed", + }; render( @@ -242,7 +255,12 @@ describe("ForkTaskButton", () => { }); it("ignores a terminal cloud status from an older run", () => { - mocks.session = { taskRunId: "old-run", cloudStatus: "completed" }; + mocks.session = { + status: "disconnected", + taskRunId: "old-run", + isCloud: true, + cloudStatus: "completed", + }; render( @@ -280,7 +298,11 @@ describe("ForkTaskButton", () => { it("disables scratch tasks as ineligible for local forking", () => { mocks.useWorkspace.mockReturnValue({ mode: "local", isScratch: true }); - mocks.session = { status: "connected", isPromptPending: false }; + mocks.session = { + status: "connected", + isCloud: false, + isPromptPending: false, + }; render( @@ -305,11 +327,19 @@ describe("ForkTaskButton", () => { it.each([ { - session: { status: "disconnected", isPromptPending: false }, + session: { + status: "disconnected", + isCloud: false, + isPromptPending: false, + }, reason: "Reconnect the local task before forking it", }, { - session: { status: "connected", isPromptPending: true }, + session: { + status: "connected", + isCloud: false, + isPromptPending: true, + }, reason: "Wait for the local task to finish before forking it", }, ])("explains unavailable local state: $reason", ({ session, reason }) => { @@ -337,6 +367,80 @@ describe("ForkTaskButton", () => { ); }); + it("forks a connected local task without a latest run", async () => { + const child = { id: "task-2" } as Task; + mocks.useWorkspace.mockReturnValue({ mode: "worktree", isScratch: false }); + mocks.session = { + status: "connected", + taskRunId: "live-local-run", + isCloud: false, + isPromptPending: false, + }; + mocks.forkTask.mockResolvedValue({ + success: true, + data: { task: child, workspace: null }, + }); + + const task = { ...cloudTask, latest_run: undefined } as Task; + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => + expect(mocks.forkTask).toHaveBeenCalledWith( + task, + expect.objectContaining({ source: { kind: "local" } }), + ), + ); + }); + + it("uses the connected cloud session instead of a stale local latest run", async () => { + const child = { id: "task-2" } as Task; + mocks.useWorkspace.mockReturnValue({ mode: "worktree", isScratch: false }); + mocks.session = { + status: "connected", + taskRunId: "live-cloud-run", + isCloud: true, + cloudStatus: "completed", + isPromptPending: false, + }; + mocks.forkTask.mockResolvedValue({ + success: true, + data: { task: child, workspace: null }, + }); + const task = { + ...cloudTask, + latest_run: { + ...cloudTask.latest_run, + id: "stale-local-run", + environment: "local", + }, + } as Task; + + render( + + + , + ); + fireEvent.click(screen.getByRole("button", { name: "Fork task" })); + + await waitFor(() => + expect(mocks.forkTask).toHaveBeenCalledWith( + task, + expect.objectContaining({ + source: { + kind: "cloud", + taskRunId: "live-cloud-run", + status: "completed", + }, + }), + ), + ); + }); + it("links a forked task back to its source", () => { const child = { ...cloudTask, diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 93eb491abf..453beac18a 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -1,7 +1,10 @@ import { GitBranch, Spinner } from "@phosphor-icons/react"; import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; import { getErrorTitle } from "@posthog/core/task-detail/taskInput"; -import type { TaskForkService } from "@posthog/core/task-detail/taskForkService"; +import type { + ForkTaskOptions, + TaskForkService, +} from "@posthog/core/task-detail/taskForkService"; import { useService } from "@posthog/di/react"; import { useHostTRPC } from "@posthog/host-router/react"; import { Button } from "@posthog/quill"; @@ -26,31 +29,52 @@ export function ForkTaskButton({ task }: { task: Task }) { const queryClient = useQueryClient(); const { invalidateTasks } = useCreateTask(); const workspace = useWorkspace(task.id); - const { cloudStatus, isPromptPending, sessionStatus, sessionTaskRunId } = - useSessionSelector( - task.id, - (session) => ({ - cloudStatus: session?.cloudStatus, - isPromptPending: session?.isPromptPending ?? false, - sessionStatus: session?.status, - sessionTaskRunId: session?.taskRunId, - }), - shallow, - ); + const { + cloudStatus, + isPromptPending, + sessionIsCloud, + sessionStatus, + sessionTaskRunId, + } = useSessionSelector( + task.id, + (session) => ({ + cloudStatus: session?.cloudStatus, + isPromptPending: session?.isPromptPending ?? false, + sessionIsCloud: session?.isCloud, + sessionStatus: session?.status, + sessionTaskRunId: session?.taskRunId, + }), + shallow, + ); const [isForking, setIsForking] = useState(false); const run = task.latest_run; - const isCloud = workspace?.mode === "cloud" || run?.environment === "cloud"; + const hasLiveSession = + sessionStatus === "connected" || sessionStatus === "connecting"; + const isCloud = + hasLiveSession && sessionIsCloud !== undefined + ? sessionIsCloud + : workspace + ? workspace.mode === "cloud" + : run?.environment === "cloud"; + const cloudTaskRunId = isCloud + ? hasLiveSession && sessionIsCloud + ? sessionTaskRunId + : run?.id + : undefined; const currentCloudStatus = - run && sessionTaskRunId === run.id - ? (cloudStatus ?? run.status) + cloudTaskRunId === sessionTaskRunId + ? (cloudStatus ?? (run?.id === cloudTaskRunId ? run?.status : undefined)) : run?.status; let disabledReason: string | null = null; if (task.runtime === "pi") { disabledReason = "Pi tasks cannot be forked yet"; - } else if (!run) { + } else if (isCloud && !cloudTaskRunId) { disabledReason = "This task has no run to fork"; - } else if (isCloud && !isTerminalStatus(currentCloudStatus)) { + } else if ( + isCloud && + (!currentCloudStatus || !isTerminalStatus(currentCloudStatus)) + ) { disabledReason = "Wait for the cloud run to finish before forking it"; } else if (!isCloud && (!workspace || workspace.isScratch)) { disabledReason = "Only repository-backed local tasks can be forked"; @@ -67,10 +91,22 @@ export function ForkTaskButton({ task }: { task: Task }) { const tooltip = disabledReason ?? "Fork task"; const handleFork = async () => { + let source: ForkTaskOptions["source"]; + if (isCloud) { + if (!cloudTaskRunId || !currentCloudStatus) return; + source = { + kind: "cloud", + taskRunId: cloudTaskRunId, + status: currentCloudStatus, + }; + } else { + source = { kind: "local" }; + } + setIsForking(true); try { const result = await taskForkService.forkTask(task, { - sourceRunStatus: currentCloudStatus, + source, }); if (!result.success) { toast.error("Could not fork task", { description: result.error }); From 1b3bed5c76f347c9913089b6e698c064bdeb8936 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 16:10:54 +0300 Subject: [PATCH 09/10] Fix task fork readiness and cache stability Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- packages/core/src/sessions/sessionService.ts | 4 ++ .../src/task-detail/taskForkService.test.ts | 26 ++++++---- .../core/src/task-detail/taskForkService.ts | 48 +++++++++++++++++-- .../components/ForkTaskButton.test.tsx | 20 ++++---- .../task-detail/components/ForkTaskButton.tsx | 29 +++++++---- .../tasks/useTaskCrudMutations.test.tsx | 21 ++++++++ .../features/tasks/useTaskCrudMutations.ts | 44 +++++++++-------- 7 files changed, 137 insertions(+), 55 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 26f78e5638..6b424fbcbc 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1608,6 +1608,10 @@ export class SessionService { }); } + getSessionByTaskId(taskId: string): AgentSession | undefined { + return this.d.store.getSessionByTaskId(taskId); + } + async loadLogsOnly(params: { taskId: string; taskRunId: string; diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index 913ef8ca4d..3f6c28d717 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ITaskCreationHost } from "./taskCreationHost"; import { TaskForkService } from "./taskForkService"; import type { TaskService } from "./taskService"; +import type { SessionService } from "../sessions/sessionService"; const createTask = (overrides: Partial = {}): Task => ({ id: "source-task", @@ -60,20 +61,24 @@ describe("TaskForkService", () => { const taskService = { createTask: vi.fn(), } as unknown as TaskService; - const service = new TaskForkService(host, 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 an active cloud run", async () => { + it("rejects a cloud run while its current turn is active", async () => { const task = createTask({ latest_run: createRun({ environment: "cloud", status: "in_progress" }), }); @@ -82,21 +87,27 @@ describe("TaskForkService", () => { source: { kind: "cloud", taskRunId: "source-run", - status: "in_progress", }, }); expect(result).toEqual({ success: false, - error: "Wait for the cloud run to finish before forking it", + error: "Wait for the current cloud turn to finish before forking it", failedStep: "validation", }); }); - it("uses the effective cloud status when the persisted run is stale", async () => { + 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" }), }); @@ -105,7 +116,6 @@ describe("TaskForkService", () => { source: { kind: "cloud", taskRunId: "source-run", - status: "completed", }, }); @@ -141,7 +151,6 @@ describe("TaskForkService", () => { createRun({ id: "live-cloud-run", environment: "cloud", - status: "completed", }), ); const task = createTask({ @@ -152,7 +161,6 @@ describe("TaskForkService", () => { source: { kind: "cloud", taskRunId: "live-cloud-run", - status: "completed", }, }); @@ -201,7 +209,6 @@ describe("TaskForkService", () => { source: { kind: "cloud", taskRunId: "source-run", - status: "completed", }, }); @@ -246,7 +253,6 @@ describe("TaskForkService", () => { source: { kind: "cloud", taskRunId: "source-run", - status: "completed", }, }); diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts index f1ebb52ecc..c38a07fb1e 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -1,15 +1,19 @@ -import type { TaskCreationOutput } from "@posthog/shared"; +import type { AgentSession, TaskCreationOutput } from "@posthog/shared"; import { executionModeSchema, isTerminalStatus, type Task, - type TaskRun, + 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"; @@ -20,11 +24,37 @@ export interface ForkTaskOptions { | { kind: "cloud"; taskRunId: string; - status: TaskRun["status"]; }; 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) + ); +} + @injectable() export class TaskForkService { constructor( @@ -32,6 +62,8 @@ export class TaskForkService { private readonly host: ITaskCreationHost, @inject(TASK_SERVICE) private readonly taskService: TaskService, + @inject(SESSION_SERVICE) + private readonly sessionService: SessionService, ) {} async forkTask( @@ -100,9 +132,15 @@ export class TaskForkService { if (sourceRun.environment !== "cloud") { return this.validationError("The source run is not a cloud run"); } - if (!isTerminalStatus(options.source.status)) { + if ( + !canForkCloudRun( + sourceRun.id, + sourceRun.status, + this.sessionService.getSessionByTaskId(sourceTask.id), + ) + ) { return this.validationError( - "Wait for the cloud run to finish before forking it", + "Wait for the current cloud turn to finish before forking it", ); } diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx index 1081f8e43e..51d64dc079 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx @@ -5,7 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ forkTask: vi.fn(), - invalidateTasks: vi.fn(), + seedTask: vi.fn(), navigateToTaskDetail: vi.fn(), openTask: vi.fn(), setQueryData: vi.fn(), @@ -43,7 +43,7 @@ vi.mock("@posthog/ui/router/useOpenTask", () => ({ openTask: mocks.openTask, })); vi.mock("@posthog/ui/features/tasks/useTaskCrudMutations", () => ({ - useCreateTask: () => ({ invalidateTasks: mocks.invalidateTasks }), + useCreateTask: () => ({ seedTask: mocks.seedTask }), })); vi.mock("@posthog/ui/features/provisioning/store", () => ({ useProvisioningStore: { @@ -107,12 +107,11 @@ describe("ForkTaskButton", () => { source: { kind: "cloud", taskRunId: "run-1", - status: "completed", }, }), ), ); - expect(mocks.invalidateTasks).toHaveBeenCalledWith( + expect(mocks.seedTask).toHaveBeenCalledWith( expect.objectContaining({ id: "task-2" }), ); expect(mocks.openTask).toHaveBeenCalledWith( @@ -230,12 +229,14 @@ describe("ForkTaskButton", () => { expect(mocks.openTask).not.toHaveBeenCalled(); }); - it("uses the live cloud status before the persisted run status", () => { + it("enables a live cloud run after its current turn completes", () => { mocks.session = { status: "connected", taskRunId: "run-1", isCloud: true, - cloudStatus: "completed", + cloudStatus: "in_progress", + isPromptPending: false, + agentIdleForRunId: "run-1", }; render( @@ -254,12 +255,14 @@ describe("ForkTaskButton", () => { expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); }); - it("ignores a terminal cloud status from an older run", () => { + it("ignores idle evidence from an older cloud run", () => { mocks.session = { status: "disconnected", taskRunId: "old-run", isCloud: true, - cloudStatus: "completed", + cloudStatus: "in_progress", + isPromptPending: false, + agentIdleForRunId: "old-run", }; render( @@ -434,7 +437,6 @@ describe("ForkTaskButton", () => { source: { kind: "cloud", taskRunId: "live-cloud-run", - status: "completed", }, }), ), diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx index 453beac18a..d53044aa68 100644 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx @@ -5,11 +5,12 @@ import type { ForkTaskOptions, TaskForkService, } from "@posthog/core/task-detail/taskForkService"; +import { canForkCloudRun } from "@posthog/core/task-detail/taskForkService"; import { useService } from "@posthog/di/react"; import { useHostTRPC } from "@posthog/host-router/react"; import { Button } from "@posthog/quill"; import type { Workspace } from "@posthog/shared"; -import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; +import type { Task } from "@posthog/shared/domain-types"; import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; import { toastError } from "@posthog/ui/features/notifications/errorDetails"; import { useProvisioningStore } from "@posthog/ui/features/provisioning/store"; @@ -27,9 +28,10 @@ export function ForkTaskButton({ task }: { task: Task }) { const taskForkService = useService(TASK_FORK_SERVICE); const trpc = useHostTRPC(); const queryClient = useQueryClient(); - const { invalidateTasks } = useCreateTask(); + const { seedTask } = useCreateTask(); const workspace = useWorkspace(task.id); const { + agentIdleForRunId, cloudStatus, isPromptPending, sessionIsCloud, @@ -39,6 +41,7 @@ export function ForkTaskButton({ task }: { task: Task }) { task.id, (session) => ({ cloudStatus: session?.cloudStatus, + agentIdleForRunId: session?.agentIdleForRunId, isPromptPending: session?.isPromptPending ?? false, sessionIsCloud: session?.isCloud, sessionStatus: session?.status, @@ -65,17 +68,24 @@ export function ForkTaskButton({ task }: { task: Task }) { cloudTaskRunId === sessionTaskRunId ? (cloudStatus ?? (run?.id === cloudTaskRunId ? run?.status : undefined)) : run?.status; + const canForkCurrentCloudRun = cloudTaskRunId + ? canForkCloudRun(cloudTaskRunId, currentCloudStatus, { + agentIdleForRunId, + cloudStatus, + isCloud: sessionIsCloud, + isPromptPending, + taskRunId: sessionTaskRunId ?? "", + }) + : false; let disabledReason: string | null = null; if (task.runtime === "pi") { disabledReason = "Pi tasks cannot be forked yet"; } else if (isCloud && !cloudTaskRunId) { disabledReason = "This task has no run to fork"; - } else if ( - isCloud && - (!currentCloudStatus || !isTerminalStatus(currentCloudStatus)) - ) { - disabledReason = "Wait for the cloud run to finish before forking it"; + } else if (isCloud && !canForkCurrentCloudRun) { + disabledReason = + "Wait for the current cloud turn to finish before forking it"; } else if (!isCloud && (!workspace || workspace.isScratch)) { disabledReason = "Only repository-backed local tasks can be forked"; } else if (!isCloud && sessionStatus !== "connected") { @@ -93,11 +103,10 @@ export function ForkTaskButton({ task }: { task: Task }) { const handleFork = async () => { let source: ForkTaskOptions["source"]; if (isCloud) { - if (!cloudTaskRunId || !currentCloudStatus) return; + if (!cloudTaskRunId || !canForkCurrentCloudRun) return; source = { kind: "cloud", taskRunId: cloudTaskRunId, - status: currentCloudStatus, }; } else { source = { kind: "local" }; @@ -128,7 +137,7 @@ export function ForkTaskButton({ task }: { task: Task }) { }), ); } - invalidateTasks(result.data.task); + seedTask(result.data.task); void openTask(result.data.task); if (result.data.provisioningError) { toastError( diff --git a/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx b/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx index 76bde2513e..196e344a63 100644 --- a/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx +++ b/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx @@ -255,3 +255,24 @@ 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(); + }); +}); diff --git a/packages/ui/src/features/tasks/useTaskCrudMutations.ts b/packages/ui/src/features/tasks/useTaskCrudMutations.ts index 628dde456f..751f9ba34e 100644 --- a/packages/ui/src/features/tasks/useTaskCrudMutations.ts +++ b/packages/ui/src/features/tasks/useTaskCrudMutations.ts @@ -42,27 +42,29 @@ 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) => { + // 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() }); }; @@ -87,7 +89,7 @@ export function useCreateTask() { }) as unknown as Promise, ); - return { ...mutation, invalidateTasks }; + return { ...mutation, invalidateTasks, seedTask }; } interface DeleteTaskOptions { From 8f92dc5b00fc7751c01c27364052ce937c140a52 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 21 Jul 2026 16:58:02 +0300 Subject: [PATCH 10/10] Fix local task fork continuity Preserve complete Claude transcripts, retain optimistic fork rows until server reconciliation, and move fork lineage actions into the task context menu. Generated-By: PostHog Code Task-Id: 7e31d410-f947-4069-8832-9f40ae85705a --- .../claude/claude-agent.resume-model.test.ts | 19 + .../agent/src/adapters/claude/claude-agent.ts | 4 +- .../src/context-menu/context-menu.test.ts | 24 + .../core/src/context-menu/context-menu.ts | 17 + packages/core/src/context-menu/schemas.ts | 4 + packages/core/src/sessions/sessionService.ts | 6 +- .../sessionServiceForkLocalTask.test.ts | 46 ++ .../src/task-detail/taskCreationSaga.test.ts | 5 +- .../core/src/task-detail/taskCreationSaga.ts | 2 +- .../src/task-detail/taskForkService.test.ts | 48 +- .../core/src/task-detail/taskForkService.ts | 45 +- .../core/src/tasks/contextMenuActions.test.ts | 6 + packages/core/src/tasks/contextMenuActions.ts | 6 + .../sidebar/components/SidebarMenu.tsx | 17 +- .../components/ForkTaskButton.test.tsx | 464 ------------------ .../task-detail/components/ForkTaskButton.tsx | 196 -------- .../components/TaskHeaderActions.test.tsx | 4 - .../components/TaskHeaderActions.tsx | 3 - .../src/features/tasks/retainedTasks.test.ts | 34 ++ .../ui/src/features/tasks/retainedTasks.ts | 38 ++ .../src/features/tasks/useTaskContextMenu.ts | 16 + .../tasks/useTaskCrudMutations.test.tsx | 18 + .../features/tasks/useTaskCrudMutations.ts | 11 +- .../src/features/tasks/useTaskFork.test.tsx | 108 ++++ packages/ui/src/features/tasks/useTaskFork.ts | 75 +++ packages/ui/src/features/tasks/useTasks.ts | 15 +- 26 files changed, 499 insertions(+), 732 deletions(-) delete mode 100644 packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx delete mode 100644 packages/ui/src/features/task-detail/components/ForkTaskButton.tsx create mode 100644 packages/ui/src/features/tasks/retainedTasks.test.ts create mode 100644 packages/ui/src/features/tasks/retainedTasks.ts create mode 100644 packages/ui/src/features/tasks/useTaskFork.test.tsx create mode 100644 packages/ui/src/features/tasks/useTaskFork.ts 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/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 6b424fbcbc..98f6915453 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1571,6 +1571,7 @@ export class SessionService { : undefined, model: source.model ?? this.d.DEFAULT_GATEWAY_MODEL, }); + const { rawEntries } = await this.fetchSessionLogs(undefined, taskRun.id); const session = createBaseSession( taskRun.id, @@ -1579,7 +1580,10 @@ export class SessionService { ); session.channel = result.channel; session.status = "connected"; - session.events = [...source.events]; + session.events = + rawEntries.length > 0 + ? convertStoredEntriesToEvents(rawEntries) + : [...source.events]; session.adapter = source.adapter; session.model = source.model; session.executionMode = source.executionMode; diff --git a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts index dc98242c17..d7fd726fc5 100644 --- a/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts +++ b/packages/core/src/sessions/sessionServiceForkLocalTask.test.ts @@ -61,6 +61,7 @@ function createHarness(updateTaskRun: ReturnType) { 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" }], @@ -80,6 +81,7 @@ function createHarness(updateTaskRun: ReturnType) { }, logs: { cloneLocalLogs: { mutate: cloneLocalLogs }, + readLocalLogs: { query: readLocalLogs }, }, }, fetchAuthState: vi.fn().mockResolvedValue({ @@ -114,6 +116,7 @@ function createHarness(updateTaskRun: ReturnType) { track, flushLogs, cloneLocalLogs, + readLocalLogs, fork, }; } @@ -205,4 +208,47 @@ describe("SessionService.forkLocalTask", () => { 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/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 99d530dc23..c9da88e531 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -1339,7 +1339,10 @@ describe("TaskCreationSaga", () => { expect(result.success).toBe(false); expect(deleteTaskMock).toHaveBeenCalledWith(createdTask.id); expect(mockHost.clearProvisioning).toHaveBeenCalledWith(createdTask.id); - expect(onTaskReady).not.toHaveBeenCalled(); + expect(onTaskReady).toHaveBeenCalledWith({ + task: createdTask, + workspace: null, + }); }); it("still rolls back a worktree task when a later (non-provisioning) step fails", async () => { diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index cf71a21f8c..416aae82ef 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -131,7 +131,7 @@ export class TaskCreationSaga extends Saga< if (hasProvisioning) { this.deps.host.setProvisioningActive(task.id); - if (!input.forkFrom && this.deps.onTaskReady) { + if (this.deps.onTaskReady) { this.deps.onTaskReady({ task, workspace }); } } diff --git a/packages/core/src/task-detail/taskForkService.test.ts b/packages/core/src/task-detail/taskForkService.test.ts index 3f6c28d717..dc4660a1bc 100644 --- a/packages/core/src/task-detail/taskForkService.test.ts +++ b/packages/core/src/task-detail/taskForkService.test.ts @@ -1,10 +1,10 @@ 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"; -import type { SessionService } from "../sessions/sessionService"; const createTask = (overrides: Partial = {}): Task => ({ id: "source-task", @@ -79,16 +79,14 @@ describe("TaskForkService", () => { }); 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, { - source: { - kind: "cloud", - taskRunId: "source-run", - }, - }); + const result = await service.forkTask(task); expect(result).toEqual({ success: false, @@ -112,12 +110,7 @@ describe("TaskForkService", () => { latest_run: createRun({ environment: "cloud", status: "in_progress" }), }); - const result = await service.forkTask(task, { - source: { - kind: "cloud", - taskRunId: "source-run", - }, - }); + const result = await service.forkTask(task); expect(result.success).toBe(true); expect(taskService.createTask).toHaveBeenCalledOnce(); @@ -129,7 +122,7 @@ describe("TaskForkService", () => { ]); const task = createTask({ latest_run: undefined }); - await service.forkTask(task, { source: { kind: "local" } }); + await service.forkTask(task); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ @@ -156,13 +149,14 @@ describe("TaskForkService", () => { 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, { - source: { - kind: "cloud", - taskRunId: "live-cloud-run", - }, - }); + const result = await service.forkTask(task); expect(result.success).toBe(true); expect(taskService.createTask).toHaveBeenCalledWith( @@ -205,12 +199,7 @@ describe("TaskForkService", () => { }), }); - await service.forkTask(task, { - source: { - kind: "cloud", - taskRunId: "source-run", - }, - }); + await service.forkTask(task); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ @@ -249,12 +238,7 @@ describe("TaskForkService", () => { }), }); - await service.forkTask(task, { - source: { - kind: "cloud", - taskRunId: "source-run", - }, - }); + await service.forkTask(task); expect(taskService.createTask).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/core/src/task-detail/taskForkService.ts b/packages/core/src/task-detail/taskForkService.ts index c38a07fb1e..86498fc685 100644 --- a/packages/core/src/task-detail/taskForkService.ts +++ b/packages/core/src/task-detail/taskForkService.ts @@ -19,12 +19,6 @@ import type { ITaskCreationHost } from "./taskCreationHost"; import type { CreateTaskResult, TaskService } from "./taskService"; export interface ForkTaskOptions { - source: - | { kind: "local" } - | { - kind: "cloud"; - taskRunId: string; - }; onTaskReady?: (output: TaskCreationOutput) => void; } @@ -55,6 +49,11 @@ export function canForkCloudRun( ); } +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( @@ -68,14 +67,22 @@ export class TaskForkService { async forkTask( sourceTask: Task, - options: ForkTaskOptions, + options: ForkTaskOptions = {}, ): Promise { if (sourceTask.runtime === "pi") { return this.validationError("Pi tasks cannot be forked yet"); } const sourceWorkspace = await this.host.getWorkspace(sourceTask.id); - if (options.source.kind === "local") { + 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" || @@ -107,17 +114,21 @@ export class TaskForkService { ); } + 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 !== options.source.taskRunId) { + 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, - options.source.taskRunId, - ); + sourceRun = await client.getTaskRun(sourceTask.id, sourceTaskRunId); } catch (error) { return { success: false, @@ -132,13 +143,7 @@ export class TaskForkService { if (sourceRun.environment !== "cloud") { return this.validationError("The source run is not a cloud run"); } - if ( - !canForkCloudRun( - sourceRun.id, - sourceRun.status, - this.sessionService.getSessionByTaskId(sourceTask.id), - ) - ) { + if (!canForkCloudRun(sourceRun.id, sourceRun.status, sourceSession)) { return this.validationError( "Wait for the current cloud turn to finish before forking it", ); 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/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/components/ForkTaskButton.test.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx deleted file mode 100644 index 51d64dc079..0000000000 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.test.tsx +++ /dev/null @@ -1,464 +0,0 @@ -import type { Task } from "@posthog/shared/domain-types"; -import { Theme } from "@radix-ui/themes"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -const mocks = vi.hoisted(() => ({ - forkTask: vi.fn(), - seedTask: vi.fn(), - navigateToTaskDetail: vi.fn(), - openTask: vi.fn(), - setQueryData: vi.fn(), - setProvisioningFailed: vi.fn(), - session: null as Record | null, - toast: vi.fn(), - toastError: vi.fn(), - useWorkspace: 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/workspace/useWorkspace", () => ({ - useWorkspace: mocks.useWorkspace, -})); -vi.mock("@posthog/ui/features/sessions/useSession", () => ({ - useSessionSelector: ( - _taskId: string, - select: (session: Record | undefined) => unknown, - ) => select(mocks.session ?? undefined), -})); -vi.mock("@posthog/ui/router/navigationBridge", () => ({ - navigateToTaskDetail: mocks.navigateToTaskDetail, -})); -vi.mock("@posthog/ui/router/useOpenTask", () => ({ - openTask: mocks.openTask, -})); -vi.mock("@posthog/ui/features/tasks/useTaskCrudMutations", () => ({ - useCreateTask: () => ({ seedTask: mocks.seedTask }), -})); -vi.mock("@posthog/ui/features/provisioning/store", () => ({ - useProvisioningStore: { - getState: () => ({ setFailed: mocks.setProvisioningFailed }), - }, -})); -vi.mock("@posthog/ui/features/notifications/errorDetails", () => ({ - toastError: mocks.toastError, -})); -vi.mock("@posthog/ui/primitives/Tooltip", () => ({ - Tooltip: ({ - children, - content, - }: { - children: React.ReactNode; - content: string; - }) =>
{children}
, -})); -vi.mock("@posthog/ui/primitives/toast", () => ({ - toast: { error: mocks.toast }, -})); - -import { ForkedFromTaskButton, ForkTaskButton } from "./ForkTaskButton"; - -const cloudTask = { - id: "task-1", - title: "Source task", - latest_run: { - id: "run-1", - environment: "cloud", - status: "completed", - state: {}, - }, -} as Task; - -describe("ForkTaskButton", () => { - beforeEach(() => { - vi.clearAllMocks(); - mocks.useWorkspace.mockReturnValue({ mode: "cloud" }); - mocks.session = null; - }); - - it("forks a completed cloud task and opens the child", async () => { - const child = { id: "task-2" } as Task; - mocks.forkTask.mockResolvedValue({ - success: true, - data: { task: child, workspace: null }, - }); - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.forkTask).toHaveBeenCalledWith( - cloudTask, - expect.objectContaining({ - source: { - kind: "cloud", - taskRunId: "run-1", - }, - }), - ), - ); - expect(mocks.seedTask).toHaveBeenCalledWith( - expect.objectContaining({ id: "task-2" }), - ); - expect(mocks.openTask).toHaveBeenCalledWith( - expect.objectContaining({ id: "task-2" }), - ); - expect(mocks.navigateToTaskDetail).not.toHaveBeenCalledWith("task-2"); - }); - - it("caches the forked workspace before opening a local child", async () => { - const child = { id: "task-2" } as Task; - const workspace = { taskId: "task-2", mode: "worktree" }; - mocks.useWorkspace.mockReturnValue({ mode: "worktree" }); - mocks.session = { - status: "connected", - taskRunId: "run-1", - isCloud: false, - }; - mocks.forkTask.mockResolvedValue({ - success: true, - data: { task: child, workspace }, - }); - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => expect(mocks.openTask).toHaveBeenCalledWith(child)); - const updateWorkspaces = mocks.setQueryData.mock.calls[0]?.[1] as ( - workspaces: Record, - ) => Record; - expect(updateWorkspaces({ existing: {} })).toEqual({ - existing: {}, - "task-2": workspace, - }); - expect(mocks.setQueryData.mock.invocationCallOrder[0]).toBeLessThan( - mocks.openTask.mock.invocationCallOrder[0], - ); - }); - - it("records and reports a provisioning failure on the child task", async () => { - const child = { id: "task-2" } as Task; - mocks.forkTask.mockResolvedValue({ - success: true, - data: { - task: child, - workspace: null, - provisioningError: "git checkout failed", - }, - }); - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.setProvisioningFailed).toHaveBeenCalledWith( - "task-2", - "git checkout failed", - ), - ); - expect( - mocks.setProvisioningFailed.mock.invocationCallOrder[0], - ).toBeLessThan(mocks.openTask.mock.invocationCallOrder[0]); - expect(mocks.toastError).toHaveBeenCalledWith( - "Failed to create workspace", - "git checkout failed", - ); - }); - - it("reports a failed fork and leaves the button ready to retry", async () => { - mocks.forkTask.mockResolvedValue({ - success: false, - error: "fork failed", - failedStep: "agent_session", - }); - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.toast).toHaveBeenCalledWith("Could not fork task", { - description: "fork failed", - }), - ); - expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); - expect(mocks.openTask).not.toHaveBeenCalled(); - }); - - it("reports a rejected fork and leaves the button ready to retry", async () => { - mocks.forkTask.mockRejectedValue(new Error("network failed")); - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.toast).toHaveBeenCalledWith("Could not fork task", { - description: "network failed", - }), - ); - expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); - expect(mocks.openTask).not.toHaveBeenCalled(); - }); - - it("enables a live cloud run after its current turn completes", () => { - mocks.session = { - status: "connected", - taskRunId: "run-1", - isCloud: true, - cloudStatus: "in_progress", - isPromptPending: false, - agentIdleForRunId: "run-1", - }; - - render( - - - , - ); - - expect(screen.getByRole("button", { name: "Fork task" })).toBeEnabled(); - }); - - it("ignores idle evidence from an older cloud run", () => { - mocks.session = { - status: "disconnected", - taskRunId: "old-run", - isCloud: true, - cloudStatus: "in_progress", - isPromptPending: false, - agentIdleForRunId: "old-run", - }; - - render( - - - , - ); - - expect(screen.getByRole("button", { name: "Fork task" })).toHaveAttribute( - "aria-disabled", - "true", - ); - }); - - it("disables Pi tasks with the runtime-specific reason", () => { - render( - - - , - ); - - const button = screen.getByRole("button", { name: "Fork task" }); - expect(button).toHaveAttribute("aria-disabled", "true"); - expect(button.closest("[data-tooltip]")).toHaveAttribute( - "data-tooltip", - "Pi tasks cannot be forked yet", - ); - }); - - it("disables scratch tasks as ineligible for local forking", () => { - mocks.useWorkspace.mockReturnValue({ mode: "local", isScratch: true }); - mocks.session = { - status: "connected", - isCloud: false, - isPromptPending: false, - }; - - render( - - - , - ); - - const button = screen.getByRole("button", { name: "Fork task" }); - expect(button).toHaveAttribute("aria-disabled", "true"); - expect(button.closest("[data-tooltip]")).toHaveAttribute( - "data-tooltip", - "Only repository-backed local tasks can be forked", - ); - }); - - it.each([ - { - session: { - status: "disconnected", - isCloud: false, - isPromptPending: false, - }, - reason: "Reconnect the local task before forking it", - }, - { - session: { - status: "connected", - isCloud: false, - isPromptPending: true, - }, - reason: "Wait for the local task to finish before forking it", - }, - ])("explains unavailable local state: $reason", ({ session, reason }) => { - mocks.useWorkspace.mockReturnValue({ mode: "local", isScratch: false }); - mocks.session = session; - - render( - - - , - ); - - const button = screen.getByRole("button", { name: "Fork task" }); - expect(button).toHaveAttribute("aria-disabled", "true"); - expect(button.closest("[data-tooltip]")).toHaveAttribute( - "data-tooltip", - reason, - ); - }); - - it("forks a connected local task without a latest run", async () => { - const child = { id: "task-2" } as Task; - mocks.useWorkspace.mockReturnValue({ mode: "worktree", isScratch: false }); - mocks.session = { - status: "connected", - taskRunId: "live-local-run", - isCloud: false, - isPromptPending: false, - }; - mocks.forkTask.mockResolvedValue({ - success: true, - data: { task: child, workspace: null }, - }); - - const task = { ...cloudTask, latest_run: undefined } as Task; - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.forkTask).toHaveBeenCalledWith( - task, - expect.objectContaining({ source: { kind: "local" } }), - ), - ); - }); - - it("uses the connected cloud session instead of a stale local latest run", async () => { - const child = { id: "task-2" } as Task; - mocks.useWorkspace.mockReturnValue({ mode: "worktree", isScratch: false }); - mocks.session = { - status: "connected", - taskRunId: "live-cloud-run", - isCloud: true, - cloudStatus: "completed", - isPromptPending: false, - }; - mocks.forkTask.mockResolvedValue({ - success: true, - data: { task: child, workspace: null }, - }); - const task = { - ...cloudTask, - latest_run: { - ...cloudTask.latest_run, - id: "stale-local-run", - environment: "local", - }, - } as Task; - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Fork task" })); - - await waitFor(() => - expect(mocks.forkTask).toHaveBeenCalledWith( - task, - expect.objectContaining({ - source: { - kind: "cloud", - taskRunId: "live-cloud-run", - }, - }), - ), - ); - }); - - it("links a forked task back to its source", () => { - const child = { - ...cloudTask, - latest_run: { - ...cloudTask.latest_run, - state: { forked_from_task_id: "task-parent" }, - }, - } as Task; - - render( - - - , - ); - fireEvent.click(screen.getByRole("button", { name: "Open source task" })); - - expect(mocks.navigateToTaskDetail).toHaveBeenCalledWith("task-parent"); - }); -}); diff --git a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx b/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx deleted file mode 100644 index d53044aa68..0000000000 --- a/packages/ui/src/features/task-detail/components/ForkTaskButton.tsx +++ /dev/null @@ -1,196 +0,0 @@ -import { GitBranch, Spinner } from "@phosphor-icons/react"; -import { TASK_FORK_SERVICE } from "@posthog/core/task-detail/identifiers"; -import { getErrorTitle } from "@posthog/core/task-detail/taskInput"; -import type { - ForkTaskOptions, - TaskForkService, -} from "@posthog/core/task-detail/taskForkService"; -import { canForkCloudRun } from "@posthog/core/task-detail/taskForkService"; -import { useService } from "@posthog/di/react"; -import { useHostTRPC } from "@posthog/host-router/react"; -import { Button } from "@posthog/quill"; -import type { Workspace } from "@posthog/shared"; -import type { Task } from "@posthog/shared/domain-types"; -import { useSessionSelector } from "@posthog/ui/features/sessions/useSession"; -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 { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; -import { Tooltip } from "@posthog/ui/primitives/Tooltip"; -import { toast } from "@posthog/ui/primitives/toast"; -import { navigateToTaskDetail } from "@posthog/ui/router/navigationBridge"; -import { openTask } from "@posthog/ui/router/useOpenTask"; -import { useQueryClient } from "@tanstack/react-query"; -import { useState } from "react"; -import { shallow } from "zustand/shallow"; - -export function ForkTaskButton({ task }: { task: Task }) { - const taskForkService = useService(TASK_FORK_SERVICE); - const trpc = useHostTRPC(); - const queryClient = useQueryClient(); - const { seedTask } = useCreateTask(); - const workspace = useWorkspace(task.id); - const { - agentIdleForRunId, - cloudStatus, - isPromptPending, - sessionIsCloud, - sessionStatus, - sessionTaskRunId, - } = useSessionSelector( - task.id, - (session) => ({ - cloudStatus: session?.cloudStatus, - agentIdleForRunId: session?.agentIdleForRunId, - isPromptPending: session?.isPromptPending ?? false, - sessionIsCloud: session?.isCloud, - sessionStatus: session?.status, - sessionTaskRunId: session?.taskRunId, - }), - shallow, - ); - const [isForking, setIsForking] = useState(false); - const run = task.latest_run; - const hasLiveSession = - sessionStatus === "connected" || sessionStatus === "connecting"; - const isCloud = - hasLiveSession && sessionIsCloud !== undefined - ? sessionIsCloud - : workspace - ? workspace.mode === "cloud" - : run?.environment === "cloud"; - const cloudTaskRunId = isCloud - ? hasLiveSession && sessionIsCloud - ? sessionTaskRunId - : run?.id - : undefined; - const currentCloudStatus = - cloudTaskRunId === sessionTaskRunId - ? (cloudStatus ?? (run?.id === cloudTaskRunId ? run?.status : undefined)) - : run?.status; - const canForkCurrentCloudRun = cloudTaskRunId - ? canForkCloudRun(cloudTaskRunId, currentCloudStatus, { - agentIdleForRunId, - cloudStatus, - isCloud: sessionIsCloud, - isPromptPending, - taskRunId: sessionTaskRunId ?? "", - }) - : false; - - let disabledReason: string | null = null; - if (task.runtime === "pi") { - disabledReason = "Pi tasks cannot be forked yet"; - } else if (isCloud && !cloudTaskRunId) { - disabledReason = "This task has no run to fork"; - } else if (isCloud && !canForkCurrentCloudRun) { - disabledReason = - "Wait for the current cloud turn to finish before forking it"; - } else if (!isCloud && (!workspace || workspace.isScratch)) { - disabledReason = "Only repository-backed local tasks can be forked"; - } else if (!isCloud && sessionStatus !== "connected") { - disabledReason = - sessionStatus === "disconnected" || sessionStatus === "error" - ? "Reconnect the local task before forking it" - : "Wait for the local task to connect before forking it"; - } else if (!isCloud && isPromptPending) { - disabledReason = "Wait for the local task to finish before forking it"; - } - - const canFork = disabledReason === null; - const tooltip = disabledReason ?? "Fork task"; - - const handleFork = async () => { - let source: ForkTaskOptions["source"]; - if (isCloud) { - if (!cloudTaskRunId || !canForkCurrentCloudRun) return; - source = { - kind: "cloud", - taskRunId: cloudTaskRunId, - }; - } else { - source = { kind: "local" }; - } - - setIsForking(true); - try { - const result = await taskForkService.forkTask(task, { - source, - }); - if (!result.success) { - toast.error("Could not fork task", { description: result.error }); - return; - } - - if (result.data.provisioningError) { - useProvisioningStore - .getState() - .setFailed(result.data.task.id, result.data.provisioningError); - } - - if (result.data.workspace) { - queryClient.setQueryData>( - trpc.workspace.getAll.queryKey(), - (workspaces) => ({ - ...workspaces, - [result.data.task.id]: result.data.workspace as Workspace, - }), - ); - } - seedTask(result.data.task); - void openTask(result.data.task); - if (result.data.provisioningError) { - toastError( - getErrorTitle("workspace_creation"), - result.data.provisioningError, - ); - } - } catch (error) { - toast.error("Could not fork task", { - description: error instanceof Error ? error.message : String(error), - }); - } finally { - setIsForking(false); - } - }; - - return ( - - - - - - ); -} - -export function ForkedFromTaskButton({ task }: { task: Task }) { - const parentTaskId = task.latest_run?.state?.forked_from_task_id; - if (typeof parentTaskId !== "string") return null; - - return ( - - - - - - ); -} diff --git a/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx b/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx index 464d6d0665..f80d03734d 100644 --- a/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx +++ b/packages/ui/src/features/task-detail/components/TaskHeaderActions.test.tsx @@ -79,10 +79,6 @@ vi.mock("@posthog/ui/features/diff-stats/DiffStatsBadge", () => ({ vi.mock("@posthog/ui/primitives/Tooltip", () => ({ Tooltip: ({ children }: { children: ReactNode }) => children, })); -vi.mock("./ForkTaskButton", () => ({ - ForkTaskButton: () => null, - ForkedFromTaskButton: () => null, -})); import { TaskHeaderActions } from "./TaskHeaderActions"; diff --git a/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx b/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx index c3966ac717..d69a198bd9 100644 --- a/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx +++ b/packages/ui/src/features/task-detail/components/TaskHeaderActions.tsx @@ -26,7 +26,6 @@ import { import { Tooltip } from "@posthog/ui/primitives/Tooltip"; import { Flex } from "@radix-ui/themes"; import { useState } from "react"; -import { ForkedFromTaskButton, ForkTaskButton } from "./ForkTaskButton"; const CLOUD_HANDOFF_FLAG = "phc-cloud-handoff"; @@ -159,8 +158,6 @@ export function TaskHeaderActions({ task }: { task: Task }) { )} - - {workspaceLoaded && ( <> 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 196e344a63..fc0280a7ac 100644 --- a/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx +++ b/packages/ui/src/features/tasks/useTaskCrudMutations.test.tsx @@ -275,4 +275,22 @@ describe("useCreateTask.seedTask", () => { 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 751f9ba34e..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"); @@ -43,6 +44,7 @@ export function useCreateTask() { const queryClient = useQueryClient(); 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 @@ -68,6 +70,13 @@ export function useCreateTask() { queryClient.invalidateQueries({ queryKey: taskKeys.lists() }); }; + const removeSeededTask = (taskId: string) => { + releaseRetainedTask(taskId); + queryClient.setQueriesData({ queryKey: taskKeys.lists() }, (old) => + removeTaskFromList(old, taskId), + ); + }; + const mutation = useAuthenticatedMutation( ( client, @@ -89,7 +98,7 @@ export function useCreateTask() { }) as unknown as Promise, ); - return { ...mutation, invalidateTasks, seedTask }; + 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,