fix(sidebar): keep created subtasks during hydration (#1626)

* fix(sidebar): keep created subtasks during hydration

* fix(sidebar): preserve placeholder subtasks during hydration

* fix(sidebar): mark websocket placeholder snapshots

apps/web/components/task/task-switcher-direct-subtasks.test.tsx

@@ -0,0 +1,53 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { cleanup, render, screen } from "@testing-library/react";
+import { StateProvider } from "@/components/state-provider";
+import { ToastProvider } from "@/components/toast-provider";
+import { TaskSwitcher, type TaskSwitcherItem } from "./task-switcher";
+import type { GroupedSidebarList } from "@/lib/sidebar/apply-view";
+
+afterEach(() => cleanup());
+
+function Providers({ children }: { children: React.ReactNode }) {
+  return (
+    <StateProvider>
+      <ToastProvider>{children}</ToastProvider>
+    </StateProvider>
+  );
+}
+
+function task(id: string, parentTaskId?: string): TaskSwitcherItem {
+  return { id, title: id, state: "IN_PROGRESS", parentTaskId };
+}
+
+function groupedWithDirectSubtasks(childCount: number): GroupedSidebarList {
+  const parent = task("Parent");
+  const children = Array.from({ length: childCount }, (_, index) =>
+    task(`Child ${index + 1}`, parent.id),
+  );
+  return {
+    groups: [{ key: "__all__", label: "All", tasks: [parent] }],
+    subTasksByParentId: new Map([[parent.id, children]]),
+  };
+}
+
+describe("TaskSwitcher direct subtasks", () => {
+  it("renders a parent with 17 direct subtasks", () => {
+    const { container } = render(
+      <Providers>
+        <TaskSwitcher
+          grouped={groupedWithDirectSubtasks(17)}
+          activeTaskId={null}
+          selectedTaskId={null}
+          onSelectTask={vi.fn()}
+          onToggleSubtasks={vi.fn()}
+          collapsedSubtaskParentIds={[]}
+        />
+      </Providers>,
+    );
+
+    expect(container.querySelectorAll("[data-testid='sortable-task-block']")).toHaveLength(18);
+    expect(screen.queryByText("Parent")).not.toBeNull();
+    expect(screen.queryByText("Child 1")).not.toBeNull();
+    expect(screen.queryByText("Child 17")).not.toBeNull();
+  });
+});

apps/web/hooks/domains/kanban/use-all-workflow-snapshots-inflight.test.ts

@@ -0,0 +1,173 @@
+import { describe, expect, it, vi } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+
+type Workflow = { id: string; workspaceId: string; name: string };
+type SnapshotTask = {
+  id: string;
+  workflowStepId: string;
+  title: string;
+  position: number;
+  state: "IN_PROGRESS";
+  parentTaskId?: string;
+};
+type MockState = {
+  connection: { status: string };
+  workflows: { items: Workflow[] };
+  kanbanMulti: {
+    snapshots: Record<
+      string,
+      {
+        workflowId: string;
+        workflowName: string;
+        steps: [];
+        tasks: SnapshotTask[];
+        isPlaceholder?: boolean;
+      }
+    >;
+    isLoading: boolean;
+  };
+  clearKanbanMulti: ReturnType<typeof vi.fn>;
+  setKanbanMultiLoading: ReturnType<typeof vi.fn>;
+  setWorkflowSnapshot: ReturnType<typeof vi.fn>;
+};
+
+const WORKFLOW_ID = "wf-A";
+const WORKSPACE_ID = "ws-A";
+const STEP_ID = "step-1";
+const PARENT_TASK_ID = "parent-task";
+
+const mocks = vi.hoisted(() => ({
+  clearKanbanMulti: vi.fn(),
+  fetchWorkflowSnapshot: vi.fn(),
+  setKanbanMultiLoading: vi.fn(),
+  setWorkflowSnapshot: vi.fn(),
+  state: undefined as MockState | undefined,
+}));
+
+vi.mock("@/components/state-provider", () => ({
+  useAppStore: (selector: (state: MockState) => unknown) => selector(mocks.state!),
+  useAppStoreApi: () => ({ getState: () => mocks.state! }),
+}));
+
+vi.mock("@/lib/api", () => ({
+  fetchWorkflowSnapshot: (...args: unknown[]) => mocks.fetchWorkflowSnapshot(...args),
+}));
+
+import { useAllWorkflowSnapshots } from "./use-all-workflow-snapshots";
+
+function resetState() {
+  vi.clearAllMocks();
+  mocks.state = {
+    connection: { status: "connected" },
+    workflows: { items: [{ id: WORKFLOW_ID, workspaceId: WORKSPACE_ID, name: "A" }] },
+    kanbanMulti: { snapshots: {}, isLoading: false },
+    clearKanbanMulti: mocks.clearKanbanMulti,
+    setKanbanMultiLoading: mocks.setKanbanMultiLoading,
+    setWorkflowSnapshot: mocks.setWorkflowSnapshot,
+  };
+}
+
+function staleSnapshotResponse() {
+  return {
+    steps: [{ id: STEP_ID, name: "Doing", color: null, position: 0 }],
+    tasks: [],
+  };
+}
+
+function setLightweightSnapshot(task: SnapshotTask) {
+  mocks.state!.kanbanMulti.snapshots[WORKFLOW_ID] = {
+    workflowId: WORKFLOW_ID,
+    workflowName: "A",
+    steps: [],
+    tasks: [task],
+    isPlaceholder: true,
+  };
+}
+
+describe("useAllWorkflowSnapshots lightweight snapshots", () => {
+  it("does not treat a lightweight websocket snapshot as boot-hydrated", async () => {
+    resetState();
+    mocks.fetchWorkflowSnapshot.mockResolvedValueOnce(staleSnapshotResponse());
+    setLightweightSnapshot({
+      id: "child-before-hydration",
+      workflowStepId: STEP_ID,
+      title: "Child before hydration",
+      position: 0,
+      state: "IN_PROGRESS",
+    });
+
+    renderHook(() => useAllWorkflowSnapshots(WORKSPACE_ID));
+
+    await waitFor(() => expect(mocks.fetchWorkflowSnapshot).toHaveBeenCalledTimes(1));
+  });
+
+  it("preserves tasks already present in a lightweight snapshot before fetch starts", async () => {
+    resetState();
+    mocks.fetchWorkflowSnapshot.mockResolvedValueOnce(staleSnapshotResponse());
+    setLightweightSnapshot({
+      id: "child-before-fetch",
+      workflowStepId: STEP_ID,
+      title: "Child before fetch",
+      position: 0,
+      state: "IN_PROGRESS",
+      parentTaskId: PARENT_TASK_ID,
+    });
+
+    renderHook(() => useAllWorkflowSnapshots(WORKSPACE_ID));
+
+    await waitFor(() =>
+      expect(mocks.setWorkflowSnapshot).toHaveBeenCalledWith(
+        WORKFLOW_ID,
+        expect.objectContaining({
+          tasks: [
+            expect.objectContaining({
+              id: "child-before-fetch",
+              parentTaskId: PARENT_TASK_ID,
+            }),
+          ],
+        }),
+      ),
+    );
+  });
+});
+
+describe("useAllWorkflowSnapshots in-flight websocket tasks", () => {
+  it("preserves tasks created while the workflow snapshot fetch is in flight", async () => {
+    resetState();
+    let resolveFetch: (value: unknown) => void = () => {};
+    mocks.fetchWorkflowSnapshot.mockReturnValueOnce(
+      new Promise((resolve) => {
+        resolveFetch = resolve;
+      }),
+    );
+
+    renderHook(() => useAllWorkflowSnapshots(WORKSPACE_ID));
+    await waitFor(() =>
+      expect(mocks.fetchWorkflowSnapshot).toHaveBeenCalledWith(WORKFLOW_ID, expect.anything()),
+    );
+
+    setLightweightSnapshot({
+      id: "child-created-during-fetch",
+      workflowStepId: STEP_ID,
+      title: "Child created during fetch",
+      position: 0,
+      state: "IN_PROGRESS",
+      parentTaskId: PARENT_TASK_ID,
+    });
+    resolveFetch(staleSnapshotResponse());
+
+    await waitFor(() =>
+      expect(mocks.setWorkflowSnapshot).toHaveBeenCalledWith(
+        WORKFLOW_ID,
+        expect.objectContaining({
+          tasks: [
+            expect.objectContaining({
+              id: "child-created-during-fetch",
+              parentTaskId: PARENT_TASK_ID,
+            }),
+          ],
+        }),
+      ),
+    );
+  });
+});

apps/web/hooks/domains/kanban/use-all-workflow-snapshots.ts

@@ -2,7 +2,7 @@ import { useEffect, useRef, type MutableRefObject } from "react";
 import { fetchWorkflowSnapshot } from "@/lib/api";
 import { useAppStore, useAppStoreApi } from "@/components/state-provider";
 import { toKanbanTask } from "@/lib/kanban/map-task";
-import type { KanbanState } from "@/lib/state/slices/kanban/types";
+import type { KanbanState, WorkflowSnapshotData } from "@/lib/state/slices/kanban/types";
 import type { Task } from "@/lib/types/http";
 import type { StoreApi } from "zustand";
 import type { AppState } from "@/lib/state/store";
@@ -10,6 +10,10 @@ import type { AppState } from "@/lib/state/store";
 type KanbanTask = KanbanState["tasks"][number];
 type Workflow = { id: string; name: string };
 
+function isBootHydratedSnapshot(snapshot: WorkflowSnapshotData | undefined): boolean {
+  return !!snapshot && snapshot.isPlaceholder !== true;
+}
+
 async function fetchAndWriteSnapshot(
   wf: Workflow,
   store: StoreApi<AppState>,
@@ -17,6 +21,8 @@ async function fetchAndWriteSnapshot(
   myGen: number,
 ): Promise<void> {
   try {
+    const snapshotAtFetchStart = store.getState().kanbanMulti.snapshots[wf.id];
+    const taskIdsAtFetchStart = new Set((snapshotAtFetchStart?.tasks ?? []).map((t) => t.id));
     const snapshot = await fetchWorkflowSnapshot(wf.id, { cache: "no-store" });
     if (fetchGenRef.current !== myGen) return;
 
@@ -51,12 +57,20 @@ async function fetchAndWriteSnapshot(
         return mapped;
       })
       .filter((t): t is KanbanTask => t !== null);
+    const snapshotTaskIds = new Set(tasks.map((t) => t.id));
+    const preserveExistingPlaceholderTasks = snapshotAtFetchStart?.isPlaceholder === true;
+    const inFlightCreatedTasks = (existingSnapshot?.tasks ?? []).filter(
+      (task) =>
+        (preserveExistingPlaceholderTasks || !taskIdsAtFetchStart.has(task.id)) &&
+        !snapshotTaskIds.has(task.id) &&
+        stepIds.has(task.workflowStepId),
+    );
 
     store.getState().setWorkflowSnapshot(wf.id, {
       workflowId: wf.id,
       workflowName: wf.name,
       steps,
-      tasks,
+      tasks: [...tasks, ...inFlightCreatedTasks],
     });
   } catch (err) {
     console.error(
@@ -112,7 +126,9 @@ export function useAllWorkflowSnapshots(workspaceId: string | null) {
     }
     if (
       lastFetchedRef.current === "" &&
-      workspaceWorkflows.every((wf) => store.getState().kanbanMulti.snapshots[wf.id])
+      workspaceWorkflows.every((wf) =>
+        isBootHydratedSnapshot(store.getState().kanbanMulti.snapshots[wf.id]),
+      )
     ) {
       lastFetchedRef.current = key;
       return;

apps/web/lib/sidebar/apply-view-direct-subtasks.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from "vitest";
+import type { TaskSwitcherItem } from "@/components/task/task-switcher";
+import { DEFAULT_VIEW } from "@/lib/state/slices/ui/sidebar-view-builtins";
+import { applyView, countGroupTasks } from "./apply-view";
+
+function task(overrides: Partial<TaskSwitcherItem>): TaskSwitcherItem {
+  return {
+    id: overrides.id ?? "task",
+    title: overrides.title ?? "Task",
+    state: "IN_PROGRESS",
+    ...overrides,
+  };
+}
+
+describe("applyView direct subtasks", () => {
+  it("keeps all 17 direct subtasks in the default all-tasks view", () => {
+    const parent = task({ id: "parent", title: "Parent" });
+    const children = Array.from({ length: 17 }, (_, index) =>
+      task({
+        id: `child-${index + 1}`,
+        title: `Child ${index + 1}`,
+        parentTaskId: parent.id,
+      }),
+    );
+
+    const grouped = applyView([parent, ...children], DEFAULT_VIEW);
+    const rootIds = grouped.groups.flatMap((group) => group.tasks.map((item) => item.id));
+    const childIds = grouped.subTasksByParentId.get(parent.id)?.map((item) => item.id);
+
+    expect(rootIds).toEqual([parent.id]);
+    expect(childIds).toHaveLength(17);
+    expect(countGroupTasks(grouped.groups[0].tasks, grouped.subTasksByParentId)).toBe(18);
+  });
+});

apps/web/lib/state/slices/kanban/types.ts

@@ -79,6 +79,7 @@ export type WorkflowSnapshotData = {
   workflowName: string;
   steps: KanbanState["steps"];
   tasks: KanbanState["tasks"];
+  isPlaceholder?: boolean;
 };
 
 export type KanbanMultiState = {

apps/web/lib/ws/handlers/tasks-created-snapshot.test.ts

@@ -0,0 +1,105 @@
+import { describe, expect, it, vi } from "vitest";
+import type { StoreApi } from "zustand";
+import type { AppState } from "@/lib/state/store";
+import { registerTasksHandlers } from "./tasks";
+
+type Listener = (state: AppState) => void;
+
+function makeStore(initial: Partial<AppState> = {}) {
+  let state = {
+    kanban: { workflowId: "wf1", steps: [], tasks: [] },
+    kanbanMulti: { snapshots: {}, isLoading: false },
+    workflows: {
+      activeId: "wf1",
+      items: [
+        { id: "wf1", workspaceId: "ws1", name: "Active Flow" },
+        { id: "wf2", workspaceId: "ws1", name: "Sibling Flow" },
+      ],
+    },
+    tasks: {
+      activeTaskId: null,
+      activeSessionId: null,
+      pinnedSessionId: null,
+      lastSessionByTaskId: {},
+    },
+    taskSessionsByTask: { itemsByTaskId: {}, loadedByTaskId: {}, loadingByTaskId: {} },
+    environmentIdBySessionId: {},
+    removeTaskFromSidebarPrefs: vi.fn(),
+    setTaskDeletedNotification: vi.fn(),
+    ...initial,
+  } as unknown as AppState;
+
+  const listeners = new Set<Listener>();
+  return {
+    getState: () => state,
+    setState: (updater: AppState | ((s: AppState) => AppState)) => {
+      const next =
+        typeof updater === "function" ? (updater as (s: AppState) => AppState)(state) : updater;
+      state = { ...state, ...next };
+      for (const listener of listeners) listener(state);
+    },
+    subscribe: (listener: Listener) => {
+      listeners.add(listener);
+      return () => listeners.delete(listener);
+    },
+    destroy: vi.fn(),
+    getInitialState: vi.fn(),
+  } as unknown as StoreApi<AppState> & { getState: () => AppState };
+}
+
+function makeCreatedMessage(payload: Record<string, unknown>) {
+  return {
+    id: "msg-1",
+    type: "notification" as const,
+    action: "task.created" as const,
+    payload,
+  } as Parameters<NonNullable<ReturnType<typeof registerTasksHandlers>["task.created"]>>[0];
+}
+
+describe("task.created sidebar snapshots", () => {
+  it("keeps a created subtask visible when its workflow snapshot has not loaded yet", () => {
+    const store = makeStore();
+    const handlers = registerTasksHandlers(store);
+
+    handlers["task.created"]!(
+      makeCreatedMessage({
+        task_id: "child-before-snapshot",
+        workflow_id: "wf2",
+        workflow_step_id: "step2",
+        title: "Child before snapshot",
+        state: "IN_PROGRESS",
+        parent_id: "parent-task",
+        is_ephemeral: false,
+      }),
+    );
+
+    const snapshot = store.getState().kanbanMulti.snapshots.wf2;
+    expect(snapshot.workflowName).toBe("Sibling Flow");
+    expect(snapshot.steps).toEqual([]);
+    expect(snapshot.isPlaceholder).toBe(true);
+    expect(snapshot.tasks.map((task) => task.id)).toEqual(["child-before-snapshot"]);
+    expect(snapshot.tasks[0].parentTaskId).toBe("parent-task");
+  });
+
+  it("keeps the task even when workflow metadata has not loaded yet", () => {
+    const store = makeStore({
+      workflows: { activeId: null, items: [] },
+    } as Partial<AppState>);
+    const handlers = registerTasksHandlers(store);
+
+    handlers["task.created"]!(
+      makeCreatedMessage({
+        task_id: "child-before-workflows",
+        workflow_id: "wf-late",
+        workflow_step_id: "step-late",
+        title: "Child before workflows",
+        state: "IN_PROGRESS",
+        is_ephemeral: false,
+      }),
+    );
+
+    const snapshot = store.getState().kanbanMulti.snapshots["wf-late"];
+    expect(snapshot.workflowName).toBe("wf-late");
+    expect(snapshot.tasks.map((task) => task.id)).toEqual(["child-before-workflows"]);
+  });
+});

apps/web/lib/ws/handlers/tasks.ts

@@ -64,7 +64,26 @@ function upsertMultiTask(
   payload: TaskEventPayload,
 ): AppState {
   const snapshot = state.kanbanMulti.snapshots[workflowId];
-  if (!snapshot) return state;
+  if (!snapshot) {
+    const workflowName =
+      state.workflows?.items.find((item) => item.id === workflowId)?.name ?? workflowId;
+    return {
+      ...state,
+      kanbanMulti: {
+        ...state.kanbanMulti,
+        snapshots: {
+          ...state.kanbanMulti.snapshots,
+          [workflowId]: {
+            workflowId,
+            workflowName,
+            steps: [],
+            tasks: upsertTask([], task, payload),
+            isPlaceholder: true,
+          },
+        },
+      },
+    };
+  }
   return {
     ...state,
     kanbanMulti: {
@@ -108,9 +127,7 @@ function upsertTaskInBothKanbans(
     };
   }
 
-  if (state.kanbanMulti.snapshots[wfId]) {
-    next = upsertMultiTask(next, wfId, nextTask, payload);
-  }
+  next = upsertMultiTask(next, wfId, nextTask, payload);
 
   return next;
 }