From 6534421733a427adc5e2a5233db4c7ba06729d86 Mon Sep 17 00:00:00 2001
From: Gregory Schier <gschier1990@gmail.com>
Date: Sun, 8 Mar 2026 16:37:25 -0700
Subject: [PATCH] Start extracting Tree component

---
 .../components/EnvironmentEditDialog.tsx      |  17 +++
 apps/yaak-client/components/Sidebar.tsx       |  21 ++-
 .../components/core/PairEditor.tsx            |   3 +-
 .../yaak-client/components/core/Tabs/Tabs.tsx |   3 +-
 .../yaak-client/components/core/tree/Tree.tsx | 140 +++++++++---------
 .../components/core/tree/TreeDropMarker.tsx   |   8 +-
 .../components/core/tree/TreeItem.tsx         |  59 ++++----
 .../components/core/tree/TreeItemList.tsx     |   2 +-
 .../yaak-client/components/core/tree/atoms.ts |  26 ----
 .../components/core/tree/common.ts            |  26 +++-
 .../components/core/tree/context.ts           |  60 ++++++++
 apps/yaak-proxy/main.tsx                      |   9 +-
 apps/yaak-proxy/store.ts                      |   2 +-
 .../yaak-proxy-lib/bindings/gen_rpc.ts        |   8 +-
 crates-proxy/yaak-proxy-lib/src/lib.rs        |  21 +++
 packages/model-store/src/index.ts             |  15 +-
 packages/ui/src/components/DropMarker.tsx     |  34 +++++
 packages/ui/src/index.ts                      |   2 +
 packages/ui/src/lib/dnd.ts                    |  39 +++++
 19 files changed, 344 insertions(+), 151 deletions(-)
 create mode 100644 apps/yaak-client/components/core/tree/context.ts
 create mode 100644 packages/ui/src/components/DropMarker.tsx
 create mode 100644 packages/ui/src/lib/dnd.ts

diff --git a/apps/yaak-client/components/EnvironmentEditDialog.tsx b/apps/yaak-client/components/EnvironmentEditDialog.tsx
index f7faf6f8..ddce40b6 100644
--- 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}
diff --git a/apps/yaak-client/components/Sidebar.tsx b/apps/yaak-client/components/Sidebar.tsx
index 53abdf9b..ee28654c 100644
--- 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"
diff --git a/apps/yaak-client/components/core/PairEditor.tsx b/apps/yaak-client/components/core/PairEditor.tsx
index b6601932..b37a01ad 100644
--- 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';
diff --git a/apps/yaak-client/components/core/Tabs/Tabs.tsx b/apps/yaak-client/components/core/Tabs/Tabs.tsx
index 4a616d80..731cd76a 100644
--- 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';
diff --git a/apps/yaak-client/components/core/tree/Tree.tsx b/apps/yaak-client/components/core/tree/Tree.tsx
index 73b9254d..e5dd008a 100644
--- 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;
diff --git a/apps/yaak-client/components/core/tree/TreeDropMarker.tsx b/apps/yaak-client/components/core/tree/TreeDropMarker.tsx
index 71947940..4cbc667c 100644
--- 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;
diff --git a/apps/yaak-client/components/core/tree/TreeItem.tsx b/apps/yaak-client/components/core/tree/TreeItem.tsx
index aab0338b..ddc362d8 100644
--- 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"
diff --git a/apps/yaak-client/components/core/tree/TreeItemList.tsx b/apps/yaak-client/components/core/tree/TreeItemList.tsx
index ea09faae..1cb006a5 100644
--- 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>[];
diff --git a/apps/yaak-client/components/core/tree/atoms.ts b/apps/yaak-client/components/core/tree/atoms.ts
index b65a1f18..735226bd 100644
--- 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,
-);
diff --git a/apps/yaak-client/components/core/tree/common.ts b/apps/yaak-client/components/core/tree/common.ts
index 9bc112e7..75b944d8 100644
--- 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;
   }
diff --git a/apps/yaak-client/components/core/tree/context.ts b/apps/yaak-client/components/core/tree/context.ts
new file mode 100644
index 00000000..8c4e4476
--- /dev/null
+++ 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);
+}
diff --git a/apps/yaak-proxy/main.tsx b/apps/yaak-proxy/main.tsx
index a8b05384..01199df2 100644
--- 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) =>
diff --git a/apps/yaak-proxy/store.ts b/apps/yaak-proxy/store.ts
index 66b1e54f..59782b18 100644
--- 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(
diff --git a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts
index bc133b2e..1e1bff59 100644
--- 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], };
diff --git a/crates-proxy/yaak-proxy-lib/src/lib.rs b/crates-proxy/yaak-proxy-lib/src/lib.rs
index 649b5f9b..2e3f93b5 100644
--- 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),
diff --git a/packages/model-store/src/index.ts b/packages/model-store/src/index.ts
index 922e5058..3a290c40 100644
--- 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 {
diff --git a/packages/ui/src/components/DropMarker.tsx b/packages/ui/src/components/DropMarker.tsx
new file mode 100644
index 00000000..18547781
--- /dev/null
+++ 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,
+);
diff --git a/packages/ui/src/index.ts b/packages/ui/src/index.ts
index 26b56dfc..6e1ad69d 100644
--- 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";
diff --git a/packages/ui/src/lib/dnd.ts b/packages/ui/src/lib/dnd.ts
new file mode 100644
index 00000000..8beb744f
--- /dev/null
+++ 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';
+  }
+}