diff --git a/packages/core/src/task-detail/taskCreationEffects.ts b/packages/core/src/task-detail/taskCreationEffects.ts index cce2d882ce..8573637c10 100644 --- a/packages/core/src/task-detail/taskCreationEffects.ts +++ b/packages/core/src/task-detail/taskCreationEffects.ts @@ -1,5 +1,5 @@ import type { TaskCreationInput, TaskCreationOutput } from "@posthog/shared"; -import type { TaskRun } from "@posthog/shared/domain-types"; +import type { Task, TaskRun } from "@posthog/shared/domain-types"; /** * Host-side reactions to a successful task-creation: optimistic workspace @@ -11,4 +11,11 @@ export interface TaskCreationEffects { onWorkspaceCreated(output: TaskCreationOutput): void; onCreateSuccess(output: TaskCreationOutput, input?: TaskCreationInput): void; onRunResumed(taskId: string, run: TaskRun): void; + /** + * The saga surfaced the task to the UI (onTaskReady) but a later step + * failed and rolled the task back (best-effort server-side delete), while + * caches seeded at ready-time still hold it. Undo those seeds and refetch — + * if the delete itself failed, the refetch restores the surviving task. + */ + onCreateRolledBack(task: Task): void; } diff --git a/packages/core/src/task-detail/taskService.test.ts b/packages/core/src/task-detail/taskService.test.ts index 9b147bd059..35b0b42659 100644 --- a/packages/core/src/task-detail/taskService.test.ts +++ b/packages/core/src/task-detail/taskService.test.ts @@ -131,6 +131,111 @@ describe("TaskService.resumeCloudPiRun", () => { }); }); +describe("TaskService.createTask rollback after surfacing", () => { + // Repo-less local ("channels chat box") creation: the saga creates the task + // row, provisions a scratch dir, fires onTaskReady, then connects the agent + // session — so a connect failure rolls back a task the UI already surfaced. + function makeHarness({ connectToTask }: { connectToTask: () => void }) { + const task = { id: "task-1", title: "T", description: "do the thing" }; + const api = { + createTask: vi.fn(async () => task), + deleteTask: vi.fn(async () => undefined), + }; + const host = { + getAuthenticatedClient: vi.fn(async () => api), + detectRepo: vi.fn(async () => null), + getTaskDirectory: vi.fn(async () => null), + ensureScratchDir: vi.fn(async () => "/scratch/task-1"), + track: vi.fn(), + } as unknown as ITaskCreationHost; + const sessionService = { + markTaskCreationInFlight: vi.fn(), + connectToTask: vi.fn(connectToTask), + disconnectFromTask: vi.fn(), + } as unknown as SessionService; + const effects = { + onWorkspaceCreated: vi.fn(), + onCreateSuccess: vi.fn(), + onCreateRolledBack: vi.fn(), + }; + const service = new TaskService( + host, + sessionService, + effects as unknown as TaskCreationEffects, + {} as PiRunner, + rootLogger, + ); + return { service, api, effects, task }; + } + + const input = { + content: "do the thing", + workspaceMode: "local" as const, + allowNoRepo: true, + }; + + it("fires onCreateRolledBack when a step fails after onTaskReady", async () => { + const { service, api, effects, task } = makeHarness({ + connectToTask: () => { + throw new Error("agent server down"); + }, + }); + const onTaskReady = vi.fn(); + + const result = await service.createTask(input, onTaskReady); + + expect(result.success).toBe(false); + if (result.success) throw new Error("expected agent_session failure"); + expect(result.failedStep).toBe("agent_session"); + // The UI got the task before the failure, and the saga deleted it — the + // effect is what lets hosts pull it back out of their caches. + expect(onTaskReady).toHaveBeenCalledOnce(); + expect(api.deleteTask).toHaveBeenCalledWith(task.id); + expect(effects.onCreateRolledBack).toHaveBeenCalledExactlyOnceWith(task); + }); + + it("does not fire onCreateRolledBack when creation succeeds", async () => { + const { service, effects } = makeHarness({ connectToTask: () => {} }); + + const result = await service.createTask(input, vi.fn()); + + expect(result.success).toBe(true); + expect(effects.onCreateRolledBack).not.toHaveBeenCalled(); + }); + + it("does not fire onCreateRolledBack when creation fails before surfacing", async () => { + const effects = { + onWorkspaceCreated: vi.fn(), + onCreateSuccess: vi.fn(), + onCreateRolledBack: vi.fn(), + }; + const host = { + getAuthenticatedClient: vi.fn(async () => ({ + createTask: vi.fn().mockRejectedValue(new Error("api down")), + deleteTask: vi.fn(), + })), + detectRepo: vi.fn(async () => null), + track: vi.fn(), + } as unknown as ITaskCreationHost; + const service = new TaskService( + host, + { markTaskCreationInFlight: vi.fn() } as unknown as SessionService, + effects as unknown as TaskCreationEffects, + {} as PiRunner, + rootLogger, + ); + const onTaskReady = vi.fn(); + + const result = await service.createTask(input, onTaskReady); + + expect(result.success).toBe(false); + if (result.success) throw new Error("expected task_creation failure"); + expect(result.failedStep).toBe("task_creation"); + expect(onTaskReady).not.toHaveBeenCalled(); + expect(effects.onCreateRolledBack).not.toHaveBeenCalled(); + }); +}); + describe("TaskService.createTask validation", () => { it("rejects an input with neither content nor a taskDescription", async () => { const result = await makeService().createTask({ diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index 52c76216a6..128963c232 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -150,6 +150,10 @@ export class TaskService { } } + // onTaskReady fires mid-saga, before the remaining steps run — so a later + // failure rolls back (deletes) a task the UI has already been handed and + // surfaced into feeds and caches. Track it so hosts can undo that. + let surfacedTask: Task | undefined; const creator = new TaskCreationSaga( { posthogClient, @@ -157,7 +161,10 @@ export class TaskService { sessionService: this.sessionService, piRunner: this.piRunner, track: (event, props) => this.host.track(event, props), - onTaskReady, + onTaskReady: (output) => { + surfacedTask = output.task; + onTaskReady?.(output); + }, }, this.log, ); @@ -166,6 +173,12 @@ export class TaskService { if (result.success) { this.effects.onWorkspaceCreated(result.data); this.effects.onCreateSuccess(result.data, input); + } else if (surfacedTask) { + this.log.info("Task creation rolled back after task was surfaced", { + taskId: surfacedTask.id, + failedStep: result.failedStep, + }); + this.effects.onCreateRolledBack(surfacedTask); } return result; diff --git a/packages/ui/src/features/canvas/components/ChannelHomeComposer.tsx b/packages/ui/src/features/canvas/components/ChannelHomeComposer.tsx index 1d11900152..5e4842d91a 100644 --- a/packages/ui/src/features/canvas/components/ChannelHomeComposer.tsx +++ b/packages/ui/src/features/canvas/components/ChannelHomeComposer.tsx @@ -333,9 +333,12 @@ export const ChannelHomeComposer = forwardRef< const created = await handleSubmit(content); if (!created) { - // Creation failed — onTaskCreated never fired, so this id is still - // queued. Pull its row and give the full structured prompt (chips and - // attachments, not just flattened text) back so the user can retry. + // Creation failed. If it failed before the task surfaced, this id is + // still queued — pull its "Starting…" row (no-ops otherwise: a failure + // after onTaskCreated fired already retired the id, and the spliced + // card is rolled back centrally via onCreateRolledBack). Give the full + // structured prompt (chips and attachments, not just flattened text) + // back so the user can retry. pendingIdsRef.current = pendingIdsRef.current.filter((p) => p !== id); onPendingEnd(id); editor.insertEditorContent(content); diff --git a/packages/ui/src/features/canvas/components/WebsiteNewTask.test.tsx b/packages/ui/src/features/canvas/components/WebsiteNewTask.test.tsx index 5a4954b069..472396f2de 100644 --- a/packages/ui/src/features/canvas/components/WebsiteNewTask.test.tsx +++ b/packages/ui/src/features/canvas/components/WebsiteNewTask.test.tsx @@ -12,11 +12,14 @@ if (typeof globalThis.ResizeObserver === "undefined") { } as unknown as typeof ResizeObserver; } -const { track, useFolderInstructions, taskInputProps } = vi.hoisted(() => ({ - track: vi.fn(), - useFolderInstructions: vi.fn(), - taskInputProps: vi.fn(), -})); +const { track, useFolderInstructions, taskInputProps, router, openTaskInput } = + vi.hoisted(() => ({ + track: vi.fn(), + useFolderInstructions: vi.fn(), + taskInputProps: vi.fn(), + router: { state: { location: { pathname: "/website/chan-1/new" } } }, + openTaskInput: vi.fn(), + })); // TaskInput is a huge hook-heavy component; stub it down to just the surface // this test cares about — a button that fires onContextChipClick when wired. @@ -59,8 +62,10 @@ vi.mock("@posthog/ui/shell/analytics", () => ({ track })); vi.mock("@tanstack/react-query", () => ({ useQueryClient: () => ({ setQueryData: vi.fn() }), })); +vi.mock("@posthog/ui/router/useOpenTask", () => ({ openTaskInput })); vi.mock("@tanstack/react-router", () => ({ useNavigate: () => vi.fn(), + useRouter: () => router, // The view reads the matched route so it can pass new-task prefill through. useRouterState: ({ select, @@ -147,3 +152,41 @@ describe("WebsiteNewTask context panel", () => { expect(screen.getByRole("button", { name: "context-chip" })).toBeDisabled(); }); }); + +// onTaskCreated navigates onto the task before creation fully settles; when a +// later step fails the task is rolled back and that page is a dead end. +describe("WebsiteNewTask creation rollback", () => { + beforeEach(() => { + useFolderInstructions.mockReturnValue({ data: undefined }); + taskInputProps.mockReset(); + openTaskInput.mockReset(); + }); + + function failCreation() { + const props = taskInputProps.mock.lastCall?.[0] as { + onTaskCreationFailed: (task: { id: string }, promptText: string) => void; + }; + props.onTaskCreationFailed({ id: "task-9" }, "ship the thing"); + } + + it("returns the user to the composer when parked on the dead task", () => { + renderNewTask(); + router.state.location.pathname = "/website/chan-1/tasks/task-9"; + + failCreation(); + + expect(openTaskInput).toHaveBeenCalledExactlyOnceWith({ + channelId: "chan-1", + initialPrompt: "ship the thing", + }); + }); + + it("leaves navigation alone when the user already moved elsewhere", () => { + renderNewTask(); + router.state.location.pathname = "/website/chan-2"; + + failCreation(); + + expect(openTaskInput).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/ui/src/features/canvas/components/WebsiteNewTask.tsx b/packages/ui/src/features/canvas/components/WebsiteNewTask.tsx index c5161e64ca..8be9bfd243 100644 --- a/packages/ui/src/features/canvas/components/WebsiteNewTask.tsx +++ b/packages/ui/src/features/canvas/components/WebsiteNewTask.tsx @@ -14,10 +14,11 @@ import { useSetHeaderContent } from "@posthog/ui/hooks/useSetHeaderContent"; import { ResizableSidebar } from "@posthog/ui/primitives/ResizableSidebar"; import { toast } from "@posthog/ui/primitives/toast"; import { useAppView } from "@posthog/ui/router/useAppView"; +import { openTaskInput } from "@posthog/ui/router/useOpenTask"; import { track } from "@posthog/ui/shell/analytics"; import { Flex } from "@radix-ui/themes"; import { useQueryClient } from "@tanstack/react-query"; -import { useNavigate } from "@tanstack/react-router"; +import { useNavigate, useRouter } from "@tanstack/react-router"; import { useCallback, useMemo, useState } from "react"; // A channel's "New task" view. Reuses /code's TaskInput, but routes the created @@ -27,6 +28,7 @@ import { useCallback, useMemo, useState } from "react"; export function WebsiteNewTask({ channelId }: { channelId: string }) { const spacesLayout = useChannelsLayout(); const navigate = useNavigate(); + const router = useRouter(); const view = useAppView(); const queryClient = useQueryClient(); const { fileTask } = useChannelTaskMutations(); @@ -109,11 +111,24 @@ export function WebsiteNewTask({ channelId }: { channelId: string }) { [channelId, fileTask, navigate, queryClient], ); + // onTaskCreated above navigates onto the task optimistically; if creation + // then fails and rolls the task back, that page is a dead end. Return the + // user to this composer with their prompt — unless they've already moved on. + const onTaskCreationFailed = useCallback( + (task: Task, promptText: string) => { + if (router.state.location.pathname.endsWith(`/tasks/${task.id}`)) { + openTaskInput({ channelId, initialPrompt: promptText }); + } + }, + [router, channelId], + ); + return (
({ + resolveService: vi.fn(), +})); + +import { resolveService } from "@posthog/di/container"; +import { HOST_TRPC_CLIENT } from "@posthog/host-router/client"; +import { IMPERATIVE_QUERY_CLIENT } from "../../shell/queryClient"; +import { navigationTaskBinder } from "./taskBinderImpl"; + +function makeHost(overrides?: { + workspaces?: Record; + folders?: unknown[]; +}) { + return { + workspace: { + getAll: { query: vi.fn(async () => overrides?.workspaces ?? {}) }, + create: { mutate: vi.fn(async () => undefined) }, + ensureScratchDir: { + mutate: vi.fn(async () => ({ path: "/scratch/t1" })), + }, + }, + folders: { + getFolders: { query: vi.fn(async () => overrides?.folders ?? []) }, + getRepositoryByRemoteUrl: { query: vi.fn(async () => null) }, + addFolder: { mutate: vi.fn(async () => undefined) }, + }, + }; +} + +const queryClientFake = { invalidateQueries: vi.fn(async () => undefined) }; + +function wire(host: ReturnType) { + vi.mocked(resolveService).mockImplementation((token) => { + if (token === HOST_TRPC_CLIENT) return host; + if (token === IMPERATIVE_QUERY_CLIENT) return queryClientFake; + // Anything else (e.g. the shell logger's HOST_LOGGER) is unbound in this + // harness; throwing matches an unbound container and the logger no-ops. + throw new Error(`unbound token in test: ${String(token)}`); + }); +} + +describe("navigationTaskBinder.ensureWorkspaceForTask", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("provisions a scratch dir for a repo-less local task with no workspace", async () => { + const host = makeHost(); + wire(host); + const task = { id: "t1", repository: null } as unknown as Task; + + await navigationTaskBinder.ensureWorkspaceForTask(task); + + expect(host.workspace.ensureScratchDir.mutate).toHaveBeenCalledWith({ + taskId: "t1", + }); + expect(host.workspace.create.mutate).not.toHaveBeenCalled(); + expect(queryClientFake.invalidateQueries).toHaveBeenCalled(); + }); + + it("leaves an existing scratch workspace alone", async () => { + const host = makeHost({ + workspaces: { t1: { taskId: "t1", isScratch: true } }, + }); + wire(host); + const task = { id: "t1", repository: null } as unknown as Task; + + await navigationTaskBinder.ensureWorkspaceForTask(task); + + expect(host.workspace.ensureScratchDir.mutate).not.toHaveBeenCalled(); + }); + + it("does not provision scratch for a repo task with no resolvable directory", async () => { + const host = makeHost(); + wire(host); + const task = { + id: "t1", + repository: "posthog/code", + } as unknown as Task; + + await navigationTaskBinder.ensureWorkspaceForTask(task); + + expect(host.workspace.ensureScratchDir.mutate).not.toHaveBeenCalled(); + expect(host.workspace.create.mutate).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/ui/src/features/navigation/taskBinderImpl.ts b/packages/ui/src/features/navigation/taskBinderImpl.ts index 11f80f1c60..6ff885f2ff 100644 --- a/packages/ui/src/features/navigation/taskBinderImpl.ts +++ b/packages/ui/src/features/navigation/taskBinderImpl.ts @@ -127,6 +127,19 @@ export const navigationTaskBinder: NavigationTaskBinder = { mode: "cloud", }); invalidateWorkspaces(); + } else if (!repoKey) { + // Repo-less task with no scratch dir on this machine (creation failed + // partway leaving the task row behind, or created elsewhere): provision + // the scratch dir the way creation does so the task view resolves a cwd + // instead of asking for a repo folder it never wanted. + try { + await hostClient().workspace.ensureScratchDir.mutate({ + taskId: task.id, + }); + invalidateWorkspaces(); + } catch (error) { + log.error("Failed to provision scratch dir on task open:", error); + } } else { log.warn("No directory resolved on task open, workspace not created", { taskId: task.id, diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index ae2adb4e78..c9f4e386bd 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/packages/ui/src/features/task-detail/components/TaskInput.tsx @@ -111,6 +111,12 @@ import { type WorkspaceMode, WorkspaceModeSelect } from "./WorkspaceModeSelect"; interface TaskInputProps { sessionId?: string; onTaskCreated?: (task: Task) => void; + /** + * Creation failed after onTaskCreated fired and the task was rolled back + * (see useTaskCreation's onTaskCreationFailed). Callers whose onTaskCreated + * navigates use this to route the user off the now-dead task. + */ + onTaskCreationFailed?: (task: Task, promptText: string) => void; initialPrompt?: string; initialPromptKey?: string; initialCloudRepository?: string; @@ -159,6 +165,7 @@ interface TaskInputProps { export function TaskInput({ sessionId = "task-input", onTaskCreated, + onTaskCreationFailed, initialPrompt, initialPromptKey, initialCloudRepository, @@ -946,6 +953,7 @@ export function TaskInput({ fastMode: runtime === "pi" ? undefined : currentFastMode, onTaskCreated, onTaskCreatedEffect: handleAutoresearchTaskCreated, + onTaskCreationFailed, environmentId: selectedEnvironment, sandboxEnvironmentId: effectiveWorkspaceMode === "cloud" && selectedCloudEnvId diff --git a/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx b/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx index 1cc7283356..307e067b06 100644 --- a/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx +++ b/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx @@ -17,7 +17,10 @@ import { useRestoreTask } from "../../suspension/useRestoreTask"; import { useSuspendedTaskIds } from "../../suspension/useSuspendedTaskIds"; import { useBranchMismatchDialog } from "../../workspace/useBranchMismatchDialog"; import { useWorkspaceLoaded } from "../../workspace/useWorkspace"; -import { useCreateWorkspace } from "../../workspace/useWorkspaceMutations"; +import { + useCreateWorkspace, + useEnsureScratchWorkspace, +} from "../../workspace/useWorkspaceMutations"; import { BranchMismatchDialog } from "../BranchMismatchDialog"; import { WorkspaceSetupPrompt } from "./WorkspaceSetupPrompt"; @@ -102,6 +105,24 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { requestFocus(taskId); }, [taskId, requestFocus]); + // A repo-less task's workspace is a synthetic scratch dir. When it's + // missing (creation failed partway leaving the task row behind, opened on + // another machine, dir cleaned up), the folder picker is the wrong ask — + // provision the scratch dir the way creation does and let the session + // resolve it as cwd. + const scratchSetup = useEnsureScratchWorkspace(); + const needsScratchWorkspace = + !!localWorkspaces && + !repoKey && + !repoPath && + !isCloud && + !isSuspended && + isWorkspaceLoaded; + const ensureScratch = scratchSetup.ensure; + useEffect(() => { + if (needsScratchWorkspace) ensureScratch(taskId); + }, [needsScratchWorkspace, ensureScratch, taskId]); + // Once a retry provisions a workspace, drop the stale provisioning error so // the guard in ensureWorkspaceForTask stops firing and the retry prompt hides. useEffect(() => { @@ -140,13 +161,15 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { ); } + // Repo task with no local mapping → folder picker. Repo-less task → the + // scratch provisioning above; the picker is only its last-resort fallback. if ( localWorkspaces && !repoPath && !isCloud && !isSuspended && isWorkspaceLoaded && - !hasDirectoryMapping && + (repoKey ? !hasDirectoryMapping : scratchSetup.isError) && !isCreatingWorkspace ) { return ( diff --git a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts index 9bbfaf88cd..760d3631ad 100644 --- a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts +++ b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts @@ -109,6 +109,15 @@ interface UseTaskCreationOptions { * this does not suppress the pending-task view. */ onTaskCreatedEffect?: (task: Task) => void; + /** + * Called when creation failed after onTaskCreated already fired — the saga + * surfaced the task, then a later step failed and rolled it back (the + * server-side delete is best-effort, so the task row may survive as a + * run-less draft). Query caches are cleaned centrally + * (TaskCreationEffects.onCreateRolledBack); use this to undo caller-side + * effects like navigation onto the dead task. + */ + onTaskCreationFailed?: (task: Task, promptText: string) => void; } interface UseTaskCreationReturn { @@ -197,6 +206,7 @@ export function useTaskCreation({ allowNoRepo, onTaskCreated, onTaskCreatedEffect, + onTaskCreationFailed, }: UseTaskCreationOptions): UseTaskCreationReturn { const [isCreatingTask, setIsCreatingTask] = useState(false); const hostClient = useHostTRPCClient(); @@ -344,6 +354,7 @@ export function useTaskCreation({ } let createdTaskId: string | undefined; + let surfacedTask: Task | undefined; try { if (!contentOverride) { @@ -428,6 +439,7 @@ export function useTaskCreation({ clearTaskInputReportAssociation(); } createdTaskId = output.task.id; + surfacedTask = output.task; if (pendingTaskKey) { pendingTaskPromptStoreApi.move(pendingTaskKey, output.task.id); } @@ -523,6 +535,10 @@ export function useTaskCreation({ } openTaskInput({ initialPrompt: plainPromptText }); } + if (surfacedTask) { + titleAttachmentStoreApi.clear(surfacedTask.id); + onTaskCreationFailed?.(surfacedTask, plainPromptText); + } } return result.success; } catch (error) { @@ -576,6 +592,7 @@ export function useTaskCreation({ invalidateTasks, onTaskCreated, onTaskCreatedEffect, + onTaskCreationFailed, hostClient, trpc, queryClient, diff --git a/packages/ui/src/features/task-detail/taskCreationEffectsImpl.test.ts b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.test.ts new file mode 100644 index 0000000000..c8a9c4ca6e --- /dev/null +++ b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.test.ts @@ -0,0 +1,62 @@ +import { QueryClient } from "@tanstack/react-query"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@posthog/di/container", () => ({ + resolveService: vi.fn(), +})); + +import { resolveService } from "@posthog/di/container"; +import type { Task } from "@posthog/shared/domain-types"; +import { channelFeedQueryRoot } from "../canvas/hooks/useChannelFeed"; +import { taskKeys } from "../tasks/taskKeys"; +import { taskCreationEffects } from "./taskCreationEffectsImpl"; + +const deadTask = { id: "task-dead", title: "Rolled back" } as Task; +const otherTask = { id: "task-live", title: "Unrelated" } as Task; + +describe("taskCreationEffects.onCreateRolledBack", () => { + let client: QueryClient; + + beforeEach(() => { + client = new QueryClient(); + vi.mocked(resolveService).mockReturnValue(client); + }); + + it("removes the rolled-back task from list, feed, and detail caches", () => { + client.setQueryData(taskKeys.list(), [deadTask, otherTask]); + client.setQueryData( + [...channelFeedQueryRoot, "channel-1"], + [deadTask, otherTask], + ); + client.setQueryData(taskKeys.detail(deadTask.id), deadTask); + + taskCreationEffects.onCreateRolledBack(deadTask); + + expect(client.getQueryData(taskKeys.list())).toEqual([otherTask]); + expect(client.getQueryData([...channelFeedQueryRoot, "channel-1"])).toEqual( + [otherTask], + ); + expect(client.getQueryData(taskKeys.detail(deadTask.id))).toBeUndefined(); + }); + + it("marks list and feed queries stale so the server copy settles", () => { + client.setQueryData(taskKeys.list(), [deadTask]); + client.setQueryData([...channelFeedQueryRoot, "channel-1"], [deadTask]); + + taskCreationEffects.onCreateRolledBack(deadTask); + + expect(client.getQueryState(taskKeys.list())?.isInvalidated).toBe(true); + expect( + client.getQueryState([...channelFeedQueryRoot, "channel-1"]) + ?.isInvalidated, + ).toBe(true); + }); + + it("leaves other tasks' detail caches alone", () => { + client.setQueryData(taskKeys.detail(otherTask.id), otherTask); + + taskCreationEffects.onCreateRolledBack(deadTask); + + expect(client.getQueryData(taskKeys.detail(otherTask.id))).toBe(otherTask); + }); +}); diff --git a/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts index 5340024b44..4e18403436 100644 --- a/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts +++ b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts @@ -1,4 +1,5 @@ import type { TaskCreationEffects } from "@posthog/core/task-detail/taskCreationEffects"; +import { removeTaskFromList } from "@posthog/core/tasks/taskDelete"; import { resolveService } from "@posthog/di/container"; import type { TaskCreationInput, @@ -10,6 +11,7 @@ import { IMPERATIVE_QUERY_CLIENT, type ImperativeQueryClient, } from "../../shell/queryClient"; +import { channelFeedQueryRoot } from "../canvas/hooks/useChannelFeed"; import { useDraftStore } from "../message-editor/draftStore"; import { useSettingsStore } from "../settings/settingsStore"; import { taskKeys } from "../tasks/taskKeys"; @@ -44,6 +46,27 @@ export const taskCreationEffects: TaskCreationEffects = { void client.invalidateQueries({ queryKey: taskKeys.allSummaries() }); }, + onCreateRolledBack(task: Task): void { + const client = queryClient(); + // Ready-time invalidations may still have refetches in flight that were + // answered while the task existed; cancel them so a late response can't + // splice the dead task back in after the removal below. + void client.cancelQueries({ queryKey: taskKeys.lists() }); + void client.cancelQueries({ queryKey: channelFeedQueryRoot }); + client.setQueriesData({ queryKey: taskKeys.lists() }, (tasks) => + removeTaskFromList(tasks, task.id), + ); + client.setQueriesData({ queryKey: channelFeedQueryRoot }, (tasks) => + removeTaskFromList(tasks, task.id), + ); + // Dropped (not just invalidated): a lingering detail entry makes the task + // route treat the server's 404 as non-authoritative and render the dead + // task instead of redirecting. + client.removeQueries({ queryKey: taskKeys.detail(task.id) }); + void client.invalidateQueries({ queryKey: taskKeys.lists() }); + void client.invalidateQueries({ queryKey: channelFeedQueryRoot }); + }, + onCreateSuccess(output: TaskCreationOutput, input?: TaskCreationInput): void { if (!input) return; diff --git a/packages/ui/src/features/workspace/useWorkspaceMutations.test.tsx b/packages/ui/src/features/workspace/useWorkspaceMutations.test.tsx index 678385dece..2f8b6b639b 100644 --- a/packages/ui/src/features/workspace/useWorkspaceMutations.test.tsx +++ b/packages/ui/src/features/workspace/useWorkspaceMutations.test.tsx @@ -1,13 +1,16 @@ import type { Workspace, WorkspaceInfo } from "@posthog/shared"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, renderHook } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import type { ReactNode } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; const WORKSPACE_QUERY_KEY = ["workspace", "getAll"]; const WORKTREES_FILTER = { queryKey: ["worktrees", "/repo"] }; -const createFn = vi.hoisted(() => vi.fn()); +const { createFn, scratchFn } = vi.hoisted(() => ({ + createFn: vi.fn(), + scratchFn: vi.fn(), +})); vi.mock("@posthog/host-router/react", () => ({ useHostTRPC: () => ({ @@ -30,11 +33,20 @@ vi.mock("@posthog/host-router/react", () => ({ ...options, }), }, + ensureScratchDir: { + mutationOptions: (options: Record) => ({ + mutationFn: (input: unknown) => scratchFn(input), + ...options, + }), + }, }, }), })); -import { useEnsureWorkspace } from "./useWorkspaceMutations"; +import { + useEnsureScratchWorkspace, + useEnsureWorkspace, +} from "./useWorkspaceMutations"; const created = { taskId: "t1", mode: "worktree" } as unknown as WorkspaceInfo; @@ -88,4 +100,54 @@ describe("useWorkspaceMutations", () => { }); expect(invalidateSpy).toHaveBeenCalledWith(WORKTREES_FILTER); }); + + it("useEnsureScratchWorkspace provisions the scratch dir and invalidates workspaces", async () => { + scratchFn.mockResolvedValue({ path: "/scratch/t1" }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + const { result } = renderHook(() => useEnsureScratchWorkspace(), { + wrapper, + }); + + await act(async () => { + result.current.ensure("t1"); + }); + + expect(scratchFn).toHaveBeenCalledExactlyOnceWith({ taskId: "t1" }); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: WORKSPACE_QUERY_KEY, + }); + }); + + it("useEnsureScratchWorkspace fires once — later ensure() calls no-op", async () => { + scratchFn.mockResolvedValue({ path: "/scratch/t1" }); + const { result } = renderHook(() => useEnsureScratchWorkspace(), { + wrapper, + }); + + await act(async () => { + result.current.ensure("t1"); + }); + await act(async () => { + result.current.ensure("t1"); + }); + + expect(scratchFn).toHaveBeenCalledTimes(1); + }); + + it("useEnsureScratchWorkspace reports failure without refiring", async () => { + scratchFn.mockRejectedValue(new Error("disk full")); + const { result } = renderHook(() => useEnsureScratchWorkspace(), { + wrapper, + }); + + await act(async () => { + result.current.ensure("t1"); + }); + await act(async () => { + result.current.ensure("t1"); + }); + + expect(scratchFn).toHaveBeenCalledTimes(1); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); }); diff --git a/packages/ui/src/features/workspace/useWorkspaceMutations.ts b/packages/ui/src/features/workspace/useWorkspaceMutations.ts index d8a385bc13..c131467aea 100644 --- a/packages/ui/src/features/workspace/useWorkspaceMutations.ts +++ b/packages/ui/src/features/workspace/useWorkspaceMutations.ts @@ -5,7 +5,7 @@ import { import { useHostTRPC } from "@posthog/host-router/react"; import type { Workspace, WorkspaceMode } from "@posthog/shared"; import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { useCallback } from "react"; +import { useCallback, useRef } from "react"; function useInvalidateWorkspaceCaches() { const trpc = useHostTRPC(); @@ -45,6 +45,40 @@ export function useCreateWorkspace(): { isPending: boolean } { return { isPending: mutation.isPending }; } +/** + * Provision the synthetic scratch workspace for a repo-less task (what task + * creation does via the saga). ensure() fires at most once per hook instance, + * so a mount effect can call it unconditionally without looping on failure. + */ +export function useEnsureScratchWorkspace(): { + ensure: (taskId: string) => void; + isError: boolean; +} { + const trpc = useHostTRPC(); + const invalidateCaches = useInvalidateWorkspaceCaches(); + const firedRef = useRef(false); + + const mutation = useMutation( + trpc.workspace.ensureScratchDir.mutationOptions({ + onSuccess: () => { + void invalidateCaches(); + }, + }), + ); + + const { mutate } = mutation; + const ensure = useCallback( + (taskId: string) => { + if (firedRef.current) return; + firedRef.current = true; + mutate({ taskId }); + }, + [mutate], + ); + + return { ensure, isError: mutation.isError }; +} + export function useDeleteWorkspace(): { isPending: boolean } { const trpc = useHostTRPC(); const invalidateCaches = useInvalidateWorkspaceCaches();