Start extracting Tree component

2026-03-11 17:46:41 -05:00 · 2026-03-08 16:37:25 -07:00
parent 6e11894f79
commit 6534421733
19 changed files with 344 additions and 151 deletions
--- a/apps/yaak-client/components/EnvironmentEditDialog.tsx
+++ b/apps/yaak-client/components/EnvironmentEditDialog.tsx
@@ -15,12 +15,15 @@ import { resolvedModelName } from '../lib/resolvedModelName';
 import { showColorPicker } from '../lib/showColorPicker';
 import { Banner } from './core/Banner';
 import type { ContextMenuProps, DropdownItem } from './core/Dropdown';
+import { ContextMenu } from './core/Dropdown';
 import { Icon } from '@yaakapp-internal/ui';
 import { IconButton } from './core/IconButton';
 import { IconTooltip } from './core/IconTooltip';
 import { InlineCode } from './core/InlineCode';
 import type { PairEditorHandle } from './core/PairEditor';
 import { SplitLayout } from './core/SplitLayout';
+import { atomFamily } from 'jotai/utils';
+import { atomWithKVStorage } from '../lib/atoms/atomWithKVStorage';
 import type { TreeNode } from './core/tree/common';
 import type { TreeHandle, TreeProps } from './core/tree/Tree';
 import { Tree } from './core/tree/Tree';
@@ -28,6 +31,11 @@ import { EnvironmentColorIndicator } from './EnvironmentColorIndicator';
 import { EnvironmentEditor } from './EnvironmentEditor';
 import { EnvironmentSharableTooltip } from './EnvironmentSharableTooltip';

+const collapsedFamily = atomFamily((treeId: string) => {
+  const key = ['env_collapsed', treeId ?? 'n/a'];
+  return atomWithKVStorage<Record<string, boolean>>(key, {});
+});
+
 interface Props {
  initialEnvironmentId: string | null;
  setRef?: (ref: PairEditorHandle | null) => void;
@@ -292,6 +300,13 @@ function EnvironmentEditDialogSidebar({
    [setSelectedEnvironmentId],
  );

+  const renderContextMenuFn = useCallback<NonNullable<TreeProps<TreeModel>['renderContextMenu']>>(
+    ({ items, position, onClose }) => (
+      <ContextMenu items={items as DropdownItem[]} triggerPosition={position} onClose={onClose} />
+    ),
+    [],
+  );
+
  const tree = useAtomValue(treeAtom);
  return (
    <aside className="x-theme-sidebar h-full w-full min-w-0 grid overflow-y-auto border-r border-border-subtle ">
@@ -300,10 +315,12 @@ function EnvironmentEditDialogSidebar({
          <Tree
            ref={treeRef}
            treeId={treeId}
+            collapsedAtom={collapsedFamily(treeId)}
            className="px-2 pb-10"
            hotkeys={hotkeys}
            root={tree}
            getContextMenu={getContextMenu}
+            renderContextMenu={renderContextMenuFn}
            onDragEnd={handleDragEnd}
            getItemKey={(i) => `${i.id}::${i.name}`}
            ItemLeftSlotInner={ItemLeftSlotInner}
--- a/apps/yaak-client/components/Sidebar.tsx
+++ b/apps/yaak-client/components/Sidebar.tsx
@@ -23,7 +23,7 @@ import {
 } from '@yaakapp-internal/models';
 import classNames from 'classnames';
 import { atom, useAtomValue } from 'jotai';
-import { selectAtom } from 'jotai/utils';
+import { atomFamily, selectAtom } from 'jotai/utils';
 import { memo, useCallback, useEffect, useMemo, useRef } from 'react';
 import { moveToWorkspace } from '../commands/moveToWorkspace';
 import { openFolderSettings } from '../commands/openFolderSettings';
@@ -48,7 +48,7 @@ import { resolvedModelName } from '../lib/resolvedModelName';
 import { isSidebarFocused } from '../lib/scopes';
 import { navigateToRequestOrFolderOrWorkspace } from '../lib/setWorkspaceSearchParams';
 import type { ContextMenuProps, DropdownItem } from './core/Dropdown';
-import { Dropdown } from './core/Dropdown';
+import { ContextMenu, Dropdown } from './core/Dropdown';
 import type { FieldDef } from './core/Editor/filter/extension';
 import { filter } from './core/Editor/filter/extension';
 import { evaluate, parseQuery } from './core/Editor/filter/query';
@@ -59,13 +59,19 @@ import { IconButton } from './core/IconButton';
 import { InlineCode } from './core/InlineCode';
 import type { InputHandle } from './core/Input';
 import { Input } from './core/Input';
-import { collapsedFamily, isSelectedFamily, selectedIdsFamily } from './core/tree/atoms';
+import { isSelectedFamily, selectedIdsFamily } from './core/tree/atoms';
 import type { TreeNode } from './core/tree/common';
 import type { TreeHandle, TreeProps } from './core/tree/Tree';
 import { Tree } from './core/tree/Tree';
 import type { TreeItemProps } from './core/tree/TreeItem';
+import { atomWithKVStorage } from '../lib/atoms/atomWithKVStorage';
 import { GitDropdown } from './git/GitDropdown';

+const collapsedFamily = atomFamily((treeId: string) => {
+  const key = ['sidebar_collapsed', treeId ?? 'n/a'];
+  return atomWithKVStorage<Record<string, boolean>>(key, {});
+});
+
 type SidebarModel = Workspace | Folder | HttpRequest | GrpcRequest | WebsocketRequest;
 function isSidebarLeafModel(m: AnyModel): boolean {
  const modelMap: Record<Exclude<SidebarModel['model'], 'workspace'>, null> = {
@@ -456,6 +462,13 @@ function Sidebar({ className }: { className?: string }) {
    [actions],
  );

+  const renderContextMenuFn = useCallback<NonNullable<TreeProps<SidebarModel>['renderContextMenu']>>(
+    ({ items, position, onClose }) => (
+      <ContextMenu items={items as DropdownItem[]} triggerPosition={position} onClose={onClose} />
+    ),
+    [],
+  );
+
  const hotkeys = useMemo<TreeProps<SidebarModel>['hotkeys']>(() => ({ actions }), [actions]);

  // Use a language compartment for the filter so we can reconfigure it when the autocompletion changes
@@ -575,11 +588,13 @@ function Sidebar({ className }: { className?: string }) {
          ref={handleTreeRefInit}
          root={tree}
          treeId={treeId}
+          collapsedAtom={collapsedFamily(treeId)}
          hotkeys={hotkeys}
          getItemKey={getItemKey}
          ItemInner={SidebarInnerItem}
          ItemLeftSlotInner={SidebarLeftSlot}
          getContextMenu={getContextMenu}
+          renderContextMenu={renderContextMenuFn}
          onActivate={handleActivate}
          getEditOptions={getEditOptions}
          className="pl-2 pr-3 pt-2 pb-2"
--- a/apps/yaak-client/components/core/PairEditor.tsx
+++ b/apps/yaak-client/components/core/PairEditor.tsx
@@ -17,9 +17,8 @@ import { useRandomKey } from '../../hooks/useRandomKey';
 import { useToggle } from '../../hooks/useToggle';
 import { languageFromContentType } from '../../lib/contentType';
 import { showDialog } from '../../lib/dialog';
-import { computeSideForDragMove } from '../../lib/dnd';
+import { computeSideForDragMove, DropMarker } from '@yaakapp-internal/ui';
 import { showPrompt } from '../../lib/prompt';
-import { DropMarker } from '../DropMarker';
 import { SelectFile } from '../SelectFile';
 import { Button } from './Button';
 import { Checkbox } from './Checkbox';
--- a/apps/yaak-client/components/core/Tabs/Tabs.tsx
+++ b/apps/yaak-client/components/core/Tabs/Tabs.tsx
@@ -22,8 +22,7 @@ import {
  useState,
 } from 'react';
 import { useKeyValue } from '../../../hooks/useKeyValue';
-import { computeSideForDragMove } from '../../../lib/dnd';
-import { DropMarker } from '../../DropMarker';
+import { computeSideForDragMove, DropMarker } from '@yaakapp-internal/ui';
 import { ErrorBoundary } from '../../ErrorBoundary';
 import type { ButtonProps } from '../Button';
 import { Button } from '../Button';
--- a/apps/yaak-client/components/core/tree/Tree.tsx
+++ b/apps/yaak-client/components/core/tree/Tree.tsx
@@ -24,19 +24,16 @@ import {
 import { useKey, useKeyPressEvent } from 'react-use';
 import type { HotKeyOptions, HotkeyAction } from '../../../hooks/useHotKey';
 import { useHotKey } from '../../../hooks/useHotKey';
-import { computeSideForDragMove } from '../../../lib/dnd';
-import { jotaiStore } from '../../../lib/jotai';
-import type { ContextMenuProps, DropdownItem } from '../Dropdown';
-import { ContextMenu } from '../Dropdown';
+import { computeSideForDragMove } from '@yaakapp-internal/ui';
+import { useStore } from 'jotai';
 import {
-  collapsedFamily,
  draggingIdsFamily,
  focusIdsFamily,
  hoveredParentFamily,
-  isCollapsedFamily,
  selectedIdsFamily,
 } from './atoms';
-import type { SelectableTreeNode, TreeNode } from './common';
+import { type CollapsedAtom, CollapsedAtomContext } from './context';
+import type { ContextMenuRenderer, JotaiStore, SelectableTreeNode, TreeNode } from './common';
 import { closestVisibleNode, equalSubtree, getSelectedItems, hasAncestor } from './common';
 import { TreeDragOverlay } from './TreeDragOverlay';
 import type { TreeItemClickEvent, TreeItemHandle, TreeItemProps } from './TreeItem';
@@ -50,8 +47,10 @@ const measuring = { droppable: { strategy: MeasuringStrategy.Always } };
 export interface TreeProps<T extends { id: string }> {
  root: TreeNode<T>;
  treeId: string;
+  collapsedAtom: CollapsedAtom;
  getItemKey: (item: T) => string;
-  getContextMenu?: (items: T[]) => ContextMenuProps['items'] | Promise<ContextMenuProps['items']>;
+  getContextMenu?: (items: T[]) => unknown[] | Promise<unknown[]>;
+  renderContextMenu?: ContextMenuRenderer;
  ItemInner: ComponentType<{ treeId: string; item: T }>;
  ItemLeftSlotInner?: ComponentType<{ treeId: string; item: T }>;
  ItemRightSlot?: ComponentType<{ treeId: string; item: T }>;
@@ -80,12 +79,14 @@ export interface TreeHandle {
 function TreeInner<T extends { id: string }>(
  {
    className,
+    collapsedAtom,
    getContextMenu,
    getEditOptions,
    getItemKey,
    hotkeys,
    onActivate,
    onDragEnd,
+    renderContextMenu,
    ItemInner,
    ItemLeftSlotInner,
    ItemRightSlot,
@@ -94,10 +95,11 @@ function TreeInner<T extends { id: string }>(
  }: TreeProps<T>,
  ref: Ref<TreeHandle>,
 ) {
+  const store = useStore();
  const treeRef = useRef<HTMLDivElement>(null);
  const selectableItems = useSelectableItems(root);
  const [showContextMenu, setShowContextMenu] = useState<{
-    items: DropdownItem[];
+    items: unknown[];
    x: number;
    y: number;
  } | null>(null);
@@ -113,11 +115,11 @@ function TreeInner<T extends { id: string }>(
  // Select the first item on first render
  // biome-ignore lint/correctness/useExhaustiveDependencies: Only used for initial render
  useEffect(() => {
-    const ids = jotaiStore.get(selectedIdsFamily(treeId));
+    const ids = store.get(selectedIdsFamily(treeId));
    const fallback = selectableItems[0];
    if (ids.length === 0 && fallback != null) {
-      jotaiStore.set(selectedIdsFamily(treeId), [fallback.node.item.id]);
-      jotaiStore.set(focusIdsFamily(treeId), {
+      store.set(selectedIdsFamily(treeId), [fallback.node.item.id]);
+      store.set(focusIdsFamily(treeId), {
        anchorId: fallback.node.item.id,
        lastId: fallback.node.item.id,
      });
@@ -145,7 +147,7 @@ function TreeInner<T extends { id: string }>(
  }, []);

  const ensureTabbableItem = useCallback(() => {
-    const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+    const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
    const lastSelectedItem = selectableItems.find(
      (i) => i.node.item.id === lastSelectedId && !i.node.hidden,
    );
@@ -156,23 +158,23 @@ function TreeInner<T extends { id: string }>(
      const firstItem = firstLeafItem ?? selectableItems.find((i) => !i.node.hidden);
      if (firstItem != null) {
        const id = firstItem.node.item.id;
-        jotaiStore.set(selectedIdsFamily(treeId), [id]);
-        jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
+        store.set(selectedIdsFamily(treeId), [id]);
+        store.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
      }
      return;
    }

-    const closest = closestVisibleNode(treeId, lastSelectedItem.node);
+    const closest = closestVisibleNode(store, collapsedAtom, lastSelectedItem.node);
    if (closest != null) {
      const id = closest.item.id;
-      jotaiStore.set(selectedIdsFamily(treeId), [id]);
-      jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
+      store.set(selectedIdsFamily(treeId), [id]);
+      store.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
    }
  }, [selectableItems, treeId]);

  // Ensure there's always a tabbable item after collapsed state changes
  useEffect(() => {
-    const unsub = jotaiStore.sub(collapsedFamily(treeId), ensureTabbableItem);
+    const unsub = store.sub(collapsedAtom, ensureTabbableItem);
    return unsub;
  }, [ensureTabbableItem, treeId]);

@@ -187,7 +189,7 @@ function TreeInner<T extends { id: string }>(

  const setSelected = useCallback(
    (ids: string[], focus: boolean) => {
-      jotaiStore.set(selectedIdsFamily(treeId), ids);
+      store.set(selectedIdsFamily(treeId), ids);
      // TODO: Figure out a better way than timeout
      if (!focus) return;
      setTimeout(tryFocus, 50);
@@ -202,18 +204,18 @@ function TreeInner<T extends { id: string }>(
      hasFocus: hasFocus,
      renameItem: (id) => treeItemRefs.current[id]?.rename(),
      selectItem: (id, focus) => {
-        if (jotaiStore.get(selectedIdsFamily(treeId)).includes(id)) {
+        if (store.get(selectedIdsFamily(treeId)).includes(id)) {
          // Already selected
          return;
        }
-        jotaiStore.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
+        store.set(focusIdsFamily(treeId), { anchorId: id, lastId: id });
        setSelected([id], focus === true);
      },
      showContextMenu: async () => {
        if (getContextMenu == null) return;
-        const items = getSelectedItems(treeId, selectableItems);
+        const items = getSelectedItems(store, treeId, selectableItems);
        const menuItems = await getContextMenu(items);
-        const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+        const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
        const rect = lastSelectedId ? treeItemRefs.current[lastSelectedId]?.rect() : null;
        if (rect == null) return;
        setShowContextMenu({ items: menuItems, x: rect.x, y: rect.y });
@@ -227,7 +229,7 @@ function TreeInner<T extends { id: string }>(
  const handleGetContextMenu = useMemo(() => {
    if (getContextMenu == null) return;
    return (item: T) => {
-      const items = getSelectedItems(treeId, selectableItems);
+      const items = getSelectedItems(store, treeId, selectableItems);
      const isSelected = items.find((i) => i.id === item.id);
      if (isSelected) {
        // If right-clicked an item that was in the multiple-selection, use the entire selection
@@ -236,22 +238,22 @@ function TreeInner<T extends { id: string }>(
      // If right-clicked an item that was NOT in the multiple-selection, just use that one
      // Also update the selection with it
      setSelected([item.id], false);
-      jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id }));
+      store.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id }));
      return getContextMenu([item]);
    };
  }, [getContextMenu, selectableItems, setSelected, treeId]);

  const handleSelect = useCallback<NonNullable<TreeItemProps<T>['onClick']>>(
    (item, { shiftKey, metaKey, ctrlKey }) => {
-      const anchorSelectedId = jotaiStore.get(focusIdsFamily(treeId)).anchorId;
+      const anchorSelectedId = store.get(focusIdsFamily(treeId)).anchorId;
      const selectedIdsAtom = selectedIdsFamily(treeId);
-      const selectedIds = jotaiStore.get(selectedIdsAtom);
+      const selectedIds = store.get(selectedIdsAtom);

      // Mark the item as the last one selected
-      jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id }));
+      store.set(focusIdsFamily(treeId), (prev) => ({ ...prev, lastId: item.id }));

      if (shiftKey) {
-        const validSelectableItems = getValidSelectableItems(treeId, selectableItems);
+        const validSelectableItems = getValidSelectableItems(store, collapsedAtom, selectableItems);
        const anchorIndex = validSelectableItems.findIndex(
          (i) => i.node.item.id === anchorSelectedId,
        );
@@ -260,7 +262,7 @@ function TreeInner<T extends { id: string }>(
        // Nothing was selected yet, so just select this item
        if (selectedIds.length === 0 || anchorIndex === -1 || currIndex === -1) {
          setSelected([item.id], true);
-          jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id }));
+          store.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id }));
          return;
        }

@@ -293,7 +295,7 @@ function TreeInner<T extends { id: string }>(
      } else {
        // Select single
        setSelected([item.id], true);
-        jotaiStore.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id }));
+        store.set(focusIdsFamily(treeId), (prev) => ({ ...prev, anchorId: item.id }));
      }
    },
    [selectableItems, setSelected, treeId],
@@ -313,8 +315,8 @@ function TreeInner<T extends { id: string }>(

  const selectPrevItem = useCallback(
    (e: TreeItemClickEvent) => {
-      const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
-      const validSelectableItems = getValidSelectableItems(treeId, selectableItems);
+      const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
+      const validSelectableItems = getValidSelectableItems(store, collapsedAtom, selectableItems);
      const index = validSelectableItems.findIndex((i) => i.node.item.id === lastSelectedId);
      const item = validSelectableItems[index - 1];
      if (item != null) {
@@ -326,8 +328,8 @@ function TreeInner<T extends { id: string }>(

  const selectNextItem = useCallback(
    (e: TreeItemClickEvent) => {
-      const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
-      const validSelectableItems = getValidSelectableItems(treeId, selectableItems);
+      const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
+      const validSelectableItems = getValidSelectableItems(store, collapsedAtom, selectableItems);
      const index = validSelectableItems.findIndex((i) => i.node.item.id === lastSelectedId);
      const item = validSelectableItems[index + 1];
      if (item != null) {
@@ -339,7 +341,7 @@ function TreeInner<T extends { id: string }>(

  const selectParentItem = useCallback(
    (e: TreeItemClickEvent) => {
-      const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+      const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
      const lastSelectedItem =
        selectableItems.find((i) => i.node.item.id === lastSelectedId)?.node ?? null;
      if (lastSelectedItem?.parent != null) {
@@ -378,8 +380,8 @@ function TreeInner<T extends { id: string }>(
      if (!isTreeFocused()) return;
      e.preventDefault();

-      const collapsed = jotaiStore.get(collapsedFamily(treeId));
-      const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+      const collapsed = store.get(collapsedAtom);
+      const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
      const lastSelectedItem = selectableItems.find((i) => i.node.item.id === lastSelectedId);

      if (
@@ -387,7 +389,7 @@ function TreeInner<T extends { id: string }>(
        lastSelectedItem?.node.children != null &&
        collapsed[lastSelectedItem.node.item.id] === true
      ) {
-        jotaiStore.set(isCollapsedFamily({ treeId, itemId: lastSelectedId }), false);
+        store.set(collapsedAtom, { ...collapsed, [lastSelectedId]: false });
      } else {
        selectNextItem(e);
      }
@@ -404,8 +406,8 @@ function TreeInner<T extends { id: string }>(
      if (!isTreeFocused()) return;
      e.preventDefault();

-      const collapsed = jotaiStore.get(collapsedFamily(treeId));
-      const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+      const collapsed = store.get(collapsedAtom);
+      const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
      const lastSelectedItem = selectableItems.find((i) => i.node.item.id === lastSelectedId);

      if (
@@ -413,7 +415,7 @@ function TreeInner<T extends { id: string }>(
        lastSelectedItem?.node.children != null &&
        collapsed[lastSelectedItem.node.item.id] !== true
      ) {
-        jotaiStore.set(isCollapsedFamily({ treeId, itemId: lastSelectedId }), true);
+        store.set(collapsedAtom, { ...collapsed, [lastSelectedId]: true });
      } else {
        selectParentItem(e);
      }
@@ -425,7 +427,7 @@ function TreeInner<T extends { id: string }>(
  useKeyPressEvent('Escape', async () => {
    if (!treeRef.current?.contains(document.activeElement)) return;
    clearDragState();
-    const lastSelectedId = jotaiStore.get(focusIdsFamily(treeId)).lastId;
+    const lastSelectedId = store.get(focusIdsFamily(treeId)).lastId;
    if (lastSelectedId == null) return;
    setSelected([lastSelectedId], false);
  });
@@ -435,7 +437,7 @@ function TreeInner<T extends { id: string }>(
      const over = e.over;
      if (!over) {
        // Clear the drop indicator when hovering outside the tree
-        jotaiStore.set(hoveredParentFamily(treeId), {
+        store.set(hoveredParentFamily(treeId), {
          parentId: null,
          parentDepth: null,
          childIndex: null,
@@ -452,7 +454,7 @@ function TreeInner<T extends { id: string }>(
      // Root is anything past the end of the list, so set it to the end
      const hoveringRoot = over.id === root.item.id;
      if (hoveringRoot) {
-        jotaiStore.set(hoveredParentFamily(treeId), {
+        store.set(hoveredParentFamily(treeId), {
          parentId: root.item.id,
          parentDepth: root.depth,
          index: selectableItems.length,
@@ -466,7 +468,7 @@ function TreeInner<T extends { id: string }>(
        return;
      }

-      const draggingItems = jotaiStore.get(draggingIdsFamily(treeId));
+      const draggingItems = store.get(draggingIdsFamily(treeId));
      for (const id of draggingItems) {
        const item = selectableItems.find((i) => i.node.item.id === id)?.node ?? null;
        if (item == null) {
@@ -499,7 +501,7 @@ function TreeInner<T extends { id: string }>(
      const parentDepth = hoveredParent?.depth ?? null;
      const index = hoveredIndex;
      const childIndex = hoveredChildIndex;
-      const existing = jotaiStore.get(hoveredParentFamily(treeId));
+      const existing = store.get(hoveredParentFamily(treeId));
      if (
        !(
          parentId === existing.parentId &&
@@ -508,7 +510,7 @@ function TreeInner<T extends { id: string }>(
          childIndex === existing.childIndex
        )
      ) {
-        jotaiStore.set(hoveredParentFamily(treeId), {
+        store.set(hoveredParentFamily(treeId), {
          parentId,
          parentDepth,
          index,
@@ -521,12 +523,12 @@ function TreeInner<T extends { id: string }>(

  const handleDragStart = useCallback(
    function handleDragStart(e: DragStartEvent) {
-      const selectedItems = getSelectedItems(treeId, selectableItems);
+      const selectedItems = getSelectedItems(store, treeId, selectableItems);
      const isDraggingSelectedItem = selectedItems.find((i) => i.id === e.active.id);

      // If we started dragging an already-selected item, we'll use that
      if (isDraggingSelectedItem) {
-        jotaiStore.set(
+        store.set(
          draggingIdsFamily(treeId),
          selectedItems.map((i) => i.id),
        );
@@ -534,7 +536,7 @@ function TreeInner<T extends { id: string }>(
        // If we started dragging a non-selected item, only drag that item
        const activeItem = selectableItems.find((i) => i.node.item.id === e.active.id)?.node.item;
        if (activeItem != null) {
-          jotaiStore.set(draggingIdsFamily(treeId), [activeItem.id]);
+          store.set(draggingIdsFamily(treeId), [activeItem.id]);
          // Also update selection to just be this one
          handleSelect(activeItem, {
            shiftKey: false,
@@ -548,13 +550,13 @@ function TreeInner<T extends { id: string }>(
  );

  const clearDragState = useCallback(() => {
-    jotaiStore.set(hoveredParentFamily(treeId), {
+    store.set(hoveredParentFamily(treeId), {
      parentId: null,
      parentDepth: null,
      index: null,
      childIndex: null,
    });
-    jotaiStore.set(draggingIdsFamily(treeId), []);
+    store.set(draggingIdsFamily(treeId), []);
  }, [treeId]);

  const handleDragEnd = useCallback(
@@ -564,8 +566,8 @@ function TreeInner<T extends { id: string }>(
        index: hoveredIndex,
        parentId: hoveredParentId,
        childIndex: hoveredChildIndex,
-      } = jotaiStore.get(hoveredParentFamily(treeId));
-      const draggingItems = jotaiStore.get(draggingIdsFamily(treeId));
+      } = store.get(hoveredParentFamily(treeId));
+      const draggingItems = store.get(draggingIdsFamily(treeId));
      clearDragState();

      // Dropped outside the tree?
@@ -624,6 +626,7 @@ function TreeInner<T extends { id: string }>(
  > = {
    getItemKey,
    getContextMenu: handleGetContextMenu,
+    renderContextMenu,
    onClick: handleClick,
    getEditOptions,
    ItemInner,
@@ -646,15 +649,14 @@ function TreeInner<T extends { id: string }>(
  const sensors = useSensors(useSensor(PointerSensor, { activationConstraint: { distance: 6 } }));

  return (
-    <>
+    <CollapsedAtomContext.Provider value={collapsedAtom}>
      <TreeHotKeys treeId={treeId} hotkeys={hotkeys} selectableItems={selectableItems} />
-      {showContextMenu && (
-        <ContextMenu
-          items={showContextMenu.items}
-          triggerPosition={showContextMenu}
-          onClose={handleCloseContextMenu}
-        />
-      )}
+      {showContextMenu &&
+        renderContextMenu?.({
+          items: showContextMenu.items,
+          position: showContextMenu,
+          onClose: handleCloseContextMenu,
+        })}
      <DndContext
        sensors={sensors}
        collisionDetection={pointerWithin}
@@ -707,7 +709,7 @@ function TreeInner<T extends { id: string }>(
          getItemKey={getItemKey}
        />
      </DndContext>
-    </>
+    </CollapsedAtomContext.Provider>
  );
 }

@@ -757,10 +759,11 @@ function TreeHotKey<T extends { id: string }>({
  enable,
  ...options
 }: TreeHotKeyProps<T>) {
+  const store = useStore();
  useHotKey(
    action,
    () => {
-      onDone(getSelectedItems(treeId, selectableItems));
+      onDone(getSelectedItems(store, treeId, selectableItems));
    },
    {
      ...options,
@@ -802,10 +805,11 @@ function TreeHotKeys<T extends { id: string }>({
 }

 function getValidSelectableItems<T extends { id: string }>(
-  treeId: string,
+  store: JotaiStore,
+  collapsedAtom: CollapsedAtom,
  selectableItems: SelectableTreeNode<T>[],
 ) {
-  const collapsed = jotaiStore.get(collapsedFamily(treeId));
+  const collapsed = store.get(collapsedAtom);
  return selectableItems.filter((i) => {
    if (i.node.hidden) return false;
    let p = i.node.parent;
--- a/apps/yaak-client/components/core/tree/TreeDropMarker.tsx
+++ b/apps/yaak-client/components/core/tree/TreeDropMarker.tsx
@@ -1,9 +1,10 @@
 import classNames from 'classnames';
 import { useAtomValue } from 'jotai';
 import { memo } from 'react';
-import { DropMarker } from '../../DropMarker';
-import { hoveredParentDepthFamily, isCollapsedFamily, isIndexHoveredFamily } from './atoms';
+import { DropMarker } from '@yaakapp-internal/ui';
+import { hoveredParentDepthFamily, isIndexHoveredFamily } from './atoms';
 import type { TreeNode } from './common';
+import { useIsCollapsed } from './context';

 export const TreeDropMarker = memo(function TreeDropMarker<T extends { id: string }>({
  className,
@@ -16,10 +17,9 @@ export const TreeDropMarker = memo(function TreeDropMarker<T extends { id: strin
  node: TreeNode<T> | null;
  className?: string;
 }) {
-  const itemId = node?.item.id;
  const isHovered = useAtomValue(isIndexHoveredFamily({ treeId, index }));
  const parentDepth = useAtomValue(hoveredParentDepthFamily(treeId));
-  const collapsed = useAtomValue(isCollapsedFamily({ treeId, itemId }));
+  const collapsed = useIsCollapsed(node?.item.id);

  // Only show if we're hovering over this index
  if (!isHovered) return null;
--- a/apps/yaak-client/components/core/tree/TreeItem.tsx
+++ b/apps/yaak-client/components/core/tree/TreeItem.tsx
@@ -1,8 +1,7 @@
 import type { DragMoveEvent } from '@dnd-kit/core';
 import { useDndContext, useDndMonitor, useDraggable, useDroppable } from '@dnd-kit/core';
 import classNames from 'classnames';
-import { useAtomValue } from 'jotai';
-import { selectAtom } from 'jotai/utils';
+import { useAtomValue, useStore } from 'jotai';
 import type {
  MouseEvent,
  PointerEvent,
@@ -10,12 +9,11 @@ import type {
  KeyboardEvent as ReactKeyboardEvent,
 } from 'react';
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { computeSideForDragMove } from '../../../lib/dnd';
-import { jotaiStore } from '../../../lib/jotai';
-import type { ContextMenuProps, DropdownItem } from '../Dropdown';
-import { ContextMenu } from '../Dropdown';
+import { computeSideForDragMove } from '@yaakapp-internal/ui';
+import type { ContextMenuRenderer } from './common';
 import { Icon } from '@yaakapp-internal/ui';
-import { collapsedFamily, isCollapsedFamily, isLastFocusedFamily, isSelectedFamily } from './atoms';
+import { isLastFocusedFamily, isSelectedFamily } from './atoms';
+import { useCollapsedAtom, useIsAncestorCollapsed, useIsCollapsed, useSetCollapsed } from './context';
 import type { TreeNode } from './common';
 import { getNodeKey } from './common';
 import type { TreeProps } from './Tree';
@@ -29,12 +27,12 @@ export interface TreeItemClickEvent {

 export type TreeItemProps<T extends { id: string }> = Pick<
  TreeProps<T>,
-  'ItemInner' | 'ItemLeftSlotInner' | 'ItemRightSlot' | 'treeId' | 'getEditOptions' | 'getItemKey'
+  'ItemInner' | 'ItemLeftSlotInner' | 'ItemRightSlot' | 'treeId' | 'getEditOptions' | 'getItemKey' | 'renderContextMenu'
 > & {
  node: TreeNode<T>;
  className?: string;
  onClick?: (item: T, e: TreeItemClickEvent) => void;
-  getContextMenu?: (item: T) => ContextMenuProps['items'] | Promise<ContextMenuProps['items']>;
+  getContextMenu?: (item: T) => unknown[] | Promise<unknown[]>;
  depth: number;
  setRef?: (item: T, n: TreeItemHandle | null) => void;
 };
@@ -56,16 +54,20 @@ function TreeItem_<T extends { id: string }>({
  ItemLeftSlotInner,
  ItemRightSlot,
  getContextMenu,
+  renderContextMenu,
  onClick,
  getEditOptions,
  className,
  depth,
  setRef,
 }: TreeItemProps<T>) {
+  const store = useStore();
  const listItemRef = useRef<HTMLLIElement>(null);
  const draggableRef = useRef<HTMLButtonElement>(null);
+  const collapsedAtom = useCollapsedAtom();
  const isSelected = useAtomValue(isSelectedFamily({ treeId, itemId: node.item.id }));
-  const isCollapsed = useAtomValue(isCollapsedFamily({ treeId, itemId: node.item.id }));
+  const isCollapsed = useIsCollapsed(node.item.id);
+  const setCollapsed = useSetCollapsed(node.item.id);
  const isLastSelected = useAtomValue(isLastFocusedFamily({ treeId, itemId: node.item.id }));
  const [editing, setEditing] = useState<boolean>(false);
  const [dropHover, setDropHover] = useState<null | 'drop' | 'animate'>(null);
@@ -110,19 +112,10 @@ function TreeItem_<T extends { id: string }>({
    return ids;
  }, [node]);

-  const isAncestorCollapsedAtom = useMemo(
-    () =>
-      selectAtom(
-        collapsedFamily(treeId),
-        (collapsed) => ancestorIds.some((id) => collapsed[id]),
-        (a, b) => a === b,
-      ),
-    [ancestorIds, treeId],
-  );
-  const isAncestorCollapsed = useAtomValue(isAncestorCollapsedAtom);
+  const isAncestorCollapsed = useIsAncestorCollapsed(ancestorIds);

  const [showContextMenu, setShowContextMenu] = useState<{
-    items: DropdownItem[];
+    items: unknown[];
    x: number;
    y: number;
  } | null>(null);
@@ -133,8 +126,8 @@ function TreeItem_<T extends { id: string }>({
  );

  const toggleCollapsed = useCallback(() => {
-    jotaiStore.set(isCollapsedFamily({ treeId, itemId: node.item.id }), (prev) => !prev);
-  }, [node.item.id, treeId]);
+    setCollapsed((prev) => !prev);
+  }, [setCollapsed]);

  const handleSubmitNameEdit = useCallback(
    async (el: HTMLInputElement) => {
@@ -207,12 +200,13 @@ function TreeItem_<T extends { id: string }>({
      const side = computeSideForDragMove(node.item.id, e);
      const isFolder = node.children != null;
      const hasChildren = (node.children?.length ?? 0) > 0;
-      const isCollapsed = jotaiStore.get(isCollapsedFamily({ treeId, itemId: node.item.id }));
-      if (isCollapsed && isFolder && hasChildren && side === 'after') {
+      const collapsedMap = store.get(collapsedAtom);
+      const itemCollapsed = !!collapsedMap[node.item.id];
+      if (itemCollapsed && isFolder && hasChildren && side === 'after') {
        setDropHover('animate');
        clearTimeout(startedHoverTimeout.current);
        startedHoverTimeout.current = setTimeout(() => {
-          jotaiStore.set(isCollapsedFamily({ treeId, itemId: node.item.id }), false);
+          store.set(collapsedAtom, { ...store.get(collapsedAtom), [node.item.id]: false });
          clearDropHover();
          // Force re-measure everything because all containers below the folder have been pushed down
          requestAnimationFrame(() => {
@@ -306,13 +300,12 @@ function TreeItem_<T extends { id: string }>({
          'grid grid-cols-[auto_minmax(0,1fr)_auto] gap-x-2 items-center rounded-md',
        )}
      >
-        {showContextMenu && (
-          <ContextMenu
-            items={showContextMenu.items}
-            triggerPosition={showContextMenu}
-            onClose={handleCloseContextMenu}
-          />
-        )}
+        {showContextMenu &&
+          renderContextMenu?.({
+            items: showContextMenu.items,
+            position: showContextMenu,
+            onClose: handleCloseContextMenu,
+          })}
        {node.children != null ? (
          <button
            type="button"
--- a/apps/yaak-client/components/core/tree/TreeItemList.tsx
+++ b/apps/yaak-client/components/core/tree/TreeItemList.tsx
@@ -8,7 +8,7 @@ import { TreeItem } from './TreeItem';

 export type TreeItemListProps<T extends { id: string }> = Pick<
  TreeProps<T>,
-  'ItemInner' | 'ItemLeftSlotInner' | 'ItemRightSlot' | 'treeId' | 'getItemKey' | 'getEditOptions'
+  'ItemInner' | 'ItemLeftSlotInner' | 'ItemRightSlot' | 'treeId' | 'getItemKey' | 'getEditOptions' | 'renderContextMenu'
 > &
  Pick<TreeItemProps<T>, 'onClick' | 'getContextMenu'> & {
    nodes: SelectableTreeNode<T>[];
--- a/apps/yaak-client/components/core/tree/atoms.ts
+++ b/apps/yaak-client/components/core/tree/atoms.ts
@@ -1,6 +1,5 @@
 import { atom } from 'jotai';
 import { atomFamily, selectAtom } from 'jotai/utils';
-import { atomWithKVStorage } from '../../../lib/atoms/atomWithKVStorage';

 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 export const selectedIdsFamily = atomFamily((_treeId: string) => {
@@ -75,28 +74,3 @@ export const hoveredParentDepthFamily = atomFamily((treeId: string) =>
  ),
 );

-export const collapsedFamily = atomFamily((workspaceId: string) => {
-  const key = ['sidebar_collapsed', workspaceId ?? 'n/a'];
-  return atomWithKVStorage<Record<string, boolean>>(key, {});
-});
-
-export const isCollapsedFamily = atomFamily(
-  ({ treeId, itemId = 'n/a' }: { treeId: string; itemId: string | undefined }) =>
-    atom(
-      // --- getter ---
-      (get) => !!get(collapsedFamily(treeId))[itemId],
-
-      // --- setter ---
-      (get, set, next: boolean | ((prev: boolean) => boolean)) => {
-        const a = collapsedFamily(treeId);
-        const prevMap = get(a);
-        const prevValue = !!prevMap[itemId];
-        const value = typeof next === 'function' ? next(prevValue) : next;
-
-        if (value === prevValue) return; // no-op
-
-        set(a, { ...prevMap, [itemId]: value });
-      },
-    ),
-  (a, b) => a.treeId === b.treeId && a.itemId === b.itemId,
-);
--- a/apps/yaak-client/components/core/tree/common.ts
+++ b/apps/yaak-client/components/core/tree/common.ts
@@ -1,5 +1,15 @@
-import { jotaiStore } from '../../../lib/jotai';
-import { collapsedFamily, selectedIdsFamily } from './atoms';
+import type { createStore } from 'jotai';
+import type { ReactNode } from 'react';
+import type { CollapsedAtom } from './context';
+import { selectedIdsFamily } from './atoms';
+
+export type JotaiStore = ReturnType<typeof createStore>;
+
+export type ContextMenuRenderer = (props: {
+  items: unknown[];
+  position: { x: number; y: number };
+  onClose: () => void;
+}) => ReactNode;

 export interface TreeNode<T extends { id: string }> {
  children?: TreeNode<T>[];
@@ -18,10 +28,11 @@ export interface SelectableTreeNode<T extends { id: string }> {
 }

 export function getSelectedItems<T extends { id: string }>(
+  store: JotaiStore,
  treeId: string,
  selectableItems: SelectableTreeNode<T>[],
 ) {
-  const selectedItemIds = jotaiStore.get(selectedIdsFamily(treeId));
+  const selectedItemIds = store.get(selectedIdsFamily(treeId));
  return selectableItems
    .filter((i) => selectedItemIds.includes(i.node.item.id))
    .map((i) => i.node.item);
@@ -60,8 +71,8 @@ export function hasAncestor<T extends { id: string }>(node: TreeNode<T>, ancesto
  return hasAncestor(node.parent, ancestorId);
 }

-export function isVisibleNode<T extends { id: string }>(treeId: string, node: TreeNode<T>) {
-  const collapsed = jotaiStore.get(collapsedFamily(treeId));
+export function isVisibleNode<T extends { id: string }>(store: JotaiStore, collapsedAtom: CollapsedAtom, node: TreeNode<T>) {
+  const collapsed = store.get(collapsedAtom);
  let p = node.parent;
  while (p) {
    if (collapsed[p.item.id]) return false; // any collapsed ancestor hides this node
@@ -71,12 +82,13 @@ export function isVisibleNode<T extends { id: string }>(treeId: string, node: Tr
 }

 export function closestVisibleNode<T extends { id: string }>(
-  treeId: string,
+  store: JotaiStore,
+  collapsedAtom: CollapsedAtom,
  node: TreeNode<T>,
 ): TreeNode<T> | null {
  let n: TreeNode<T> | null = node;
  while (n) {
-    if (isVisibleNode(treeId, n) && !n.hidden) return n;
+    if (isVisibleNode(store, collapsedAtom, n) && !n.hidden) return n;
    if (n.parent == null) return null;
    n = n.parent;
  }
--- a/apps/yaak-client/components/core/tree/context.ts
+++ b/apps/yaak-client/components/core/tree/context.ts
@@ -0,0 +1,60 @@
+import type { WritableAtom } from 'jotai';
+import { useAtomValue, useStore } from 'jotai';
+import { selectAtom } from 'jotai/utils';
+import { createContext, useCallback, useContext, useMemo } from 'react';
+
+type CollapsedMap = Record<string, boolean>;
+type SetAction = CollapsedMap | ((prev: CollapsedMap) => CollapsedMap);
+export type CollapsedAtom = WritableAtom<CollapsedMap, [SetAction], void>;
+
+export const CollapsedAtomContext = createContext<CollapsedAtom | null>(null);
+
+export function useCollapsedAtom(): CollapsedAtom {
+  const atom = useContext(CollapsedAtomContext);
+  if (!atom) throw new Error('CollapsedAtomContext not provided');
+  return atom;
+}
+
+export function useIsCollapsed(itemId: string | undefined) {
+  const collapsedAtom = useCollapsedAtom();
+  const derivedAtom = useMemo(
+    () => selectAtom(collapsedAtom, (map) => !!map[itemId ?? 'n/a'], Object.is),
+    [collapsedAtom, itemId],
+  );
+  return useAtomValue(derivedAtom);
+}
+
+export function useSetCollapsed(itemId: string | undefined) {
+  const collapsedAtom = useCollapsedAtom();
+  const store = useStore();
+  return useCallback(
+    (next: boolean | ((prev: boolean) => boolean)) => {
+      const key = itemId ?? 'n/a';
+      const prevMap = store.get(collapsedAtom);
+      const prevValue = !!prevMap[key];
+      const value = typeof next === 'function' ? next(prevValue) : next;
+      if (value === prevValue) return;
+      store.set(collapsedAtom, { ...prevMap, [key]: value });
+    },
+    [collapsedAtom, itemId, store],
+  );
+}
+
+export function useCollapsedMap() {
+  const collapsedAtom = useCollapsedAtom();
+  return useAtomValue(collapsedAtom);
+}
+
+export function useIsAncestorCollapsed(ancestorIds: string[]) {
+  const collapsedAtom = useCollapsedAtom();
+  const derivedAtom = useMemo(
+    () =>
+      selectAtom(
+        collapsedAtom,
+        (collapsed) => ancestorIds.some((id) => collapsed[id]),
+        (a, b) => a === b,
+      ),
+    [collapsedAtom, ancestorIds],
+  );
+  return useAtomValue(derivedAtom);
+}
--- a/apps/yaak-proxy/main.tsx
+++ b/apps/yaak-proxy/main.tsx
@@ -7,11 +7,18 @@ import { StrictMode, useState } from "react";
 import { createRoot } from "react-dom/client";
 import "./main.css";
 import { listen, rpc } from "./rpc";
-import { applyChange, dataAtom, httpExchangesAtom } from "./store";
+import { applyChange, dataAtom, httpExchangesAtom, replaceAll } from "./store";

 const queryClient = new QueryClient();
 const jotaiStore = createStore();

+// Load initial models from the database
+rpc("list_models", {}).then((res) => {
+  jotaiStore.set(dataAtom, (prev) =>
+    replaceAll(prev, "http_exchange", res.httpExchanges),
+  );
+});
+
 // Subscribe to model change events from the backend
 listen("model_write", (payload) => {
  jotaiStore.set(dataAtom, (prev) =>
--- a/apps/yaak-proxy/store.ts
+++ b/apps/yaak-proxy/store.ts
@@ -5,7 +5,7 @@ type ProxyModels = {
  http_exchange: HttpExchange;
 };

-export const { dataAtom, applyChange, listAtom, orderedListAtom } =
+export const { dataAtom, applyChange, replaceAll, listAtom, orderedListAtom } =
  createModelStore<ProxyModels>(["http_exchange"]);

 export const httpExchangesAtom = orderedListAtom(
--- a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts
+++ b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts
@@ -1,5 +1,9 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
-import type { ModelPayload } from "./gen_models";
+import type { HttpExchange, ModelPayload } from "./gen_models";
+
+export type ListModelsRequest = Record<string, never>;
+
+export type ListModelsResponse = { httpExchanges: Array<HttpExchange>, };

 export type ProxyStartRequest = { port: number | null, };

@@ -9,4 +13,4 @@ export type ProxyStopRequest = Record<string, never>;

 export type RpcEventSchema = { model_write: ModelPayload, };

-export type RpcSchema = { proxy_start: [ProxyStartRequest, ProxyStartResponse], proxy_stop: [ProxyStopRequest, boolean], };
+export type RpcSchema = { proxy_start: [ProxyStartRequest, ProxyStartResponse], proxy_stop: [ProxyStopRequest, boolean], list_models: [ListModelsRequest, ListModelsResponse], };
--- a/crates-proxy/yaak-proxy-lib/src/lib.rs
+++ b/crates-proxy/yaak-proxy-lib/src/lib.rs
@@ -51,6 +51,17 @@ pub struct ProxyStartResponse {
 #[ts(export, export_to = "gen_rpc.ts")]
 pub struct ProxyStopRequest {}

+#[derive(Deserialize, TS)]
+#[ts(export, export_to = "gen_rpc.ts")]
+pub struct ListModelsRequest {}
+
+#[derive(Serialize, TS)]
+#[ts(export, export_to = "gen_rpc.ts")]
+#[serde(rename_all = "camelCase")]
+pub struct ListModelsResponse {
+    pub http_exchanges: Vec<HttpExchange>,
+}
+
 // -- Handlers --

 fn proxy_start(ctx: &ProxyCtx, req: ProxyStartRequest) -> Result<ProxyStartResponse, RpcError> {
@@ -86,6 +97,15 @@ fn proxy_stop(ctx: &ProxyCtx, _req: ProxyStopRequest) -> Result<bool, RpcError>
    Ok(handle.take().is_some())
 }

+fn list_models(ctx: &ProxyCtx, _req: ListModelsRequest) -> Result<ListModelsResponse, RpcError> {
+    ctx.db.with_conn(|db| {
+        Ok(ListModelsResponse {
+            http_exchanges: db.find_all::<HttpExchange>()
+                .map_err(|e| RpcError { message: e.to_string() })?,
+        })
+    })
+}
+
 // -- Event loop --

 fn run_event_loop(rx: std::sync::mpsc::Receiver<ProxyEvent>, db: ProxyQueryManager, events: RpcEventEmitter) {
@@ -193,6 +213,7 @@ define_rpc! {
    commands {
        proxy_start(ProxyStartRequest) -> ProxyStartResponse,
        proxy_stop(ProxyStopRequest) -> bool,
+        list_models(ListModelsRequest) -> ListModelsResponse,
    }
    events {
        model_write(ModelPayload),
--- a/packages/model-store/src/index.ts
+++ b/packages/model-store/src/index.ts
@@ -71,7 +71,20 @@ export function createModelStore<M extends ModelMap>(
    );
  }

-  return { dataAtom, applyChange, listAtom, orderedListAtom };
+  /** Replace all models of a given type. Used for initial hydration. */
+  function replaceAll<K extends keyof M & string>(
+    prev: StoreData<M>,
+    modelType: K,
+    models: M[K][],
+  ): StoreData<M> {
+    const bucket = {} as Record<string, M[K]>;
+    for (const m of models) {
+      bucket[m.id] = m;
+    }
+    return { ...prev, [modelType]: bucket };
+  }
+
+  return { dataAtom, applyChange, replaceAll, listAtom, orderedListAtom };
 }

 function shallowEqual<T>(a: T[], b: T[]): boolean {
--- a/packages/ui/src/components/DropMarker.tsx
+++ b/packages/ui/src/components/DropMarker.tsx
@@ -0,0 +1,34 @@
+import classNames from 'classnames';
+import type { CSSProperties } from 'react';
+import { memo } from 'react';
+
+interface Props {
+  className?: string;
+  style?: CSSProperties;
+  orientation?: 'horizontal' | 'vertical';
+}
+
+export const DropMarker = memo(
+  function DropMarker({ className, style, orientation = 'horizontal' }: Props) {
+    return (
+      <div
+        style={style}
+        className={classNames(
+          className,
+          'absolute pointer-events-none z-50',
+          orientation === 'horizontal' && 'w-full',
+          orientation === 'vertical' && 'w-0 top-0 bottom-0',
+        )}
+      >
+        <div
+          className={classNames(
+            'absolute bg-primary rounded-full',
+            orientation === 'horizontal' && 'left-2 right-2 -bottom-[0.1rem] h-[0.2rem]',
+            orientation === 'vertical' && '-left-[0.1rem] top-0 bottom-0 w-[0.2rem]',
+          )}
+        />
+      </div>
+    );
+  },
+  () => true,
+);
--- a/packages/ui/src/index.ts
+++ b/packages/ui/src/index.ts
@@ -9,3 +9,5 @@ export { useIsFullscreen } from "./hooks/useIsFullscreen";
 export { useDebouncedValue } from "./hooks/useDebouncedValue";
 export { useDebouncedState } from "./hooks/useDebouncedState";
 export { HEADER_SIZE_MD, HEADER_SIZE_LG, WINDOW_CONTROLS_WIDTH } from "./lib/constants";
+export { DropMarker } from "./components/DropMarker";
+export { computeSideForDragMove } from "./lib/dnd";
--- a/packages/ui/src/lib/dnd.ts
+++ b/packages/ui/src/lib/dnd.ts
@@ -0,0 +1,39 @@
+import type { DragMoveEvent } from '@dnd-kit/core';
+
+export function computeSideForDragMove(
+  id: string,
+  e: DragMoveEvent,
+  orientation: 'vertical' | 'horizontal' = 'vertical',
+): 'before' | 'after' | null {
+  if (e.over == null || e.over.id !== id) {
+    return null;
+  }
+  if (e.active.rect.current.initial == null) return null;
+
+  const overRect = e.over.rect;
+
+  if (orientation === 'horizontal') {
+    // For horizontal layouts (tabs side-by-side), use left/right logic
+    const activeLeft =
+      e.active.rect.current.translated?.left ?? e.active.rect.current.initial.left + e.delta.x;
+    const pointerX = activeLeft + e.active.rect.current.initial.width / 2;
+
+    const hoverLeft = overRect.left;
+    const hoverRight = overRect.right;
+    const hoverMiddleX = hoverLeft + (hoverRight - hoverLeft) / 2;
+
+    return pointerX < hoverMiddleX ? 'before' : 'after'; // 'before' = left, 'after' = right
+  } else {
+    // For vertical layouts, use top/bottom logic
+    const activeTop =
+      e.active.rect.current.translated?.top ?? e.active.rect.current.initial.top + e.delta.y;
+    const pointerY = activeTop + e.active.rect.current.initial.height / 2;
+
+    const hoverTop = overRect.top;
+    const hoverBottom = overRect.bottom;
+    const hoverMiddleY = (hoverBottom - hoverTop) / 2;
+    const hoverClientY = pointerY - hoverTop;
+
+    return hoverClientY < hoverMiddleY ? 'before' : 'after';
+  }
+}